分布鲁棒的混合专家训练

arXiv:2610.07207v1 公告类型：新 摘要：混合专家（MoE）Transformer 通过为每个 token 仅激活少数专家来扩展模型容量，但这种稀疏性带来了一个隐藏的可靠性问题：当路由不够完善时，负载均衡模型可能会将 token 发送给那些对所分配输入训练不足的专家。我们提出分布鲁棒混合专家训练（DRMoET），这是一种可直接替换的目标函数，将逐层专家视为内生的鲁棒性群组，并针对高损失路由结果进行优化，而不仅仅是均衡流量。DRMoET 基于经过 EMA 平滑、并按激活加权的专家损失，通过熵正则化 softmax 规则更新每层的专家分布，在保留标准 MoE 计算方式的同时，强化合理的非首选路由路径。在总参数量为 746M 和 10.3B 的规模下，按照 FLAME-MoE 配方训练时，DRMoET 在下游任务平均表现上均优于标准 FLAME-MoE 和无辅助损失均衡方法。在总参数量为 10.3B、训练 token 数为 67B 时，DRMoET 将七项任务的平均分从 0.6625 提升至 0.6767，而无辅助损失基线的平均分为 0.6431。机制分析表明，在平均损失几乎不变的情况下，专家损失方差更低；在强制进行中间-$k$ 误路由时，额外损失降低了 4.3%；同时，领域专家的专业化程度得到提升。这些结果表明，在可靠地扩展稀疏 MoE 方面，路由鲁棒性与利用率均衡同样是一个实用的优化目标。项目主页和代码地址：https://drmoet.github.io/。