Mixture-of-Depths (MoD) & Adaptive Per-Token Compute

One-line summary. Not every token needs the whole model. Mixture-of-Depths puts a tiny router in front of each transformer block that, under a fixed budget, lets only the top-k most "deserving" tokens pay for that block's attention+MLP — everyone else rides the residual connection straight past. Same quality, fewer FLOPs, faster steps. Where [[mixture-of-experts]] routes tokens sideways to different experts at full depth, MoD routes them vertically through different depths.

What it is (intuition first)

Read a sentence out loud and notice your own effort. The word "the" needs almost no thought; the word that pins down the meaning of the whole clause needs a lot. A standard transformer doesn't have that luxury — it spends exactly the same compute on every token at every layer. The word "the" gets the full attention pass, the full MLP, the full stack of layers, identical to the hardest word in the sentence. That's egalitarian and wasteful.

Mixture-of-Depths is the idea that a token should be able to skip a layer when that layer has nothing to add to it. Concretely: in front of each transformer block you place a tiny router — one small linear layer that emits a single scalar per token. Under a fixed budget (say "only 12.5% of tokens may enter this block"), the block processes the top-scoring tokens and lets the rest bypass the block entirely through the residual connection, arriving at the next layer unchanged Raposo et al. 2024, [arXiv:2404.02258].

The payoff has the same flavour as MoE — decouple a per-token cost from a global quantity — but along a different axis:

MoE decouples total parameters from active FLOPs. The model is huge; each token only touches a slice of it. Every token still goes through every layer at full depth.
MoD decouples which tokens from how much depth. The model is the normal size; the budget of token-slots per block is what's capped. Different tokens take different-length paths through the same stack.

A useful one-liner that the literature itself uses: MoE routes tokens to different experts; MoD routes tokens to different depths. They are orthogonal — you can do both at once (see MoDE below).

The crucial constraint that makes MoD different from naive early-exit. MoD uses a static compute budget. Because k is fixed a priori, the tensors have known sizes and the computation graph is static — k tokens go through the block, full stop. This is what makes it GPU-friendly and is the key engineering difference from "let any token exit whenever it's confident," which produces ragged, data-dependent shapes [arXiv:2404.02258].

Why it matters

The dominant cost in a transformer is that compute scales with sequence length × depth, applied uniformly. Self-attention is worse than linear — its query-key product is quadratic in the number of participating tokens. So if you could shrink the number of tokens that actually participate in a block, the savings compound super-linearly in exactly the most expensive operation.

MoD's headline results, at a fixed FLOP budget (isoFLOP), are:

Models can match the baseline's final loss while using a fraction of the FLOPs per forward pass, and some MoD variants reach a lower loss than the isoFLOP-optimal baseline Raposo et al. 2024, [arXiv:2404.02258].
Because each forward pass is cheaper, MoD models can be "upwards of 50% faster to step during post-training sampling" [arXiv:2404.02258].
The optimal configuration in the paper routes only 12.5% of tokens (capacity = 1/8) through each MoD block, with the other 87.5% bypassing via the residual — applied to every other block [arXiv:2404.02258].

The deeper reason it matters is philosophical alignment with how inference should work: spending equal compute on trivial and hard tokens is a known inefficiency, and MoD is one of the cleaner ways to attack it while keeping the static, batched, tensor-core-friendly shapes that GPUs love. It sits in the same family as [[test-time-compute-reasoning]] (spend more where it helps) read in reverse: spend less where it doesn't.

The lineage: early-exit, depth-adaptive, then MoD

MoD didn't appear from nowhere. The "not every token needs full depth" idea has an earlier branch — early-exit — and understanding the contrast is the whole point.

Depth-Adaptive Transformer & CALM (early-exit)

Early-exit flips MoD's question around. Instead of "which tokens may enter this block," it asks "may this token *leave* the stack now?" A token climbs the layers and, at each candidate exit point, a lightweight classifier checks a confidence measure; once the model is confident enough about the next-token prediction, the token exits early and skips all remaining layers.

The Depth-Adaptive Transformer (Elbayad et al. 2020) introduced per-token dynamic depth for sequence models, learning when to stop computing arXiv:[2004.07453].
CALM — Confident Adaptive Language Modeling (Schuster et al. 2022, a NeurIPS oral) made this rigorous for decoder LLMs. It frames three problems: what confidence measure to use, how to tie a sequence-level quality guarantee to per-token exit decisions, and how to handle "missing hidden representations due to early exits in previous tokens" Schuster et al. 2022, [arXiv:2207.07061]. It reports a potential speedup of up to ×3 while provably staying within a chosen quality tolerance [arXiv:2207.07061].

