Mixture of Depths: Teaching Transformers to Skip Work They Do Not Need

Every forward pass through a transformer applies the same amount of compute to every token, regardless of whether a given token needs it. Predicting the next word after 'The capital of France is' involves near-zero uncertainty; predicting the next word in the middle of a complex proof or an ambiguous sentence is harder. Standard transformer architectures treat both cases identically, spending the same FLOPs on every position in every layer. Mixture-of-Depths (MoD), introduced in 'Mixture-of-Depths: Dynamically allocating compute in transformer-based language models' (arXiv:2404.02258, April 2024) by David Raposo, Sam Ritter, Blake Richards, Timothy Lillicrap, Peter Conway Humphreys, and Adam Santoro at Google DeepMind, demonstrates that transformers can instead learn to dynamically route tokens: only a fraction of tokens participate in each block's self-attention and MLP computation, while the rest pass through a residual connection unchanged. The headline result is that MoD models can match the language modeling loss of isoFLOP-optimal baselines while requiring a fraction of the FLOPs per forward pass, and can be up to 50% faster during post-training sampling.

The uniform compute problem and the conditional computation idea

The Chinchilla scaling laws established how to allocate a fixed training FLOP budget between model size and training tokens. What neither Chinchilla nor the original transformer paper addressed is whether every token in every layer actually needs the same compute. Prior attempts at conditional computation for transformers often introduced dynamic computation graphs with variable tensor sizes, which are incompatible with the static graph assumptions that GPU and TPU hardware is designed to exploit. Early-exit methods (where a token stops processing after some layer) have dynamic depth, but require choosing a threshold and cannot easily let a token skip a middle layer while still being attended to by later tokens that have processed all layers. MoD takes a different approach: enforce a fixed, known capacity k at each routing block, and let the network learn which k tokens to process. Because k is set before training, the computation graph and all tensor sizes remain static. The only thing that changes at runtime is which k tokens fill those fixed-size compute slots.

How the top-k routing mechanism works

Each routing block in an MoD transformer contains a small scalar router alongside the standard self-attention and MLP. The router emits one scalar weight per token, expressing its preference for that token to participate in the full block computation versus routing around it via a residual connection. At each block, the top-k tokens by router weight are selected for full self-attention and MLP computation; all remaining tokens skip the block and their representation is passed through unchanged via the residual. The key constraint is that k is fixed before training as a fraction of the total sequence capacity T. Setting k to 12.5% of T means 87.5% of tokens route around a given block, spending no FLOPs on self-attention or MLP at that layer. Because k is constant, the training computation graph has known tensor sizes throughout, preserving hardware efficiency.

• Router: a lightweight scalar projection applied to each token's hidden state, producing a single weight used for top-k selection.
• Top-k selection: exactly k tokens per block participate in full self-attention and MLP; the rest take the residual path with zero additional FLOPs.
• Static graph: k is fixed before training, so all tensor sizes are known in advance and hardware utilization is predictable.
• Routing frequency: applying routing every other block (interleaving routing blocks with standard full-attention blocks) was crucial for strong performance.
• Optimal capacity: 12.5% capacity (87.5% of tokens routed around blocks) produced the best results; performance degraded below 12.5%.
• Stochastic routing fails: replacing the learned router with random top-k selection causes performance to drop drastically, confirming that the routing decisions themselves carry meaningful signal.

isoFLOP results and the 220M parameter finding

The paper evaluates MoD using isoFLOP comparisons: given a fixed total training FLOP budget, which model configuration achieves the lowest language modeling loss? The authors train models ranging from 60M to 3B parameters at three FLOP budgets: 6e18, 2e19, and 1e20. The core finding is that MoD transformers drag the isoFLOP optimal curve 'down and to the right': the optimal MoD model achieves a lower loss than the optimal baseline model and has more parameters, because the per-token compute savings free up budget to add model capacity. A concrete example from the paper: a 220M parameter MoD variant using the 12.5% capacity configuration slightly outperforms the isoFLOP-optimal 220M parameter baseline, while being up to 60% faster to step during training. Crucially, when measured in wall-clock time on equivalent hardware, both models take approximately the same time to train, because the faster per-step speed of the MoD model offsets the fact that it takes more optimizer steps to reach the same loss. The paper reports that MoD models can be up to 50% faster during post-training autoregressive sampling, where the compute savings directly translate to wall-clock throughput.

The autoregressive sampling problem and the predictor router

Top-k routing creates a complication for autoregressive inference. During training, the router sees the full sequence and selects the globally top-k tokens for each block; this is a non-causal operation, because deciding whether token 10 is in the top-k depends on comparing its router weight to tokens 11, 12, and beyond. During autoregressive generation, future tokens do not exist yet, making global top-k selection impossible. MoD addresses this with a predictor router: a small auxiliary network trained to predict, for each token, whether it would have been selected by the top-k mechanism at training time. The predictor is trained alongside the main model using the training-time top-k decisions as labels. The result reported in the paper is that the predictor achieves over 97% accuracy early in training, and switching from the training-time top-k router to the predictor-based router at inference time produces minimal performance degradation on the 256,000-sequence held-out evaluation set (500M tokens). The MoD models that outperform the isoFLOP-optimal baseline at training time continue to outperform it under the predictor-based autoregressive evaluation.

MoDE: combining depth routing with expert routing

The routing logic in MoD is architecturally compatible with Mixture-of-Experts (MoE) routing, and the paper explores Mixture-of-Depths-and-Experts (MoDE) as a natural extension. In MoE transformers, each token is routed to one of several expert MLPs rather than a single shared MLP; total compute per token stays roughly constant across the sequence, but each token activates a different parameter subset. In MoDE, the depth-skipping logic from MoD is layered on top of the expert-selection logic from MoE. The paper tests two variants: staged MoDE, which applies MoD routing before MoE routing so that skipped tokens bypass both attention and expert selection, and integrated MoDE, which treats the residual (no-op) path as one option among the experts so that a single routing step covers both. Both variants show performance improvements that compound the gains from MoD and MoE separately. The integrated variant is noted to outperform simply reducing expert capacity and relying on token dropping, because in the integrated variant tokens explicitly learn to choose the residual path rather than preferring an expert but getting dropped at capacity.

What MoD reveals about token-level difficulty in language modeling

One of the paper's qualitative findings concerns which tokens learn to engage with blocks. Analysis of a trained MoD model shows that some tokens consistently route to processing blocks throughout the model's depth, while others consistently take the residual path. Preliminary analysis in the paper finds that tokens that engage with blocks more frequently are correlated with output predictions that have higher entropy, a proxy for predictions that are harder to make. This suggests the network is learning something real about per-token difficulty rather than discovering an arbitrary routing pattern. The observation that high-entropy, high-uncertainty predictions receive more compute aligns with the intuition that a model should spend effort where it is needed. The MoD framework does not explicitly encode this preference; it emerges from training on the language modeling objective alone. The paper also notes that MoD opens a direction for decoupled routing of queries versus keys: a token could route to be a key (available to be attended to) without also paying the cost of computing its own query, which is a step toward more structured long-context memory mechanisms. The code and results for the experiments are associated with the arXiv submission arXiv:2404.02258. PromptingIndex covers MoD alongside PagedAttention, Multi-Head Latent Attention, speculative decoding, and Mixture of Experts as part of its series on inference efficiency techniques for transformer-based language models.