Source-linked AI summary
Mixture-of-Depths: Dynamically allocating compute in transformer-based language models
David Raposo, Sam Ritter, Blake Richards, Timothy Lillicrap, Peter Conway Humphreys, Adam Santoro
TL;DR
Transformers typically allocate the same computation to every token, motivating a more selective approach. Mixture-of-Depths learns top-k token routing under a fixed budget, skipping standard attention and MLP computation for some tokens. It preserves or improves training performance while using fewer forward-pass FLOPs, but top-k routing is non-causal during autoregressive sampling and therefore requires a predictive substitute.
Problem
Transformers spend the same compute per token even though different tokens and sequences may require different amounts of effort for accurate prediction.
Method
Mixture-of-Depths uses learned top-k routing under a fixed compute budget to decide which tokens participate in self-attention and MLP computations at each layer.
Results
MoD models match or improve baseline performance while using fewer forward-pass FLOPs, including up to 50% fewer FLOPs per forward pass and up to 1.5% higher final log probability under equivalent training FLOPs.
Takeaways & Limitations
MoD provides a way to tune compute per forward pass and inference time by learning which tokens should receive standard block computation.
Takeaways & Limitations
Top-k routing is non-causal during autoregressive sampling, so MoD requires an auxiliary classifier or router loss to mimic those decisions with minimal to no performance degradation.
Abstract
from arXiv · showhide
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.
1. Introduction
The paper addresses inefficient uniform computation in transformers by learning per-token, per-layer routing under a fixed compute budget. Mixture-of-Depths applies or skips standard attention and MLP computation dynamically, achieving comparable or improved performance with less computation.
- Transformers expend the same compute per token despite tokens and sequences requiring different amounts of effort for accurate prediction.
- Conditional computation can reduce total compute, but dynamic computation graphs and variable tensor sizes conflict with hardware optimized for static graphs.
- Mixture-of-Depths uses a fixed, user-defined compute budget while learning per-token, per-layer decisions about where to spend it.
- MoD routes tokens either through standard self-attention and MLP computation or through a residual connection, and applies this routing across model depth.
- Upwards of 50% fewer FLOPs per forward pass can achieve training-loss parity with an isoFLOP-optimal vanilla transformer, while equivalent-FLOP training can improve the final log probability objective by as much as 1.5%.
2. Background
The background situates Mixture-of-Depths within conditional computation, early exiting, adaptive-depth, token-merging, adapter-skipping, and mixture-of-experts approaches. MoD instead dynamically skips a single standard computation path while allowing tokens to bypass middle layers and later interact through attention.
- Conditional computation methods seek to reduce transformer computation by deciding when and how much processing each token receives.
- Early-exit methods stop computation for a token and skip its remaining layers, whereas MoD can skip middle layers before later updating the token.
- Related approaches include shared-weight adaptive iteration, inference-time token merging, and skipping frozen transformer blocks with fine-tuned adapters.
- MoD focuses on decoder-only transformers and proposes a predictive router for efficient inference despite the non-causal nature of top-k routing.
- MoD adapts MoE routing logic by dynamically skipping one expert computation rather than routing tokens among multiple expert MLPs.
3. Implementing Mixture-of-Depths Transformers
Mixture-of-Depths sets a fixed, reduced compute budget and learns which tokens receive expensive self-attention and MLP processing at each block. Expert-choice top-k routing sends selected tokens through computation and routes others through a residual path, preserving static tensor sizes while reducing FLOPs.
- 3.1. Defining a compute budget: A static compute budget limits the number of tokens participating in each block’s self-attention and MLP computations.The capacity is user-defined before training, so compute and tensor sizes remain predictable.
- 3.1. Defining a compute budget: Capacity, rather than realized routing outcomes, determines FLOPs in static-graph conditional computation.Inputs are padded to capacity or tokens are dropped when capacity is exceeded.
- 3.2. Routing around transformer blocks: Each token is routed either through self-attention and the MLP or around the block via a computationally cheap residual connection.Routing around the block leaves the token unchanged and saves computation.
- 3.2. Routing around transformer blocks: The model interpolates between vanilla-transformer quality and a fast but weak model by varying block capacities and routing decisions.Full capacity recovers a vanilla transformer, whereas zero capacity routes every token around every block.
- 3.3. Routing schemes: Learned routing is preferred to stochastic routing because the network can identify tokens that require more or less processing.The stochastic control significantly under-performs relative to vanilla transformers.
- 3.3. Routing schemes: Expert-choice routing selects the top-k router weights per block, avoids auxiliary balancing loss, and assigns computation to mutually exclusive token sets.With one computational path, tokens outside the selected set are routed away from self-attention and the MLP.
- 3.4. Routing implementation: A router produces a scalar per token, and a percentile threshold selects the user-defined capacity C tokens for block computation.The selected tokens form the input to self-attention and the subsequent MLP; fewer than S tokens reduce their computational cost.
- 3.5. Sampling: During autoregressive sampling, routing can use only the router’s current output, avoiding dependence on future tokens.The auxiliary routing task quickly achieves 99% accuracy.
4. Results
MoD models improve the speed–performance trade-off by routing only selected tokens through computation, with 12.5% capacity every other block emerging as the strongest configuration. Across isoFLOP budgets, MoD variants can match or outperform vanilla baselines while using fewer FLOPs per forward pass, and learned routing remains effective during autoregressive sampling.
- Training, isoFLOP comparisons: Upwards of 60% faster stepping was achieved by a 220M MoD variant that slightly outperformed an equally sized isoFLOP-optimal baseline.On equivalent hardware, both variants took approximately the same wall-clock time to train.
- Hyperparameter tuning: 12.5% capacity blocks routed every other block provided the best tested MoD configuration.The study found performance improved as capacity decreased to 12.5%, but degraded below that level.
- IsoFLOP scaling: MoD variants across 6e18, 2e19, and 1e20 training FLOPs could achieve lower loss while stepping faster than isoFLOP-optimal baselines.Lower per-forward-pass FLOPs were tightly correlated with faster wall-clock step times on equivalent hardware.
- IsoFLOP scaling: The optimal MoD model uses as many FLOPs per forward pass as the isoFLOP-optimal baseline, allowing model size to be tuned for a given routing configuration.This provides an empirical way to select the optimal model size for a fixed capacity and routing frequency.
- Routing analysis: Routing blocks process sparse token subsets, while some tokens are preferentially routed across depth and others frequently bypass blocks.Tokens engaging blocks more often were correlated with higher-entropy output predictions, which may indicate harder predictions.
- Mixture-of-Depths-and-Experts: MoD routing can be combined with MoE, and integrated MoDE performed better than reducing conventional MoE expert capacity with token dropping.Integrated MoDE explicitly learns the residual path, whereas token dropping can discard tokens that prefer an expert.
- Autoregressive evaluation: Switching from non-causal top-k routing during training to causal predictor-based routing during sampling caused little performance degradation.The predictor reached upwards of 97% accuracy soon into training, and MoD variants retained favorable performance–FLOPs trade-offs.
5. Discussion
Mixture-of-Depths improves the performance–compute trade-off by learning which tokens should receive computation, while supporting extensions across computation types and model architectures.
- Mixture-of-Depths can improve isoFLOP-optimal baseline performance while using fewer FLOPs per forward pass.This enables models that are both faster and better performing under a given training FLOP budget.
- Learned routing determines whether each token participates in self-attention and the subsequent MLP rather than dropping computation indiscriminately.Routing decisions select between computation and a residual path that leaves the token unchanged.
- Top-k routing is non-causal during training, so auxiliary classifiers or router losses learn to mimic its decisions for autoregressive sampling.The paper reports minimal to no performance degradation when applying the learned routing during sampling.
- MoD can combine with MoE through staged or integrated routing, including no-op operations alongside expert computations.The integrated approach uses one routing operation, while staged MoDE routes tokens around or toward blocks before self-attention.
- Decoupling query, key, and value routing could support a more efficient long-term memory mechanism for future self-attention.Tokens could be selected for retrieval once during memory encoding instead of undergoing full content-based lookup at every future step.
- Routing among computation types could extend beyond transformer blocks and null operations to functions such as memory lookup or tool use.Capacity constraints can offset the cost of introducing expensive computations by routing only a small number of tokens to them.