arXiv:2410.14574cs.LGcs.AI2024-10NeurIPS被引 10

给稀疏专家模型引入动量机制,提升训练稳定性和抗数据污染能力。

MomentumSMoE: Integrating Momentum into Sparse Mixture of Experts

  • 将专家表示动态类比为多目标优化梯度下降,构建动量融合框架
  • 在ImageNet-1K和WikiText-103上验证比传统SMoE更稳定、更鲁棒
  • 兼容V-MoE、GLaM等主流架构,可无缝接入Adam等优化器

稀疏混合专家(SMoE)是实现深度学习可扩展性的关键,通过仅激活少量参数即可实现参数量的指数增长,同时保持高效性。然而,研究发现SMoE存在训练不稳定、难以适应新分布的问题,导致对数据污染缺乏鲁棒性。本文首次建立SMoE中专家表示动态与多目标优化梯度下降之间的联系,并在此基础上引入动量机制,提出新型SMoE家族——MomentumSMoE。理论证明并数值实证表明,MomentumSMoE相比SMoE更具稳定性与鲁棒性。我们在ImageNet-1K图像识别和WikiText-103语言建模等多个实际任务中验证了其优势。该方法适用于多种SMoE模型,包括视觉稀疏专家模型(V-MoE)和通用语言模型(GLaM)。此外,还可轻松集成如Adam等先进动量优化方法,在几乎不增加计算成本的前提下设计性能更优的新模型,且实现简单。

原文摘要 · Abstract (English)

Sparse Mixture of Experts (SMoE) has become the key to unlocking unparalleled scalability in deep learning. SMoE has the potential to exponentially increase parameter count while maintaining the efficiency of the model by only activating a small subset of these parameters for a given sample. However, it has been observed that SMoE suffers from unstable training and has difficulty adapting to new distributions, leading to the model's lack of robustness to data contamination. To overcome these limitations, we first establish a connection between the dynamics of the expert representations in SMoEs and gradient descent on a multi-objective optimization problem. Leveraging our framework, we then integrate momentum into SMoE and propose a new family of SMoEs named MomentumSMoE. We theoretically prove and numerically demonstrate that MomentumSMoE is more stable and robust than SMoE. In particular, we verify the advantages of MomentumSMoE over SMoE on a variety of practical tasks including ImageNet-1K object recognition and WikiText-103 language modeling. We demonstrate the applicability of MomentumSMoE to many types of SMoE models, including those in the Sparse MoE model for vision (V-MoE) and the Generalist Language Model (GLaM). We also show that other advanced momentum-based optimization methods, such as Adam, can be easily incorporated into the MomentumSMoE framework for designing new SMoE models with even better performance, almost negligible additional computation cost, and simple implementations.

稀疏专家动量机制模型鲁棒性高效训练

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