Research map / Architectures and efficiency / Mixture of experts
From Sparse to Soft Mixtures 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.
Proposes Soft MoE, a fully differentiable sparse Transformer MoE that routes soft-mixed slots (weighted averages of all tokens) to experts, evaluated on JFT-pretrained vision transformers against ViT and sparse MoEs.
Problem
Sparse MoEs rely on discrete token-to-expert assignment, causing training instability, token dropping, expert imbalance, and difficulty scaling experts or finetuning.
Approach
Replace hard token-expert routing with soft assignment: each expert slot is a learned convex combination of all input tokens (dispatch weights), and each output token is a convex combination of all expert output slots (combine weights), keeping everything differentiable and avoiding sort/top-k operations.
Key results
- Soft MoE L/16 outperforms ViT H/14 on upstream, few-shot and finetuning while being about 2x faster at inference (4.8 vs 8.6 ms/img).
- Soft MoE B/16 at 1k TPU-days matches or outperforms ViT H/14 and is over 5x cheaper in wall-clock inference (1.5 vs 8.6 ms/img).
- Soft MoE Huge/14 with 128 experts has over 40x more parameters than ViT H/14 with only about 2% increased inference time (in the abstract).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace the second half of ViT S/16 MLP blocks with Soft MoE layers | ViT S/16 · JFT-4B | precision-at-1: 51.3 -> 58.1 | Helped |
| Replace the second half of ViT S/16 MLP blocks with Soft MoE layers | ViT S/16 · ImageNet-1k 10-shot | accuracy: 67.6 -> 78.8 | Helped |
| Replace the second half of ViT S/16 MLP blocks with Soft MoE layers | ViT S/16 · ImageNet-1k finetune | accuracy: 84.0 -> 86.8 | Helped |
| Replace the second half of ViT L/16 MLP blocks with Soft MoE layers | ViT L/16 · JFT-4B | precision-at-1: 59.8 -> 61.3 | Helped |
| Replace the second half of ViT L/16 MLP blocks with Soft MoE layers | ViT L/16 · ImageNet-1k 10-shot | accuracy: 81.5 -> 83.7 | Helped |
| Replace the second half of ViT L/16 MLP blocks with Soft MoE layers | ViT L/16 · ImageNet-1k finetune | accuracy: 88.5 -> 88.9 | Helped |
| Average expert outputs uniformly instead of learning combine weights (ablation) | Soft MoE S/14 256E · JFT-4B | precision-at-1: 54.3% -> 53.6% | Hurt |
| Average expert outputs uniformly instead of learning combine weights (ablation) | Soft MoE S/14 256E · ImageNet-1k 10-shot | accuracy: 74.8% -> 72.0% | Hurt |
| Average input tokens uniformly instead of learning dispatch weights (ablation) | Soft MoE S/14 256E · JFT-4B | precision-at-1: 54.3% -> 52.6% | Hurt |
| Average input tokens uniformly instead of learning dispatch weights (ablation) | Soft MoE S/14 256E · ImageNet-1k 10-shot | accuracy: 74.8% -> 71.8% | Hurt |
| Use uniform averaging for both dispatch and combine (ablation) | Soft MoE S/14 256E · JFT-4B | precision-at-1: 54.3% -> 51.8% | Hurt |
| Route tokens to experts in identity order without mixing (ablation) | Soft MoE S/14 256E · ImageNet-1k 10-shot | accuracy: 74.8% -> 69.1% | Hurt |
Surprising findings
- Softmax over layer-normalized inputs collapses to one-hot as model dimension grows; L2 normalization of inputs and router weights is needed for stability.
- Identity and uniform (fixed) routing substantially underperform Soft MoE, though they still outperform dense ViT.
Limitations
- Not directly applicable to auto-regressive decoders because mixing all tokens breaks causality.
- Lazy experts and memory consumption grow with large expert counts since parameters scale with experts.
- Main pretraining data (JFT-4B, WebLI) is proprietary, limiting reproducibility.
Experimental setup
Datasets: JFT-4B WebLI LAION-400M ImageNet-1k Cifar-100 Pets COCO
Benchmarks: JFT-4B precision-at-1 ImageNet 10-shot ImageNet finetuning ImageNet zero-shot Cifar-100 zero-shot Pets zero-shot COCO retrieval
Metrics: precision-at-1 10-shot accuracy finetune accuracy zero-shot accuracy recall TPU-v3 training time inference ms/img GFLOPs
Baselines: ViT Tokens Choice (V-MoE top-K) Experts Choice Batch Priority Routing LIT
Builds on: Vision Transformer V-MoE Expert Choice routing Batch Priority Routing
Closest papers by meaning
- A Provably Effective Method for Pruning Experts in Fine-tuned Sparse Mixture-of-Experts (ICML 2024 · similarity 0.91)
- Efficient Data Driven Mixture-of-Expert Extraction from Trained Networks (CVPR 2025 · similarity 0.90)
- Mixture of Tokens: Continuous MoE through Cross-Example Aggregation (NeurIPS 2024 · similarity 0.90)
- Dense Backpropagation Improves Training for Sparse Mixture-of-Experts (NeurIPS 2025 · similarity 0.89)
- Multilinear Mixture of Experts: Scalable Expert Specialization through Factorization (NeurIPS 2024 · similarity 0.89)
- Domain-Specific Pruning of Large Mixture-of-Experts Models with Few-shot Demonstrations (NeurIPS 2025 · similarity 0.89)
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.