That third problem is the catch with early-exit: when token t exits at layer 5, the deeper layers never computed its hidden states — but token t+1 attends back to token t at layer 9, where that state is now missing. CALM has to reconstruct or propagate the skipped state (e.g. copy the exit-layer state forward). This bookkeeping, plus the data-dependent, ragged shapes (different tokens exit at different layers, so batches aren't rectangular), is what makes early-exit awkward to batch efficiently.

How MoD fixes the ragged-shape problem

MoD's design choice is to make the budget static instead of confidence-driven. Rather than "exit when confident" (variable, data-dependent), MoD says "exactly `k` tokens may use this block" (fixed, known at compile time). The shapes stay rectangular, the kernels stay tensor-core-friendly, and you get most of the adaptivity (the router still chooses which tokens dynamically) without the ragged-tensor tax. The price is that MoD's budget is a fixed cap, not a true per-input dial — a sequence of all-hard tokens still only gets k slots per block.

How it works (the real mechanics)

The per-block router

Each MoD block has a learned router weight vector w_θ. For each token embedding x_i at layer l, the router emits a scalar:

r_i^l = w_θ^T · x_i^l

These scalars are ranked across the sequence. With a capacity C (e.g. C = S/8 for 12.5%), the block admits the tokens whose router weight exceeds the β-th percentile, β = 1 − C/S. Admitted tokens go through the block's attention + MLP; the rest skip it. The output rule is:

x_i^(l+1) = r_i^l · f(X̃^l)_i + x_i^l      if  r_i^l > P_β(R^l)     # routed: processed, then residual
x_i^(l+1) =                    x_i^l       otherwise               # bypassed: pure residual

Two subtle but load-bearing details:

The router weight `r_i` multiplies the block output before the residual add. This "puts the router weights along the gradient path" — so the router is trained by ordinary backprop through the tokens it did select, learning to predict which tokens benefit from the block [arXiv:2404.02258].
Bypassed tokens are not dropped — they pass through unchanged via the residual. Their information is preserved; they just don't get updated by this particular block.

A 12.5%-capacity self-attention block makes the quadratic query-key product roughly (0.125)² ≈ 1.56% of the vanilla cost for that block — the super-linear win mentioned above [arXiv:2404.02258].

The causal-routing trick (the genuinely hard part)

Top-k routing is non-causal, and for an autoregressive model that's a real problem. To know whether token t is in the top-k of a block, you need the router weights of all tokens in that block — including tokens that come after t. During training (full sequence visible) that's fine. During autoregressive sampling, token t is generated before t+1 exists, so you cannot compute its top-k membership. The routing decision must be made causally, from token t's own state alone.

Raposo et al. give two ways to keep MoD causal at inference [arXiv:2404.02258]:

(a) Auxiliary predictor loss. Add a binary cross-entropy loss where the router's per-token output is the logit and the (non-causal, training-time) top-k membership is the target. The router thus learns to predict "will I be selected?" from the token alone. Reported ~99% accuracy, at a small ≈0.2–0.3% performance degradation.
(b) A small auxiliary MLP "predictor." A separate tiny MLP receives the same input as the router but with a stop-gradient (so it doesn't perturb the main model), and outputs a causal prediction of top-k membership. This "empirically does not significantly impact the step speed" and sidesteps the loss-based variant's quality hit.

Either way, the trick is the same shape: learn a causal proxy for the non-causal top-k decision. This is the MoD analogue of the routing subtleties MoE wrestles with — both architectures pay a "the selection isn't naturally differentiable / isn't naturally causal" tax and engineer around it.

Routing frequency and budget

The paper's tuned recipe interleaves: route every other block at 12.5% capacity, leaving the alternate blocks at full capacity. This balances "skip enough to save real FLOPs" against "don't starve every token of updates." Capacity, block frequency, and whether to route around attention only vs the whole block are the main knobs.

MoDE — combining MoD with MoE

Because MoD picks depths and MoE picks experts, they compose. The paper calls the combination Mixture-of-Depths-and-Experts (MoDE) and gives two flavours [arXiv:2404.02258]:

Staged MoDE — route tokens around the block (skipping self-attention) before the MoE step. Lets a token skip attention entirely, then still hit the experts.
Integrated MoDE — fold MoD directly into the expert choice by adding "no-op" experts alongside the normal MLP experts. Choosing a no-op expert is the bypass. The paper notes this is "distinctly better than simply reducing the capacity of experts in conventional MoE models," because tokens explicitly learn to choose the residual path rather than preferring an expert and being silently dropped.

The reported result is that MoDE's improvements compound with MoE's — the two efficiency axes stack rather than cannibalise each other.

Key ideas & tradeoffs

Lever	What it buys	What it costs
Static top-`k` budget	GPU-friendly fixed shapes; predictable FLOPs	Budget is a cap, not a true per-input dial
Lower capacity (e.g. 12.5%)	Big FLOP cut, faster steps	Fewer token-slots; risk of starving hard tokens
Route every other block	Keeps a full-update path alive	Less savings than routing every block
Router-on-gradient-path	Router learns which tokens benefit	Adds the non-causal selection problem
Causal predictor (loss or MLP)	Autoregressive sampling works	~0.2–0.3% quality hit / extra module
MoDE (MoD + MoE)	Two efficiency axes compound	Two routers, two failure modes to stabilise

Mental model of the core tension: MoD buys speed by *unevenly spending depth*. It wins exactly when token difficulty is genuinely non-uniform (most natural-language sequences) and a router can learn to spot the easy ones. It wins less when every token is hard, or when the router can't tell them apart — then you've paid for routing machinery to skip blocks that mattered.

Honest caveats & open questions

Modest real-world adoption. This is the big honest caveat. Despite clean isoFLOP results, MoD has not become a mainstream production architecture the way MoE has. Frontier 2024–2025 open models lean on MoE (Mixtral, DeepSeek-V3, Qwen-MoE); MoD remains mostly a research line with follow-ups (router-tuning, MoD-for-vision/CNNs) rather than a standard component [Turing Post; secondary]. Treat "everyone will use MoD" as unproven, not established.
Router brittleness at inference. Causal routing is only as good as the predictor. If the router's causal proxy mispredicts, a token that needed a block gets skipped, and in autoregressive generation those errors can propagate token-by-token. Stable generation is exactly where mistakes hurt most.
The capacity cap is rigid. A fixed k means MoD cannot give a genuinely hard region of text more total depth than a budget allows — it can only reallocate within the cap. It's adaptive in allocation, not in total compute. (CALM-style early-exit is adaptive in total compute but pays the ragged-shape tax — you pick your poison.)
Retrofitting pretrained models is delicate. Inserting MoD blocks into an already-trained dense model can disrupt the representation balance; reported mitigations (gated normalization, router-only tuning) are patches signalling the difficulty, not a solved problem [secondary].
Where the savings actually land. The dramatic 1.56% figure is for one block's attention under 12.5% capacity — the end-to-end model still has full-capacity blocks, MLPs, embeddings, and the un-skipped path. The honest top-line is "match quality at lower FLOPs / up to ~50% faster sampling steps," not "10× cheaper model." Don't over-extrapolate the per-block number.
Interpretability of skips is shallow. As with MoE experts, which tokens a router learns to skip doesn't cleanly map onto "easy words." The selection often tracks surface features; "the router skips the unimportant tokens" is an aspiration, not a guarantee.

How it connects to OpenAlice

Be precise — don't overclaim a tie that isn't there.

No in-house MoD. As of this writing there is no honest tie to MoD training or inference in the OpenAlice codebase; Alice's default stack runs gpt-5.x / Codex-family providers. The relevance is conceptual literacy about the efficiency frontier, not a shipped feature.
The conceptual mirror is adaptive *agent* compute. MoD's thesis — spend compute where it's needed, skip where it isn't — is exactly the principle behind Alice's orchestration: cheap models / shallow passes for easy turns, Opus / deeper reasoning for hard ones. MoD does this at the layer level inside one forward pass; Alice's cost-ladder and [[model-routing]] do it at the whole-model level across turns. Same instinct, different abstraction, and worth keeping that distinction crisp in any internal doc.
The budget-vs-confidence tradeoff transfers. MoD's "static budget, GPU-friendly shapes" vs CALM's "confidence-driven, ragged shapes" is the same design choice Alice's dispatcher faces: cap effort per request for predictable cost, or let effort float with difficulty for best quality. Knowing both failure modes (starved-under-cap vs ragged-and-unbatchable) informs how the orchestrator sets per-task budgets under a shared quota.
Don't conflate MoD with MoE in lab notes. MoE is the one Alice actually consumes (DeepSeek-class providers); MoD is the depth-axis cousin she does not. Cite [[mixture-of-experts]] for the models we run, this article for the contrast.

Mixture-of-Depths (MoD) & Adaptive Per-Token Compute

Mixture-of-Depths (MoD) & Adaptive Per-Token Compute

What it is (intuition first)

Why it matters

The lineage: early-exit, depth-adaptive, then MoD

Depth-Adaptive Transformer & CALM (early-exit)

How MoD fixes the ragged-shape problem

How it works (the real mechanics)

The per-block router

The causal-routing trick (the genuinely hard part)

Routing frequency and budget

MoDE — combining MoD with MoE

Key ideas & tradeoffs

Honest caveats & open questions

How it connects to OpenAlice

See also

Sources (fetched & verified 2026-06-17)