提出GRIN训练法,让专家模型更高效地优化和扩展。
GRIN: GRadient-INformed MoE
- 用稀疏梯度估计改进专家路由,提升训练效果
- 16×3.8B专家模型仅激活6.6B参数,性能超7B稠密模型
- 适合追求高效大模型训练的研究者与工程师
Mixture-of-Experts(MoE)模型通过专家路由实现稀疏计算,仅激活少量专家模块,因此比稠密模型更具可扩展性。然而,离散的专家路由阻碍了标准反向传播,难以进行基于梯度的优化。为更好发挥MoE的扩展潜力,本文提出GRIN(GRadient-INformed MoE训练),引入稀疏梯度估计优化专家路由,并配置模型并行以避免令牌丢失。在自回归语言建模中,我们构建了一个top-2、16×3.8B的MoE模型。该模型仅激活6.6B参数,性能超越7B稠密模型,并达到与14B稠密模型相当的水平。在多种任务上广泛评估显示,GRIN显著提升了MoE的有效性:在MMLU上达79.4,在HellaSwag上达83.7,在HumanEval上达74.4,在MATH上达58.9。
原文摘要 · Abstract (English)
Mixture-of-Experts (MoE) models scale more effectively than dense models due to sparse computation through expert routing, selectively activating only a small subset of expert modules. However, sparse computation challenges traditional training practices, as discrete expert routing hinders standard backpropagation and thus gradient-based optimization, which are the cornerstone of deep learning. To better pursue the scaling power of MoE, we introduce GRIN (GRadient-INformed MoE training), which incorporates sparse gradient estimation for expert routing and configures model parallelism to avoid token dropping. Applying GRIN to autoregressive language modeling, we develop a top-2 16$\times$3.8B MoE model. Our model, with only 6.6B activated parameters, outperforms a 7B dense model and matches the performance of a 14B dense model trained on the same data. Extensive evaluations across diverse tasks demonstrate the potential of GRIN to significantly enhance MoE efficacy, achieving 79.4 on MMLU, 83.7 on HellaSwag, 74.4 on HumanEval, and 58.9 on MATH.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。