---
title: "Mixture-of-Depths — paper by Raposo et al. (2024)"
description: "Standard transformers spend the same compute on every token, even though some… 97% accuracy. 24 references, cited by 53 papers here."
authors: ["David Raposo", "Sam Ritter", "Blake Richards", "Timothy Lillicrap", "Peter Conway Humphreys", "Adam Santoro"]
year: 2024
venue: "arXiv:2404.02258"
source: "https://arxiv.org/abs/2404.02258"
doi: "10.48550/arxiv.2404.02258"
arxiv: "2404.02258"
licence: "CC BY 4.0 (arXiv)"
licence_url: "http://creativecommons.org/licenses/by/4.0/"
summary_licence: "Elamur Cosmic License"
canonical: "https://research.elamur.ai/free/en/papers/raposo-2024-mixture-depths-dynamically-allocating-compute/"
---
[Elamur Research](https://md.research.elamur.ai/free/en/) · [this page in HTML](https://research.elamur.ai/free/en/papers/raposo-2024-mixture-depths-dynamically-allocating-compute/)

[home](https://md.research.elamur.ai/free/en/) / [papers](https://md.research.elamur.ai/free/en/papers/) / [2024](https://md.research.elamur.ai/free/en/papers/2024/page/7/) / Mixture-of-Depths: Dynamically allocating compute in transformer-based language models

# Mixture-of-Depths: Dynamically allocating compute in transformer-based language models

**Authors:** [David Raposo](https://md.research.elamur.ai/free/en/authors/david-raposo/), Sam Ritter, Blake Richards, [Timothy Lillicrap](https://md.research.elamur.ai/free/en/authors/timothy-lillicrap/), Peter Conway Humphreys and [Adam Santoro](https://md.research.elamur.ai/free/en/authors/adam-santoro/) · **Year:** 2024 · **Venue:** arXiv:2404.02258 · **Licence:** [CC BY 4.0 (arXiv)](http://creativecommons.org/licenses/by/4.0/)

**Links:** [Publisher's page](https://arxiv.org/abs/2404.02258) · [PDF](https://arxiv.org/pdf/2404.02258) · [DOI 10.48550/arxiv.2404.02258](https://doi.org/10.48550/arxiv.2404.02258) · [arXiv 2404.02258](https://arxiv.org/abs/2404.02258) · [OpenAlex](https://openalex.org/W4393967974)

**Citations:** 4 citations (OpenAlex, 2026-10-14) · [53](https://md.research.elamur.ai/free/en/papers/raposo-2024-mixture-depths-dynamically-allocating-compute/cited-by/) in our library, of the 149,439 papers whose bibliographies we read

## Summary
Standard transformers spend the same compute on every token, even though some predictions are much easier than others. Raposo et al. (Google DeepMind, McGill/Mila) propose Mixture-of-Depths (MoD), a conditional-computation method for decoder-only language models. Most such methods produce dynamic computation graphs that suit current hardware poorly. MoD instead fixes a compute budget before training by capping each block's capacity, meaning the number of tokens allowed into its self-attention and MLP. A learned per-block router scores tokens. The top-k tokens are processed, and the rest skip the block through the residual connection. Because k is fixed, tensor sizes stay static, yet which tokens get compute varies with context and with depth. Top-k selection uses future tokens, so for autoregressive generation the authors add a causal predictor (an auxiliary classifier or router loss) that imitates it. MoD models either beat isoFLOP-optimal baselines or match them while using far fewer FLOPs per forward pass. MoD also combines with MoE. The work shows that vanilla transformers waste compute and that learned routing can recover it cheaply.

## Key points
- Routing covers both the MLP and self-attention. Skipped tokens are therefore also unavailable as keys for other tokens. Unlike early-exit methods, a token can bypass middle blocks and still be processed by later ones.
- At equal training FLOPs and equal wall-clock training time, MoD improves the final log-probability objective by up to 1.5% over vanilla transformers.
- At loss parity with an isoFLOP-optimal baseline, MoD needs only a fraction of the per-forward-pass FLOPs and can step upwards of 50% faster during sampling.
- With capacity set to T/2, the query-key product in attention costs about 25% of the vanilla cost. Replacing non-causal top-k with the causal predictor at inference barely hurts quality, and the predictor exceeds 97% accuracy early in training.
- MoDE (Mixture-of-Depths-and-Experts) comes in two variants:
  - **staged**: MoD routing first, then MoE;
  - **integrated**: a single router chooses between experts and a no-op path.

  The integrated variant outperforms MoE models that rely on token dropping.

*Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.*

## Institutions

- [McGill University](https://md.research.elamur.ai/free/en/universities/mcgill-university/) — University, Montreal, Canada · 97 works · 11,755 citations

📚 Shelves (5): [🖥️ Computer Science](https://md.research.elamur.ai/free/en/shelves/computer-science/) · 19,503 works · [🤖 Artificial Intelligence](https://md.research.elamur.ai/free/en/shelves/computer-science/artificial-intelligence/) · 18,805 works · [🎭 Arts and Humanities](https://md.research.elamur.ai/free/en/shelves/arts-and-humanities/) · 17,959 works · [🗣️ Language and Linguistics](https://md.research.elamur.ai/free/en/shelves/arts-and-humanities/language-and-linguistics/) · 17,789 works · [👁️ Computer Vision and Pattern Recognition](https://md.research.elamur.ai/free/en/shelves/computer-science/computer-vision-and-pattern-recognition/) · 3,621 works

🧭 Topics (1): [🧠 Language models & pretraining](https://md.research.elamur.ai/free/en/topics/language-models/) · 4,145 works

## Cite

Cite as: David Raposo, Sam Ritter, Blake Richards, Timothy Lillicrap, Peter Conway Humphreys and Adam Santoro. 2024. Mixture-of-Depths: Dynamically allocating compute in transformer-based language models. arXiv:2404.02258.

The paper itself, not this page. BibTeX from the library's record of the work.

```bibtex
@misc{raposo2024mixture,
  title = {{Mixture-of-Depths: Dynamically allocating compute in transformer-based language models}},
  author = {David Raposo, Sam Ritter, Blake Richards, Timothy Lillicrap, Peter Conway Humphreys and Adam Santoro},
  year = {2024},
  doi = {10.48550/arxiv.2404.02258},
  eprint = {2404.02258},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2404.02258}
}
```

- [← previous HelpSteer2-Preference: Complementing Ratings with Preferences](https://md.research.elamur.ai/free/en/papers/wang-2024-helpsteer2-preference-complementing-ratings-preferences/)
- [next → Reducing Transformer Key-Value Cache Size with Cross-Layer Attention](https://md.research.elamur.ai/free/en/papers/brandon-2024-reducing-transformer-key-value-cache/)
- [↑ back to papers of 2024](https://md.research.elamur.ai/free/en/papers/2024/page/7/)

## References

24 references, 4 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

1. J. Ainslie, T. Lei, M. de Jong, S. Ontañón, S. Brahma, Y. Zemlyanskiy, D. Uthus, M. Guo, J. Lee-Thorp, Y. Tay, Y.-H. Sung, and S. Sanghai. [Colt5: Faster long-range transformers with conditional computation](https://md.research.elamur.ai/free/en/papers/ainslie-2023-colt5-faster-long-range-transformers/), 2023.
2. A. Bapna, N. Arivazhagan, and O. Firat. Controlling computation versus quality for neural sequence models. CoRR, abs/2002.07106, 2020. URL https://arxiv.org/abs/2002.07106 . — [arXiv:2002.07106](https://arxiv.org/abs/2002.07106) [↗](https://arxiv.org/abs/2002.07106)
3. E. Bengio, P.-L. Bacon, J. Pineau, and D. Precup. Conditional computation in neural networks for faster models, 2016. [find](https://scholar.google.com/scholar?q=E%2E+Bengio%2C+P%2E%2DL%2E+Bacon%2C+J%2E+Pineau%2C+and+D%2E+Precup%2E+Conditional+computation+in+neural+networks+for+faster+models%2C+2016%2E)
4. Y. Bengio. Deep learning of representations: Looking forward, 2013. [find](https://scholar.google.com/scholar?q=Y%2E+Bengio%2E+Deep+learning+of+representations%3A+Looking+forward%2C+2013%2E)
5. Y. Bengio, N. Léonard, and A. Courville. Estimating or propagating gradients through stochastic neurons for conditional computation, 2013. [find](https://scholar.google.com/scholar?q=Y%2E+Bengio%2C+N%2E+L%C3%A9onard%2C+and+A%2E+Courville%2E+Estimating+or+propagating+gradients+through+stochastic+neurons+for+conditional+computation%2C+2013%2E)
6. D. Bolya, C.-Y. Fu, X. Dai, P. Zhang, C. Feichtenhofer, and J. Hoffman. Token merging: Your vit but faster, 2023. [find](https://scholar.google.com/scholar?q=D%2E+Bolya%2C+C%2E%2DY%2E+Fu%2C+X%2E+Dai%2C+P%2E+Zhang%2C+C%2E+Feichtenhofer%2C+and+J%2E+Hoffman%2E+Token+merging%3A+Your+vit+but+faster%2C+2023%2E)
7. K. Cho and Y. Bengio. Exponentially increasing the capacity-to-computation ratio for conditional computation in deep learning, 2014. [find](https://scholar.google.com/scholar?q=K%2E+Cho+and+Y%2E+Bengio%2E+Exponentially+increasing+the+capacity%2Dto%2Dcomputation+ratio+for+conditional+computation+in+deep+learning%2C+2014%2E)
8. M. Dehghani, S. Gouws, O. Vinyals, J. Uszkoreit, and Ł. Kaiser. Universal transformers. arXiv preprint arXiv:1807.03819, 2018. [find](https://scholar.google.com/scholar?q=M%2E+Dehghani%2C+S%2E+Gouws%2C+O%2E+Vinyals%2C+J%2E+Uszkoreit%2C+and+%C5%81%2E+Kaiser%2E+Universal+transformers%2E+arXiv+preprint+arXiv%3A1807%2E03819%2C+2018%2E)
9. M. Elbayad, J. Gu, E. Grave, and M. Auli. Depth-adaptive transformer. CoRR, abs/1910.10073, 2019. URL http://arxiv.org/abs/1910.10073 . — [arXiv:1910.10073](https://arxiv.org/abs/1910.10073) [↗](http://arxiv.org/abs/1910.10073)
10. W. Fedus, B. Zoph, and N. Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity, 2022. [find](https://scholar.google.com/scholar?q=W%2E+Fedus%2C+B%2E+Zoph%2C+and+N%2E+Shazeer%2E+Switch+transformers%3A+Scaling+to+trillion+parameter+models+with+simple+and+efficient+sparsity%2C+2022%2E)
11. A. Graves. Adaptive computation time for recurrent neural networks. CoRR, abs/1603.08983, 2016. URL http://arxiv.org/abs/1603.08983 . [↗](http://arxiv.org/abs/1603.08983)
12. M. Guo, J. Ainslie, D. Uthus, S. Ontanon, J. Ni, Y.-H. Sung, and Y. Yang. [Longt5: Efficient text-to-text transformer for long sequences](https://md.research.elamur.ai/free/en/papers/guo-2022-longt5-efficient-text-text-transformer/), 2022.
13. M. Gupta and P. Agrawal. Compression of deep learning models for text: A survey, 2021. [find](https://scholar.google.com/scholar?q=M%2E+Gupta+and+P%2E+Agrawal%2E+Compression+of+deep+learning+models+for+text%3A+A+survey%2C+2021%2E)
14. J. He, C. Zhou, X. Ma, T. Berg-Kirkpatrick, and G. Neubig. [Towards a unified view of parameter-efficient transfer learning](https://md.research.elamur.ai/free/en/papers/he-2021-towards-unified-view-parameter-efficient/). arXiv preprint arXiv:2110.04366, 2021.
15. Y. Jernite, E. Grave, A. Joulin, and T. Mikolov. Variable computation in recurrent neural networks, 2017. [find](https://scholar.google.com/scholar?q=Y%2E+Jernite%2C+E%2E+Grave%2C+A%2E+Joulin%2C+and+T%2E+Mikolov%2E+Variable+computation+in+recurrent+neural+networks%2C+2017%2E)
16. T. Lei, J. Bai, S. Brahma, J. Ainslie, K. Lee, Y. Zhou, N. Du, V. Y. Zhao, Y. Wu, B. Li, Y. Zhang, and M.-W. Chang. Conditional adapters: Parameter-efficient transfer learning with fast inference, 2023. [find](https://scholar.google.com/scholar?q=T%2E+Lei%2C+J%2E+Bai%2C+S%2E+Brahma%2C+J%2E+Ainslie%2C+K%2E+Lee%2C+Y%2E+Zhou%2C+N%2E+Du%2C+V%2E+Y%2E+Zhao%2C+Y%2E+Wu%2C+B%2E+Li%2C+Y%2E+Zhang%2C+and+M%2E%2DW%2E+Chang%2E+Conditional+adapters%3A+Parameter%2Defficient+transfer+learning+with+fast+inference%2C+202)
17. D. Lepikhin, H. Lee, Y. Xu, D. Chen, O. Firat, Y. Huang, M. Krikun, N. Shazeer, and Z. Chen. Gshard: Scaling giant models with conditional computation and automatic sharding. arXiv preprint arXiv:2006.16668, 2020. [find](https://scholar.google.com/scholar?q=D%2E+Lepikhin%2C+H%2E+Lee%2C+Y%2E+Xu%2C+D%2E+Chen%2C+O%2E+Firat%2C+Y%2E+Huang%2C+M%2E+Krikun%2C+N%2E+Shazeer%2C+and+Z%2E+Chen%2E+Gshard%3A+Scaling+giant+models+with+conditional+computation+and+automatic+sharding%2E+arXiv+preprint+arXiv%3A2006)
18. Z. Liu, Z. Xu, H.-J. Wang, T. Darrell, and E. Shelhamer. Anytime dense prediction with confidence adaptivity. arXiv preprint arXiv:2104.00749, 2021. [find](https://scholar.google.com/scholar?q=Z%2E+Liu%2C+Z%2E+Xu%2C+H%2E%2DJ%2E+Wang%2C+T%2E+Darrell%2C+and+E%2E+Shelhamer%2E+Anytime+dense+prediction+with+confidence+adaptivity%2E+arXiv+preprint+arXiv%3A2104%2E00749%2C+2021%2E)
19. T. Schuster, A. Fisch, J. Gupta, M. Dehghani, D. Bahri, V. Q. Tran, Y. Tay, and D. Metzler. Confident adaptive language modeling, 2022. [find](https://scholar.google.com/scholar?q=T%2E+Schuster%2C+A%2E+Fisch%2C+J%2E+Gupta%2C+M%2E+Dehghani%2C+D%2E+Bahri%2C+V%2E+Q%2E+Tran%2C+Y%2E+Tay%2C+and+D%2E+Metzler%2E+Confident+adaptive+language+modeling%2C+2022%2E)
20. N. Shazeer, A. Mirhoseini, K. Maziarz, A. Davis, Q. Le, G. Hinton, and J. Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. arXiv preprint arXiv:1701.06538, 2017. [find](https://scholar.google.com/scholar?q=N%2E+Shazeer%2C+A%2E+Mirhoseini%2C+K%2E+Maziarz%2C+A%2E+Davis%2C+Q%2E+Le%2C+G%2E+Hinton%2C+and+J%2E+Dean%2E+Outrageously+large+neural+networks%3A+The+sparsely%2Dgated+mixture%2Dof%2Dexperts+layer%2E+arXiv+preprint+arXiv%3A1701%2E06538%2C+2017%2E)
21. A. Simoulin and B. Crabbé. How many layers and why? An analysis of the model depth in transformers. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing: Student Research Workshop, pages… [find](https://scholar.google.com/scholar?q=A%2E+Simoulin+and+B%2E+Crabb%C3%A9%2E+How+many+layers+and+why%3F+An+analysis+of+the+model+depth+in+transformers%2E+In+Proceedings+of+the+59th+Annual+Meeting+of+the+Association+for+Computational+Linguistics+and+the+1)
22. Y. Tay, M. Dehghani, D. Bahri, and D. Metzler. [Efficient transformers: A survey](https://md.research.elamur.ai/free/en/papers/tay-2020-efficient-transformers-survey/). CoRR, abs/2009.06732, 2020. URL https://arxiv.org/abs/2009.06732 . [↗](https://arxiv.org/abs/2009.06732)
23. X. Wang, F. Yu, Z. Dou, and J. E. Gonzalez. Skipnet: Learning dynamic routing in convolutional networks. CoRR, abs/1711.09485, 2017. URL http://arxiv.org/abs/1711.09485 . [↗](http://arxiv.org/abs/1711.09485)
24. B. Zoph, I. Bello, S. Kumar, N. Du, Y. Huang, J. Dean, N. Shazeer, and W. Fedus. St-moe: Designing stable and transferable sparse expert models, 2022. [find](https://scholar.google.com/scholar?q=B%2E+Zoph%2C+I%2E+Bello%2C+S%2E+Kumar%2C+N%2E+Du%2C+Y%2E+Huang%2C+J%2E+Dean%2C+N%2E+Shazeer%2C+and+W%2E+Fedus%2E+St%2Dmoe%3A+Designing+stable+and+transferable+sparse+expert+models%2C+2022%2E)

*Updated 3 November 2026*