Research map / Architectures and efficiency / Mixture of experts
Mixture of Tokens: Continuous MoE through Cross-Example Aggregation
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 Mixture of Tokens, a continuous MoE where each expert processes a learned mixture of same-position tokens from different examples, enabling autoregressive training and a transition-tuning path to sparse MoE inference.
Problem
Sparse MoE models are discontinuous in parameters due to top-k routing, causing training instability, while existing continuous MoE designs either fail to scale or are incompatible with autoregressive decoding.
Approach
Each expert receives a softmax-weighted linear combination of tokens grouped across different examples at the same sequence position; each token's output is a mixing-weight-weighted combination of expert outputs. This keeps the layer fully continuous and avoids top-k, load-balancing losses, and token dropping, while remaining compatible with autoregressive training. Transition tuning then initiali
Key results
- MoT reaches the Transformer baseline's final loss in about one third of training steps, a stated 3x training speedup.
- MoT-Medium/32E matches Expert Choice in mixed precision (3.442 vs 3.420 loss) and beats it in bf16-only training (3.661 vs 3.728).
- Transition tuning initializes a Token Choice model from MoT weights and recovers performance with 15K extra steps (about 10% of pretraining).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use MoT instead of Expert Choice | GPT-2 Medium/32E · C4, bf16 only | loss: 3.728 ($\pm$ 0.044) -> 3.661 ($\pm$ 0.007) | Helped |
| Use MoT instead of Expert Choice | GPT-2 Medium/32E · C4, mixed precision | loss: 3.420 ($\pm$ 0.002) -> 3.442 ($\pm$ 0.002) | Hurt |
| Replace learned mixing weights with uniform weights (ablation) | MoT · C4 language modeling | loss | Hurt |
| Increase the number of mixtures processed by each expert (ablation) | MoT · C4 language modeling | loss | Helped |
| Increase the number of experts while keeping the compute budget the same (ablation) | MoT · C4 language modeling | loss | Helped |
| Use MoT instead of a dense Transformer | C4 language pretraining · C4 | compute budget to reach baseline final loss: 33% | Helped |
Surprising findings
- Expert Choice attains better loss than MoT in mixed precision but falls behind MoT when training is fully in bfloat16.
Limitations
- Does not support unbatched inference directly; requires transition tuning to a Token Choice model.
- Large memory footprint of MoT layers.
- Only evaluated at medium scales (77M and 162M dense-equivalent) with untested larger-model stability.
Experimental setup
Datasets: C4
Benchmarks: C4 validation perplexity/loss
Metrics: perplexity cross-entropy loss training compute to reach baseline loss
Baselines: Transformer dense baseline Token Choice MoE Expert Choice MoE uniform-mixing MoT ablation
Builds on: GPT-2 Token Choice MoE Expert Choice MoE Soft MoE
Closest papers by meaning
- From Sparse to Soft Mixtures of Experts (ICLR 2024 · similarity 0.90)
- Solving Token Gradient Conflict in Mixture-of-Experts for Large Vision-Language Model (ICLR 2025 · similarity 0.89)
- Pruning and Distilling Mixture-of-Experts into Dense Language Models (NeurIPS 2026 · similarity 0.89)
- Mixture Compressor for Mixture-of-Experts LLMs Gains More (ICLR 2025 · similarity 0.89)
- MoEUT: Mixture-of-Experts Universal Transformers (NeurIPS 2024 · similarity 0.89)
- MoE-CAP: Benchmarking Cost, Accuracy and Performance of Sparse Mixture-of-Experts Systems (NeurIPS 2025 · similarity 0.88)
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.