Research map / Architectures and efficiency / Mixture of experts
MomentumSMoE: Integrating Momentum into Sparse Mixture of Experts
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Connects SMoE layer dynamics to multiple-gradient descent and adds heavy-ball, Adam-style and robust momentum to SMoE/MoE layers, showing improved stability, perplexity and ImageNet robustness.
Problem
Sparse Mixture of Experts (SMoE) training is unstable and its models adapt poorly to new distributions and data contamination.
Approach
Views the SMoE layer update as gradient descent on a multi-objective problem (expert outputs as local gradients) and integrates heavy-ball momentum, AdamW-style and robust momentum into the SMoE residual dynamics, giving MomentumSMoE, AdamSMoE and Robust MomentumSMoE.
Key results
- On WikiText-103, AdamSMoE-medium reaches 33.25 test PPL vs 35.55 for the SMoE baseline; on attacked data 41.11 vs 44.19.
- On ImageNet-1K, Robust MomentumV-MoE improves IN-C top-1 from 46.98 to 48.82 and mCE from 67.14 to 64.92 over V-MoE.
- Momentum-Soft MoE improves ImageNet-1K top-1 from 73.52 to 75.47 over Soft MoE baseline.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add heavy-ball momentum | SMoE-medium · WikiText-103 clean test | perplexity: 35.55 -> 33.46 | Helped |
| Use AdamSMoE in the first layer and MomentumSMoE in subsequent layers | SMoE-medium · WikiText-103 clean test | perplexity: 35.55 -> 33.25 | Helped |
| Add heavy-ball momentum | SMoE-medium · WikiText-103 attacked test | perplexity: 44.19 -> 42.33 | Helped |
| Use AdamSMoE in the first layer and MomentumSMoE in subsequent layers | SMoE-medium · WikiText-103 attacked test | perplexity: 44.19 -> 41.11 | Helped |
| Add heavy-ball momentum | SMoE-large · WikiText-103 clean test | perplexity: 30.33 -> 29.03 | Helped |
| Add heavy-ball momentum | GLaM-medium · WikiText-103 clean test | perplexity: 37.71 -> 35.29 | Helped |
| Add heavy-ball momentum | V-MoE small · ImageNet-R | top-1 accuracy: 36.10 -> 37.45 | Helped |
| Use robust momentum | V-MoE small · ImageNet-C | mCE: 67.14 -> 64.92 | Helped |
| Add momentum to Soft MoE layers | Soft MoE-tiny · ImageNet-1K validation | top-1 accuracy: 73.52 -> 75.47 | Helped |
| Learn both momentum and step size instead of tuning them (ablation) | MomentumSMoE · WikiText-103 clean validation | perplexity: 32.29 -> 32.28 | No clear effect |
| Learn step size while fixing momentum at 0.7 (ablation) | MomentumSMoE · WikiText-103 attacked validation | perplexity: 40.94 -> 39.96 | Helped |
| Learn step size as a linear network (ablation) | MomentumSMoE · WikiText-103 clean test | perplexity: 33.46 -> 33.96 | Hurt |
Surprising findings
- Momentum gives no benefit for the small 3-layer SMoE and slightly hurts it; benefit grows with depth.
- Introducing AdamSMoE into V-MoE led to unstable training, so it was deferred to future work.
Limitations
- Momentum offers little or negative benefit for shallow models (3 layers).
- AdamSMoE is unstable in V-MoE and in the paper was used only in the first layer of language models.
- Theoretical explanation of enhanced robustness is left open.
Experimental setup
Datasets: WikiText-103 ImageNet-1K
Benchmarks: WikiText-103 WikiText-103 word-swap attacked ImageNet-1K ImageNet-R ImageNet-A ImageNet-C
Metrics: perplexity top-1 accuracy top-5 accuracy mean corruption error
Baselines: SMoE (Switch Transformer) GLaM V-MoE Soft MoE
Builds on: Switch Transformer GLaM V-MoE Soft MoE Robust Momentum Method MGDA
Closest papers by meaning
- Multilinear Mixture of Experts: Scalable Expert Specialization through Factorization (NeurIPS 2024 · similarity 0.87)
- SMMF: Square-Matricized Momentum Factorization for Memory-Efficient Optimization (AAAI 2025 · similarity 0.87)
- The Marginal Value of Momentum for Small Learning Rate SGD (ICLR 2024 · similarity 0.86)
- How to Scale Mixture-of-Experts: From muP to the Maximally Scale-Stable Parameterization (NeurIPS 2026 · similarity 0.86)
- Retraining-free Merging of Sparse MoE via Hierarchical Clustering (ICML 2025 · similarity 0.86)
- Merge, Then Compress: Demystify Efficient SMoE with Hints from Its Routing Policy (ICLR 2024 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.