Continual Pre-training of MoEs: How robust is your router?
作者: Benjamin Thérien, Charles-Étienne Joseph, Zain Sarwar, Ashwinee Panda, Anirban Das, Shi-Xiong Zhang, Stephen Rawls, Sambit Sahu, Eugene Belilovsky, Irina Rish
发布时间: 2025-03-10
来源: arxiv
研究方向: 大规模语言模型（LLM）的持续预训练与混合专家（MoE）架构
主要内容
该研究探讨了如何通过持续预训练方法来更新MoE架构的LLM，同时保持其性能和效率。研究者对比了MoE与密集模型在持续预训练过程中的表现，并分析了不同路由算法和预训练策略对MoE性能的影响。
主要贡献
1. 建立了重放和无限学习率（LR）调度对MoE Transformer LMs持续预训练中遗忘和路由不平衡动态的影响。
2. 证明了使用DeepSeek架构的惩罚平衡MoE可以以较低的计算成本匹配全重训练基线的性能。
3. 展示了惩罚平衡和Sinkhorn平衡路由算法在语言建模性能、评估基准和最大路由不平衡方面的出色鲁棒性。
4. 提供了持续预训练期间路由决策变化的全面分析，揭示了MoE如何适应新分布和遗忘旧分布。
5. 提出了一个衡量MoE延迟的新指标：最大路由不平衡（MRI）。
研究方法
1. 对比了两种流行的路由算法：惩罚平衡（PBTk）和Sinkhorn平衡（SBTk）。
2. 使用了两种流行的MoE架构：DeepSeek和Switch。
3. 在400B tokens的FineWeb和200B tokens的代码数据以及德国网络爬虫数据上预训练和持续预训练了MoE语言模型。
4. 比较了持续预训练MoE的性能与全重训练基线。
5. 分析了路由决策在持续预训练期间如何变化，以了解MoE如何适应新分布和遗忘旧分布。
实验结果
实验结果表明，MoEs在持续预训练过程中保持了其样本效率，并且可以以较低的成本匹配全重训练MoE的性能。此外，PBTk和SBTk路由算法对分布变化表现出惊人的鲁棒性，并且持续预训练MoEs在语言建模性能、评估基准和最大路由不平衡方面都表现出良好的性能。
未来工作
未来研究可以探索如何更好地处理MoE预训练中的早期层遗忘问题，以及如何进一步优化路由算法以提高MoE的效率和性能。此外，可以研究MoE在特定领域或任务中的应用，以及如何将其与其他技术（如知识蒸馏和元学习）结合使用。