MoE不只要负载均衡:DRMoET让专家更能扛住路由偏差
导语
混合专家模型(Mixture-of-Experts,MoE)通过为每个Token只激活少数专家,在控制计算量的同时扩大模型容量。但稀疏路由也带来一个容易被忽略的问题:路由器把Token送到某个专家,并不意味着这个专家已经充分学会处理这类输入。即使不同专家接收的Token数量看起来很均衡,模型仍可能在路由偏差或候选专家变化时表现不稳定。
一篇题为《Distributionally Robust Mixture-of-Experts Training》的工作提出DRMoET,试图把MoE优化目标从“分得均匀”推进到“分错了也更可靠”。
核心方法:关注高损失路由结果
传统负载均衡机制主要约束专家的使用率,目标是避免少数专家过载、其他专家闲置。DRMoET则进一步追踪每层专家在实际激活数据上的损失,并使用指数滑动平均(EMA)降低短期波动的影响。随后,方法通过带熵正则的Softmax规则更新每层的专家分布,让高损失、但仍然可能发生的路由结果获得更强训练关注。
这并不意味着模型在推理时要激活更多专家,也不是简单地把Token平均分配给所有专家。它仍保留标准MoE的稀疏计算形式,变化主要发生在训练目标上:除了训练当前最优路由,还要让潜在的非Top路由路径具备更好的处理能力。换言之,专家不再只为“最常被派来的输入”优化,而是要对一定范围内的路由不确定性保持韧性。
实验结果
作者在FLAME-MoE训练配方下测试了746M和10.3B总参数规模的模型,并与标准FLAME-MoE及无辅助损失的负载均衡方案比较。在10.3B参数、67B训练Token的设置中,DRMoET的七项任务平均得分从基线的0.6625提升至0.6767;辅助损失自由的基线为0.6431。摘要还指出,DRMoET在两个规模上都带来了下游平均成绩改善。
机制分析提供了更细的解释:DRMoET降低了专家损失的方差,而平均损失基本没有明显变化;在强制使用中间排名专家、模拟误路由的测试中,额外损失下降了4.3%;同时,专家在不同领域上的分工更加清晰。
意义与边界
这项工作提示,MoE的可靠性不能只用容量利用率或Token流量均衡来衡量。真正重要的还包括:当路由器出现不确定性时,候选专家是否具备足够的覆盖能力,以及专家之间的能力差距是否过大。DRMoET作为一种可插拔训练目标,为稀疏模型的稳定扩展提供了一个相对直接的方向。
不过,现有材料主要呈现作者在特定训练配方、规模和任务上的结果。它是否能在更多路由结构、数据分布及超大规模训练中保持收益,仍需要进一步验证。对工程实践而言,如何选择熵正则和损失平滑相关超参数,也可能影响方法的实际成本与效果。
评论
正在确认登录状态……
正在加载评论……