# MomentumSMoE: Integrating Momentum into Sparse Mixture of Experts

arXiv: https://arxiv.org/abs/2410.14574 · HTML: https://mlsanity.com/paper/2410.14574/

## Key facts

- Cited by 9 papers in the map (data through 14 Oct 2026).
- Placed in MoE theory and gating functions, one of 12 directions in Mixture-of-experts.
- Builds on 60 of its 78 references that are in the map.
- First posted on arXiv 18 Oct 2024.

## Placed in

- Mixture-of-experts / [MoE theory and gating functions](https://mlsanity.com/moe/moe-theory-and-gating-functions/)

## Builds on (most cited first, 30 of 60)

- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [Visual Instruction Tuning](https://mlsanity.com/paper/2304.08485/) (2023)
- [Momentum Contrast for Unsupervised Visual Representation Learning](https://mlsanity.com/paper/1911.05722/) (2019)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [BLIP-2: Bootstrapping Language-Image Pre-training with Frozen Image Encoders and Large Language Models](https://mlsanity.com/paper/2301.12597/) (2023)
- [PaLM: Scaling Language Modeling with Pathways](https://mlsanity.com/paper/2204.02311/) (2022)
- [Training data-efficient image transformers & distillation through attention](https://mlsanity.com/paper/2012.12877/) (2020)
- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [Pointer Sentinel Mixture Models](https://mlsanity.com/paper/1609.07843/) (2016)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [Benchmarking Neural Network Robustness to Common Corruptions and Perturbations](https://mlsanity.com/paper/1903.12261/) (2019)
- [BEiT: BERT Pre-Training of Image Transformers](https://mlsanity.com/paper/2106.08254/) (2021)
- [The Many Faces of Robustness: A Critical Analysis of Out-of-Distribution Generalization](https://mlsanity.com/paper/2006.16241/) (2020)
- [GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding](https://mlsanity.com/paper/2006.16668/) (2020)
- [Decision Transformer: Reinforcement Learning via Sequence Modeling](https://mlsanity.com/paper/2106.01345/) (2021)
- [Do ImageNet Classifiers Generalize to ImageNet?](https://mlsanity.com/paper/1902.10811/) (2019)
- [Natural Adversarial Examples](https://mlsanity.com/paper/1907.07174/) (2019)
- [On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minima](https://mlsanity.com/paper/1609.04836/) (2016)
- [GLaM: Efficient Scaling of Language Models with Mixture-of-Experts](https://mlsanity.com/paper/2112.06905/) (2021)
- [Scaling Vision with Sparse Mixture of Experts](https://mlsanity.com/paper/2106.05974/) (2021)
- [Offline Reinforcement Learning as One Big Sequence Modeling Problem](https://mlsanity.com/paper/2106.02039/) (2021)
- [Mixture-of-Experts with Expert Choice Routing](https://mlsanity.com/paper/2202.09368/) (2022)
- [Computing Nonvacuous Generalization Bounds for Deep (Stochastic) Neural Networks with Many More Parameters than Training Data](https://mlsanity.com/paper/1703.11008/) (2017)
- [Fantastic Generalization Measures and Where to Find Them](https://mlsanity.com/paper/1912.02178/) (2019)
- [VLMo: Unified Vision-Language Pre-Training with Mixture-of-Modality-Experts](https://mlsanity.com/paper/2111.02358/) (2021)
- [ST-MoE: Designing Stable and Transferable Sparse Expert Models](https://mlsanity.com/paper/2202.08906/) (2022)

## Built on by (most cited first, 9 of 9)

- [Tight Clusters Make Specialized Experts](https://mlsanity.com/paper/2502.15315/) (2025)
- [LIBMoE: A Library for comprehensive benchmarking Mixture of Experts in Large Language Models](https://mlsanity.com/paper/2411.00918/) (2024)
- [CAMEx: Curvature-aware Merging of Experts](https://mlsanity.com/paper/2502.18821/) (2025)
- [MoLEx: Mixture of Layer Experts for Finetuning with Sparse Upcycling](https://mlsanity.com/paper/2503.11144/) (2025)
- [Expert Merging in Sparse Mixture of Experts with Nash Bargaining](https://mlsanity.com/paper/2510.16138/) (2025)
- [MuonSSM: Orthogonalizing State Space Models for Sequence Modeling](https://mlsanity.com/paper/2606.30461/) (2026)
- [Selective Sinkhorn Routing for Improved Sparse Mixture of Experts](https://mlsanity.com/paper/2511.08972/) (2025)
- [Mixture of Group Experts for Learning Invariant Representations](https://mlsanity.com/paper/2504.09265/) (2025)
- [Mixture-of-Control: State-Aware Fine-Tuning for Transformer-based Models](https://mlsanity.com/paper/2606.31397/) (2026)

## Abstract (arXiv)

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.