Mixture-of-Depths: What Happens When a Transformer Skips Its Own Layers
Learn how Mixture-of-Depths lets transformers skip layers per token, cutting inference FLOPs. Compare MoD vs. MoE and early exit, plus routing and capacity tips.
The Transformer That Skips Itself
Most transformer models run every token through every layer. A period gets the same compute budget as a rare technical term. That uniformity is simple to build and easy to reason about—and it is increasingly expensive at scale.
As inference becomes a first-order line item in AI budgets, the question shifts from "how big can we make the model" to "how much of the model does each token actually need?" Mixture-of-Depths (MoD) is one answer: a learned router that lets a transformer decide, token by token, whether to process a block or skip it entirely.
The important framing is that MoD is not a new model family. It is a conditional compute technique that retrofits into existing transformer architectures, reallocating a fixed compute budget rather than adding parameters. The payoff is real but configuration-dependent efficiency at comparable quality; the catch is that routing, capacity tuning, and autoregressive decoding all become new engineering problems. This article covers how MoD works, where it differs from Mixture-of-Experts and early exit, and how to evaluate it for your stack. (See also: conditional compute techniques for LLM inference)
What Mixture-of-Depths Actually Does
At its core, MoD adds a single decision to each transformer block: does this token go through the block, or around it?
A small, learned router assigns each token either to the standard transformer block (attention plus MLP) or to a residual skip connection that carries the token's representation forward unchanged. The router is a lightweight linear projection that produces a scalar weight per token—one number expressing "how much does this token benefit from this block?" The top-k of those weights determines selection.
The core insight is that not every token needs the same depth of processing. Punctuation, function words, and predictable continuations rarely need a full block's worth of attention and feed-forward computation. Content-bearing tokens—rare terms, ambiguous references, reasoning steps—often do. MoD lets the model learn that distinction instead of assuming uniformity. (Related: how attention and feed-forward layers split work in a transformer)
Critically, MoD makes compute per block predictable and controllable. Each MoD block processes a fixed fraction of the sequence's tokens, set by a capacity budget, so total FLOPs are set by the capacity configuration rather than by the data. The router itself adds a small overhead—one linear projection per token per MoD block—which must be counted in any FLOP comparison. This is what separates MoD from pruning or distillation: you are not making the model smaller, you are making it selective about where it spends effort. (Compare: model pruning vs. distillation for inference cost reduction)
In the original formulation, MoD is applied to a subset of blocks interleaved with standard, always-on transformer blocks. The headline result is a depth-vs-width trade: at a matched FLOP budget, an MoD-transformer can be built deeper than a dense baseline and match its quality, or match quality with fewer FLOPs per forward pass. The gains are real but modest and benchmark- and scale-dependent—MoD is not a free lunch, and at large scale it has not consistently outperformed Mixture-of-Experts on quality-per-FLOP. Treat published numbers as configuration-specific, not as a general efficiency multiplier.
The Routing Mechanism, Step by Step
The forward pass through an MoD block proceeds in five steps:
- Score tokens. The router computes a scalar weight for each token in the sequence via a linear projection.
- Select top-k. Tokens with the highest router weights—up to the block's capacity—are selected to enter the block; the rest are not.
- Route through the block. Selected tokens pass through attention and MLP as usual.
- Skip the rest. Unselected tokens travel the residual connection, unchanged.
- Recombine. Selected tokens receive the block output added to their residual stream; unselected tokens pass through with the residual only. Routing is a hard top-k selection, not a continuous gate on the block output. (Soft-gated variants exist, but they are not the baseline MoD mechanism and change the auxiliary-loss design.)
An auxiliary loss keeps the router honest. Without it, the router can collapse—assigning nearly all tokens to one path—which defeats the purpose and destabilizes training. In the original work this is a binary cross-entropy loss that pushes the fraction of tokens routed into each MoD block toward the target capacity p. The loss coefficient is a real hyperparameter: too low and the router drifts off-capacity, too high and it dominates the task loss.
Capacity and the Top-k Constraint
MoD introduces a capacity factor: a fixed fraction of the sequence's tokens that may enter each MoD block. This is analogous in spirit to capacity in Mixture-of-Experts, but the units differ—MoE capacity is typically per-expert per-batch, while MoD capacity is a per-block fraction of the sequence. If a sequence has 1,024 tokens and capacity is set to 12.5%, roughly 128 tokens enter the block and the rest skip it.
Capacity is the primary tunable knob trading quality against speed. Lower capacity means more skipping and cheaper inference, but risks starving tokens that genuinely need depth. The original paper sweeps capacity and loss coefficient jointly; both interact with model depth and must be tuned together, not independently.
One deployment consideration deserves early attention: naive top-k selection is non-causal. Selecting the top-k tokens across a full sequence requires knowing future tokens, which breaks autoregressive decoding. The original paper acknowledges this as an open limitation and proposes auxiliary predictors or causal-masking workarounds; production-grade causal variants remain an active engineering problem, not a solved one. Budget for this constraint in any deployment plan. (Background: autoregressive decoding and KV cache tradeoffs)
How MoD Differs From Mixture-of-Experts and Early Exit
MoD sits inside a broader family of conditional compute techniques, and it is routinely confused with its neighbors. The distinctions matter because they imply different costs, different infrastructure, and different failure modes.
| Technique | What varies per token | Parameters | Primary cost |
|---|---|---|---|
| Dense transformer | Nothing | Fixed | Uniform compute |
| Mixture-of-Experts | Which experts process the token | Grows with expert count | Memory, load balancing |
| Early exit | When the token leaves the network | Fixed | Exit criteria, calibration |
| Mixture-of-Depths | Whether a block runs at all | Fixed | Routing, capacity tuning, causal selection |
The short version: MoE routes over width, early exit routes over exit timing, and MoD routes over depth.
MoD vs. MoE: Routing Over Depth, Not Width
Mixture-of-Experts routes each token to a subset of expert feed-forward networks. It increases parameter count (memory footprint grows with expert count) while holding per-token FLOPs roughly constant. MoD does the opposite: parameter count is unchanged, and the lever is how many tokens a given block processes. The two techniques are orthogonal and can be combined—an MoE layer can itself be gated by an MoD router—but the operational profiles differ sharply. MoE's hard problems are expert-parallel memory and load balancing; MoD's are router collapse, capacity tuning, and causal selection.
MoD vs. Early Exit: Fixed Depth vs. Variable Depth
Early exit lets a token leave the network at an intermediate layer once a confidence criterion is met. Every token still enters every layer up to its exit point, and the hard problems are exit-criterion calibration and per-token latency variance. MoD instead fixes the network depth and varies which blocks run for each token, so latency is more uniform across a batch—an important property for serving. The two are complementary: early exit changes when a token stops, MoD changes where compute is spent within a fixed depth.
Evaluating MoD for Your Stack
MoD is worth a serious look if you are compute-bound at inference and willing to invest in router tuning and causal-selection engineering. It is a poor fit if your bottleneck is memory bandwidth or KV-cache size—MoD does not shrink either—or if you need a drop-in change with no training-side work.
Concrete evaluation checklist:
- Match FLOPs, not parameters. Compare against a dense baseline at the same inference FLOP budget, and count router overhead. Report quality at matched compute, not matched parameter count.
- Tune capacity and auxiliary-loss coefficient jointly. Sweep both; a single-capacity result is not representative.
- Test on your token distribution. MoD's gains depend on how much of your traffic is low-information tokens. Code, structured data, and long-context retrieval workloads behave differently from prose.
- Validate the causal path early. If you need autoregressive serving, prototype the causal-selection variant before committing—this is the most likely source of schedule slip.
- Benchmark against MoE, not just dense. At large scale, MoE is the stronger incumbent on quality-per-FLOP; MoD's case is strongest where memory is constrained and depth is the lever you want.
The honest summary: MoD is a well-motivated, published conditional-compute technique with real but configuration-dependent gains and a known open problem in autoregressive decoding. It belongs in the evaluation set for compute-bound inference stacks—not in the "solved efficiency win" column.