MoE 的持续预训练：你的路由器有多鲁棒？
机器学习
2025-11-11 v2 人工智能
计算与语言
摘要
稀疏激活的专家混合模型（MoE）Transformer 是有前景的基础模型架构。与每次前向传播需要相同浮点运算次数（FLOPs）的密集 Transformer 相比，MoE 在训练时受益于更高的样本效率，并取得了更强的性能。因此，许多闭源和开源前沿语言模型已采用 MoE 架构。自然 practitioners 希望在不完全重新训练的情况下，利用大量新收集的数据扩展这些模型的能力。先前工作已表明，简单的重放、学习率重新升温与重新衰减的组合可以在与完全重新训练相比性能退化极小的情况下实现密集仅解码器 Transformer 的持续预训练（CPT）。然而，对于仅解码器 MoE Transformer，路由算法将如何影响持续预训练性能尚不清楚：1）MoE Transformer 的路由器是否会加剧相对于密集模型的遗忘？2）路由器在 CPT 后是否对先前分布保持均衡负载？3）应用于密集模型的相同策略是否足以持续预训练 MoE LLM？接下来，我们进行了一项大规模研究，训练了一个 500M 参数的密集 Transformer 和四个 500M 活跃/2B 总参数的 MoE Transformer。每个模型训练 600B 个 token。我们的结果表明，使用 Sinkhorn-Balanced 和 Z-and-Aux-loss-balanced 路由算法的 MoE 对分布偏移具有惊人的鲁棒性，即使在没有重放的持续预训练 MoE 中也是如此。此外，我们表明 MoE LLM 在 CPT 期间保持了其样本效率（相对于 FLOPs 匹配的密集模型），并且可以在极低的成本下匹配完全重新训练的 MoE 的性能。
引用
@article{arxiv.2503.05029,
title = {Continual Pre-training of MoEs: How robust is your router?},
author = {Benjamin Thérien and Charles-Étienne Joseph and Zain Sarwar and Ashwinee Panda and Anirban Das and Shi-Xiong Zhang and Stephen Rawls and Sambit Sahu and Eugene Belilovsky and Irina Rish},
journal= {arXiv preprint arXiv:2503.05029},
year = {2025}
}