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

arXiv: https://arxiv.org/abs/2404.02258 · HTML: https://mlsanity.com/paper/2404.02258/

## Key facts

- Cited by 258 papers in the map (data through 14 Oct 2026).
- Builds on 15 of its 19 references that are in the map.
- First posted on arXiv 2 Apr 2024.

## Builds on (most cited first, 15 of 15)

- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding](https://mlsanity.com/paper/2006.16668/) (2020)
- [Towards a Unified View of Parameter-Efficient Transfer Learning](https://mlsanity.com/paper/2110.04366/) (2021)
- [Efficient Transformers: A Survey](https://mlsanity.com/paper/2009.06732/) (2020)
- [Adaptive Computation Time for Recurrent Neural Networks](https://mlsanity.com/paper/1603.08983/) (2016)
- [SkipNet: Learning Dynamic Routing in Convolutional Networks](https://mlsanity.com/paper/1711.09485/) (2017)
- [Depth-Adaptive Transformer](https://mlsanity.com/paper/1910.10073/) (2019)
- [Conditional Computation in Neural Networks for faster models](https://mlsanity.com/paper/1511.06297/) (2015)
- [Deep Learning of Representations: Looking Forward](https://mlsanity.com/paper/1305.0445/) (2013)
- [Compression of Deep Learning Models for Text: A Survey](https://mlsanity.com/paper/2008.05221/) (2020)
- [Controlling Computation versus Quality for Neural Sequence Models](https://mlsanity.com/paper/2002.07106/) (2020)
- [Exponentially Increasing the Capacity-to-Computation Ratio for Conditional Computation in Deep Learning](https://mlsanity.com/paper/1406.7362/) (2014)
- [Anytime Dense Prediction with Confidence Adaptivity](https://mlsanity.com/paper/2104.00749/) (2021)
- [Variable Computation in Recurrent Neural Networks](https://mlsanity.com/paper/1611.06188/) (2016)

## Built on by (most cited first, 30 of 258)

- [LayerSkip: Enabling Early Exit Inference and Self-Speculative Decoding](https://mlsanity.com/paper/2404.16710/) (2024)
- [OLMoE: Open Mixture-of-Experts Language Models](https://mlsanity.com/paper/2409.02060/) (2024)
- [A Survey on Mixture of Experts in Large Language Models](https://mlsanity.com/paper/2407.06204/) (2024)
- [ShowUI: One Vision-Language-Action Model for GUI Visual Agent](https://mlsanity.com/paper/2411.17465/) (2024)
- [Mixture-of-Transformers: A Sparse and Scalable Architecture for Multi-Modal Foundation Models](https://mlsanity.com/paper/2411.04996/) (2024)
- [Learning-to-Cache: Accelerating Diffusion Transformer via Layer Caching](https://mlsanity.com/paper/2406.01733/) (2024)
- [From Decoding to Meta-Generation: Inference-time Algorithms for Large Language Models](https://mlsanity.com/paper/2406.16838/) (2024)
- [Mixture-of-Recursions: Learning Dynamic Recursive Depths for Adaptive Token-Level Computation](https://mlsanity.com/paper/2507.10524/) (2025)
- [Why do LLMs attend to the first token?](https://mlsanity.com/paper/2504.02732/) (2025)
- [MiniCache: KV Cache Compression in Depth Dimension for Large Language Models](https://mlsanity.com/paper/2405.14366/) (2024)
- [DeeR-VLA: Dynamic Inference of Multimodal Large Language Models for Efficient Robot Execution](https://mlsanity.com/paper/2411.02359/) (2024)
- [Mono-InternVL: Pushing the Boundaries of Monolithic Multimodal Large Language Models with Endogenous Visual Pre-training](https://mlsanity.com/paper/2410.08202/) (2024)
- [Relaxed Recursive Transformers: Effective Parameter Sharing with Layer-wise LoRA](https://mlsanity.com/paper/2410.20672/) (2024)
- [Complexity Experts are Task-Discriminative Learners for Any Image Restoration](https://mlsanity.com/paper/2411.18466/) (2024)
- [CogVLA: Cognition-Aligned Vision-Language-Action Model via Instruction-Driven Routing & Sparsification](https://mlsanity.com/paper/2508.21046/) (2025)
- [Dynamic Chunking for End-to-End Hierarchical Sequence Modeling](https://mlsanity.com/paper/2507.07955/) (2025)
- [MoLe-VLA: Dynamic Layer-skipping Vision Language Action Model via Mixture-of-Layers for Efficient Robot Manipulation](https://mlsanity.com/paper/2503.20384/) (2025)
- [MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts](https://mlsanity.com/paper/2407.21770/) (2024)
- [Online Video Understanding: OVBench and VideoChat-Online](https://mlsanity.com/paper/2501.00584/) (2024)
- [LoopFormer: Elastic-Depth Looped Transformers for Latent Reasoning via Shortcut Modulation](https://mlsanity.com/paper/2602.11451/) (2026)
- [LION-FS: Fast & Slow Video-Language Thinker as Online Video Assistant](https://mlsanity.com/paper/2503.03663/) (2025)
- [Inner Thinking Transformer: Leveraging Dynamic Depth Scaling to Foster Adaptive Internal Thinking](https://mlsanity.com/paper/2502.13842/) (2025)
- [VideoLLM-MoD: Efficient Video-Language Streaming with Mixture-of-Depths Vision Computation](https://mlsanity.com/paper/2408.16730/) (2024)
- [Parcae: Scaling Laws For Stable Looped Language Models](https://mlsanity.com/paper/2604.12946/) (2026)
- [Demons in the Detail: On Implementing Load Balancing Loss for Training Specialized Mixture-of-Expert Models](https://mlsanity.com/paper/2501.11873/) (2025)
- [Teaching Pretrained Language Models to Think Deeper with Retrofitted Recurrence](https://mlsanity.com/paper/2511.07384/) (2025)
- [TinyFusion: Diffusion Transformers Learned Shallow](https://mlsanity.com/paper/2412.01199/) (2024)
- [AdaMoE: Token-Adaptive Routing with Null Experts for Mixture-of-Experts Language Models](https://mlsanity.com/paper/2406.13233/) (2024)
- [Memory3: Language Modeling with Explicit Memory](https://mlsanity.com/paper/2407.01178/) (2024)
- [Large Language Model Inference Acceleration: A Comprehensive Hardware Perspective](https://mlsanity.com/paper/2410.04466/) (2024)

## Abstract (arXiv)

Transformer-based language models spread FLOPs uniformly across input sequences. In this work we demonstrate that transformers can instead learn to dynamically allocate FLOPs (or compute) to specific positions in a sequence, optimising the allocation along the sequence for different layers across the model depth. Our method enforces a total compute budget by capping the number of tokens (k) that can participate in the self-attention and MLP computations at a given layer. The tokens to be processed are determined by the network using a top-k routing mechanism. Since k is defined a priori, this simple procedure uses a static computation graph with known tensor sizes, unlike other conditional computation techniques. Nevertheless, since the identities of the k tokens are fluid, this method can expend FLOPs non-uniformly across the time and model depth dimensions. Thus, compute expenditure is entirely predictable in sum total, but dynamic and context-sensitive at the token-level. Not only do models trained in this way learn to dynamically allocate compute, they do so efficiently. These models match baseline performance for equivalent FLOPS and wall-clock times to train, but require a fraction of the FLOPs per forward pass, and can be upwards of 50% faster to step during post-training sampling.