Elamur Research⌕
☰
home / papers / 2024 / Mixture-of-Depths: Dynamically allocating compute in transformer-based language models

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

Paper · 2024 · David Raposo, Sam Ritter, Blake Richards et al.

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 UniversityMontreal · Canada97 works · 11,755 citations

 Fields, topics and concepts

📚 Shelves 5The field of science this work belongs to (OpenAlex fields).Computer Science · 19,503 worksArtificial Intelligence · 18,805 worksArts and Humanities · 17,959 worksLanguage and Linguistics · 17,789 worksComputer Vision and Pattern Recognition · 3,621 works

🧭 Topics 1Research subjects and languages the work is about.Language models & pretraining · 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. DOI 10.48550/arxiv.2404.02258 · arXiv 2404.02258

@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}
}

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, 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) ↗
  3. E. Bengio, P.-L. Bacon, J. Pineau, and D. Precup. Conditional computation in neural networks for faster models, 2016. find
  4. Y. Bengio. Deep learning of representations: Looking forward, 2013. find
  5. Y. Bengio, N. Léonard, and A. Courville. Estimating or propagating gradients through stochastic neurons for conditional computation, 2013. find
  6. D. Bolya, C.-Y. Fu, X. Dai, P. Zhang, C. Feichtenhofer, and J. Hoffman. Token merging: Your vit but faster, 2023. find
  7. K. Cho and Y. Bengio. Exponentially increasing the capacity-to-computation ratio for conditional computation in deep learning, 2014. find
  8. M. Dehghani, S. Gouws, O. Vinyals, J. Uszkoreit, and Ł. Kaiser. Universal transformers. arXiv preprint arXiv:1807.03819, 2018. find
  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) ↗
  10. W. Fedus, B. Zoph, and N. Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity, 2022. find
  11. A. Graves. Adaptive computation time for recurrent neural networks. CoRR, abs/1603.08983, 2016. URL 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, 2022.
  13. M. Gupta and P. Agrawal. Compression of deep learning models for text: A survey, 2021. find
  14. J. He, C. Zhou, X. Ma, T. Berg-Kirkpatrick, and G. Neubig. Towards a unified view of parameter-efficient transfer learning. arXiv preprint arXiv:2110.04366, 2021.
  15. Y. Jernite, E. Grave, A. Joulin, and T. Mikolov. Variable computation in recurrent neural networks, 2017. find
  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
  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
  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
  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
  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
  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
  22. Y. Tay, M. Dehghani, D. Bahri, and D. Metzler. Efficient transformers: A survey. CoRR, abs/2009.06732, 2020. URL 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 . ↗
  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

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/raposo-2024-mixture-depths-dynamically-allocating-compute/