Mixture-of-Depths: Dynamically allocating compute in transformer-based language models
David Raposo, Sam Ritter, Blake Richards, Timothy Lillicrap, Peter Conway Humphreys, Adam Santoro
cs.LG, cs.CL
2024-04-03
MoD routes 12.5% of tokens through every other block via expert-choice top-k. Under isoFLOP, loss matches a dense model, with ~60% faster steps and up to half the forward FLOPs.
A vanilla transformer spends the same FLOPs on every position in every layer. Tokens are not equal: a function word and a step that needs long-range reasoning do not need the same depth. Conditional computation tries to spend work only where it helps, but dynamic graphs and unknown tensor sizes fight current accelerators. Mixture-of-experts reroutes among MLPs and keeps total FLOPs roughly constant.
Mixture-of-Depths wants a static graph with a capacity fixed before training, while letting tokens take different paths through depth. Some layers run full self-attention and MLP; others pass the residual and skip that block's work.
Each MoD block feeds self-attention and the MLP only k tokens, with k set as capacity. A router emits a scalar per token; expert-choice keeps the top-k and residual-skips the rest. Expert-choice gives perfect load balance without an auxiliary balancing loss, and high-scoring tokens are guaranteed a seat. Router weights multiply the block output so routing sits on the gradient path.
Capacity below sequence length cuts the query-key product from T² to k². At k=T/2 that term is 25% of the dense cost. Total compute is known up front; which tokens fill the k slots still depends on context.
Top-k is non-causal: whether a token ranks in the top-k depends on later scores, which autoregressive decoding cannot see. Two workarounds. A binary cross-entropy on the router, using top-k labels, pushes scores to either side of 0.5 and raises language-modeling loss by about 0.2–0.3%. A small auxiliary MLP predicts membership without touching the main loss. That auxiliary task hits 99% accuracy quickly.
Unlike early-exit, a token can skip middle blocks and later attend to tokens that ran those blocks. Unlike MoE, the choice is compute versus skip, not which expert MLP, and routing applies to attention as well as the MLP.
Sequence length 2048, batch 128. A 6e18 FLOP sweep finds the best layout: route every other block, capacity 12.5% (256 of 2048 tokens). Routing every block hurts; stochastic routing loses badly to learned routing and to the dense baseline. Pushing capacity below 12.5% starts to degrade.
MoD drags the isoFLOP curve down and to the right: the optimum has lower loss and more parameters. A 220M MoD slightly beats the isoFLOP-optimal 220M dense model and steps about 60% faster, with similar wall-clock training time. At 2e19 and 1e20 FLOPs, models from 60M to 3B, the pattern holds: some MoD variants are both faster and lower-loss. A useful rule is to match the dense isoFLOP-optimal model's FLOPs per step, then add depth rather than width.
Autoregressive eval uses 256k held-out sequences, about 500M tokens. Switching from non-causal top-k at train time to the predictor at decode barely moves loss. MoD also stacks with MoE as MoDE, either staged (MoD then MoE) or integrated (a no-op expert). Both add MoD's gain on top of MoE. Integrated beats shrinking expert capacity and dropping overflow tokens, because the model learns to pick the skip on purpose.
This is a clean 2024 demonstration that per-token depth can live on a static graph. For latency, the smaller MoD models beside the isoFLOP optimum match dense optimal loss with fewer FLOPs per step and can step more than 50% faster at sampling. For training wall-clock, the paper holds training FLOPs and time fixed and returns a bigger, better net, or equal quality with a faster step.
It is orthogonal to MoE and can stack. Routing decides both who is updated and who is visible as keys, so the KV set is sparsified too, which is a different knob from MLP-only MoE.
Reported metrics are training loss and step speed, not MMLU-style downstream tables. Reading this as a proven general quality lift goes beyond the paper.
Non-causal top-k is structural. Decode depends on an auxiliary loss or predictor. 99% predictor accuracy sounds high; whether a 1% routing error compounds across depth on long sequences is not broken out. The "every other block, 12.5% capacity" recipe comes from a 6e18 sweep. Whether it remains optimal at larger scale is only hinted at by isoFLOP curves, not a fresh capacity sweep.
There is no public standard downstream eval, and no head-to-head against CoLT5, early-exit, or token merging on a shared latency axis. Memory and KV-cache savings are expected and not systematically measured. Training uses one hyperparameter family at length 2048; long context and expert-parallel clusters are untested.