Mixture-of-Depths

Mixture-of-Depths (MoD) is a conditional-computation technique for transformer language models in which each layer decides whether to process a given token or skip it entirely, passing the token through an unchanged residual connection. The model still spends a fixed total amount of compute per sequence, but it chooses where along the sequence and across depth that compute goes, rather than treating every token identically.^[1]

Background

Standard transformers spend the same floating-point operations (FLOPs) on every token at every layer — a punctuation mark receives the same computation as the hardest word in the sentence. Researchers have long argued this is wasteful. The most successful prior answer is Mixture-of-Experts (MoE), introduced by Shazeer and colleagues in 2017, which routes each token to a small subset of expert feed-forward networks while keeping total compute constant.^[2]

Google DeepMind's Mixture-of-Depths paper (April 2024) applied a similar routing idea along a different axis. Instead of choosing which expert processes a token, MoD chooses whether the layer processes the token at all — a single expert that can be dynamically skipped.^[1]

How It Works

At each layer, a top-k routing mechanism scores the tokens in the current sequence and admits only the top k of them into the self-attention and MLP computations; the rest bypass the layer via the residual stream.^[1] Two properties make the method hardware-friendly:

• Fixed budget. Because k is defined in advance, tensor shapes stay static — unlike most conditional-computation schemes that create dynamic computation graphs, MoD keeps total FLOPs exactly predictable.^[1]
• Context-sensitive allocation. Although the number of processed tokens is constant, their identities shift with the input, so compute flows to whatever positions the network judges most deserving, at each depth.^[1]

The paper reported two equivalent operating points: at equal training FLOPs, MoD models improved final log probability by up to about 1.5% over vanilla transformers; alternatively, MoD could match baseline training loss while using a fraction of the FLOPs per forward pass — upwards of 50% fewer — making sampling steps correspondingly faster.^[3]

Why It Matters for AI

Inference cost is the binding constraint on putting large language models into production. Techniques like MoD, together with MoE and speculative decoding, attack that cost from different angles: MoE widens the model without proportionally widening per-token compute, while MoD thins out the computation for tokens that need less of it. The two are conceptually complementary — routing across expert sub-networks (the width dimension) and across layers (the depth dimension). MoD also opens a natural speed–quality trade-off knob: raising or lowering k moves a deployment along the curve between latency and accuracy.^[1]

The idea has influenced a broader research programme on dynamic and adaptive computation, where models decide how much reasoning effort a prompt deserves rather than applying one uniform budget to all inputs.

>Key Takeaways

• MoD lets a transformer skip entire layers for selected tokens while holding the total compute budget fixed, unlike most conditional-computation methods.^[1]
• A top-k router decides which tokens are processed at each layer; skipped tokens pass through unchanged via the residual connection.
• At equal training FLOPs it improved log-likelihood by up to ~1.5%; at equal quality it cut forward-pass FLOPs by upwards of 50%.^[3]
• It complements, rather than replaces, Mixture-of-Experts: MoD sparsifies across depth, MoE across width.
• For cost-sensitive deployments — including Malaysia's SME and public-sector AI use — such efficiency gains lower the unit cost of inference.

See Also

• Mixture-of-Experts
• Transformer Architecture
• Speculative Decoding
• Quantization

References