arXiv:2503.05029cs.LGcs.AI2025-03被引 5

MoE模型持续预训练效果稳定,无需重训也能高效保持性能。

Continual Pre-training of MoEs: How robust is your router?

  • 采用路由算法优化,实现无重放缓存的持续预训练
  • 在6000亿词元上训练,性能接近全量重训结果
  • 适合资源有限但需持续更新的大模型应用

稀疏激活的专家混合(MoE)Transformer是基础模型的有前景架构。与相同前向计算量的密集模型相比,MoE在训练时具有更高的样本效率并能实现更强性能。许多前沿语言模型已采用MoE结构。实践中,人们希望在不完全重训的前提下,利用大量新数据扩展模型能力。已有研究证明,仅通过回放、学习率重暖和衰减即可实现密集解码器模型的持续预训练(CPT),且性能损失极小。然而,对于解码器型MoE模型,其路由机制如何影响持续预训练尚不清楚:1)路由是否加剧遗忘?2)模型能否维持对历史分布的负载均衡?3)适用于密集模型的策略是否足以支持MoE模型的持续预训练?本文针对500M参数密集模型及四款500M活跃/2B总参数的MoE模型展开大规模实验,每模型均训练600B词元。结果表明,无论使用Sinkhorn-Balanced还是Z-and-Aux-loss-balanced路由算法,MoE对分布偏移表现出惊人鲁棒性,甚至在无回放条件下仍表现良好。此外,MoE在持续预训练中保持相对于等量FLOP密集模型的样本效率,并以极低成本达到全量重训模型的性能。

原文摘要 · Abstract (English)

Sparsely-activated Mixture of Experts (MoE) transformers are promising architectures for foundation models. Compared to dense transformers that require the same amount of floating-point operations (FLOPs) per forward pass, MoEs benefit from improved sample efficiency at training time and achieve much stronger performance. Many closed-source and open-source frontier language models have thus adopted an MoE architecture. Naturally, practitioners will want to extend the capabilities of these models with large amounts of newly collected data without completely re-training them. Prior work has shown that a simple combination of replay, learning rate re-warming, and re-decaying can enable the continual pre-training (CPT) of dense decoder-only transformers with minimal performance degradation compared to full re-training. In the case of decoder-only MoE transformers, however, it is unclear how the routing algorithm will impact continual pre-training performance: 1) do the MoE transformer's routers exacerbate forgetting relative to a dense model?; 2) do the routers maintain a balanced load on previous distributions after CPT?; 3) are the same strategies applied to dense models sufficient to continually pre-train MoE LLMs? In what follows, we conduct a large-scale study training a 500M parameter dense transformer and four 500M-active/2B-total parameter MoE transformers. Each model is trained for 600B tokens. Our results establish a surprising robustness to distribution shifts for MoEs using both Sinkhorn-Balanced and Z-and-Aux-loss-balanced routing algorithms, even in MoEs continually pre-trained without replay. Moreover, we show that MoE LLMs maintain their sample efficiency (relative to a FLOP-matched dense model) during CPT and that they can match the performance of a fully re-trained MoE at a fraction of the cost.

MoE持续学习大模型

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。