Mixture of Depths: Budget Compute Per Token

Mixture of depths gives a transformer layer a fixed token budget: some positions receive the expensive block while the rest take a residual bypass. The engineering question is not merely which scores are largest, but whether routing, capacity, causality, quality, and latency remain legible together.

This guide builds one deterministic router from tensor shapes outward. Its numbers are worked fixtures, not reported production gains; the point is to expose the exact ledger you should replace with measurements from a target model and serving stack.

Mixture of depths is a capacity contract

At layer (l), a router emits one score per eligible token. Capacity (C_l) says how many positions enter attention and the feed-forward block; the remaining positions carry their residual state forward. The batch and sequence shapes stay stable even though expensive work touches a subset. This is dynamic transformer depth under an explicit quota, not an unbounded promise that every token independently chooses its ideal path.

The Mixture-of-Depths paper introduces the central method and its static compute framing. Before adopting a checkpoint or reproducing training, write the local contract: whether special tokens are always selected, whether capacity is per sequence or flattened batch, how padding is excluded, what router scores can observe, and how bypassed positions rejoin the stream. Those choices change both semantics and kernels.

For a worked batch of two sequences with eight eligible positions each, a 25% ratio gives four routed positions per sequence. Record requested ratio, integer rounding, protected-token reservations, actual selected count, and padding. Mixture of depths becomes testable when the same input, weights, mask, and tie policy produce the same indices and the ledger reconciles every token.

Define router inputs without future leakage

A causal model cannot let an earlier token's route depend on future token content unavailable at generation time. Training implementations that score a complete sequence must explain how their router remains compatible with autoregressive inference. If capacity is chosen across a whole sequence, decode may need a different policy, an online budget, or a checkpoint designed for that mismatch. Token compute routing is part of model semantics, not a serving-only optimization.

Mask padding before ranking. Reserve capacity for any mandatory boundary tokens, then select among the remaining eligible scores. Pin tie-breaking to stable position order or a documented deterministic key. The PyTorch topk documentation warns that indices for tied elements are not guaranteed stable; a test that uses equal scores should therefore detect whether the chosen backend satisfies your reproducibility requirement.

Inspect router features. If score magnitude correlates with token frequency, position, punctuation, language, or prompt template, the model may spend depth according to shortcuts. Log aggregate route rates by declared slices without retaining sensitive text. Mixture of depths needs a protected trace mode that can map selected indices back to synthetic fixtures and an operational mode that records bounded histograms instead of prompts.

Gather, transform, and scatter without semantic drift

After selection, gather chosen hidden states in a deterministic order, run the block, multiply or gate as defined by the checkpoint, and scatter results back to their original positions. Bypassed tokens use the declared residual path. Verify shapes, dtypes, gradient flow, mask interpretation, and position mapping. A fast kernel that permutes equal-score positions may preserve total work while changing the model.

The reference artifact below isolates top-k layer routing from learned weights. It asserts capacity, stable ties, zero and full capacity, invalid budgets, repeatability, and the boolean selection mask. Extend it with a tiny real block: compare the sparse output against a dense control where unselected updates are explicitly zeroed. Forward values, selected gradients, and bypass gradients should match within named tolerance.

Mixture of depths adds communication choices under tensor or sequence parallelism. Scores may be local while a capacity decision is global; gathering token shards can cost more than skipped computation. Mixture-of-experts inference offers an adjacent routing ledger, but expert routing moves tokens among parameter owners while this method skips layer work by position. Keep those traffic models distinct even if kernels share primitives.

Runnable artifact: The deterministic router fixture makes tie order, exact capacity, empty input, invalid budgets, and repeated selection observable before a learned gate or fused kernel enters the test.

Save this worked fixture as mixture-depth-router.test.mjs and run node mixture-depth-router.test.mjs. Expected final line: PASS: 10 router assertions.

import assert from "node:assert/strict";
function route(scores, capacity) {
  if (!Number.isInteger(capacity) || capacity < 0 || capacity > scores.length) throw new Error("capacity");
  const ranked = scores.map((score,index)=>({score,index})).sort((a,b)=>b.score-a.score || a.index-b.index);
  const chosen = new Set(ranked.slice(0,capacity).map(x=>x.index));
  return scores.map((_,index)=>chosen.has(index));
}
let n=0; const check=fn=>{fn();n++};
check(()=>assert.deepEqual(route([.2,.9,.4],1),[false,true,false]));
check(()=>assert.deepEqual(route([.2,.9,.4],2),[false,true,true]));
check(()=>assert.deepEqual(route([.5,.5,.1],1),[true,false,false]));
check(()=>assert.deepEqual(route([],0),[]));
check(()=>assert.deepEqual(route([1],0),[false]));
check(()=>assert.deepEqual(route([1],1),[true]));
check(()=>assert.equal(route([.1,.2,.3],2).filter(Boolean).length,2));
check(()=>assert.throws(()=>route([1],-1),/capacity/));
check(()=>assert.throws(()=>route([1],2),/capacity/));
check(()=>assert.deepEqual(route([.3,.2,.1],2),route([.3,.2,.1],2)));
assert.equal(n,10); console.log("PASS: 10 router assertions");

Train the router against the budget it will serve

Fixed capacity creates competition among tokens. Training needs a rule for selection, a gradient path for router learning, and any auxiliary regularization. Measure score distributions, saturation, protected-token selection, and route churn between checkpoints. If deployment changes capacity from the training setting, treat that as a model intervention and evaluate it rather than a free latency knob.

The Adaptive Computation Time study provides broader context for models that allocate varying computation. Mixture of depths differs in mechanism and should retain its own capacity semantics, yet the common warning is useful: compute allocation and task behavior must be evaluated together. A lower average operation count cannot justify a route policy that erases rare but important tokens.

Train controls at several ratios, including the dense path. Compare loss convergence, router entropy, selected-position patterns, and downstream metrics with matched tokens and optimizer schedules. Inspect whether conditional token compute collapses onto fixed positions or common syntax.

The target is not maximum router novelty. It is a bounded allocation that learns meaningful differentiation without creating an unmeasured quality tax for specific cohorts.

Measure quality by where compute disappears

Aggregate benchmark scores can conceal a route failure concentrated in long identifiers, low-resource languages, code delimiters, citations, or late-sequence corrections. Report task quality beside route rate for length buckets, position buckets, token classes, language, safety category, and tool structure. Compare matched dense and routed checkpoints; a dense model with sparse execution hacked on afterward is not a valid control unless trained for that path.

Use ablations that replace learned selection with first-k, last-k, uniform-stride, seeded-random, and oracle-inspired heuristics. If the learned router barely beats position rules, its scores may not carry the intended information. Test-time compute stop rules asks a related question at the response level; mixture of depths makes the decision inside the network and needs token-level attribution.

Check route stability under paraphrase and harmless formatting. Stability is not always required, but large unexplained changes may predict brittle latency and quality. Inspect routed special tokens and the first few positions because early attention anchors can carry disproportionate value. Preserve representative traces as fixtures, labeling all numerical results as local measurements rather than universal properties of the architecture.

Translate skipped operations into real latency

Theoretical block FLOPs do not equal wall-clock savings. Small or irregular gathers may underutilize accelerators; scatter and routing add memory traffic; decode has too few positions to fill a large sparse kernel; and distributed selection may communicate scores or states. Benchmark prefill and decode separately across batch, length, capacity, dtype, compilation, and parallel topology.

Record dense baseline, router time, selection time, gather/scatter, attention, MLP, synchronization, memory peak, tokens per second, and p50/p95 latency. Mixture of depths succeeds operationally only when the selected capacity maps to a useful kernel shape. Admission control should use measured service curves rather than assuming a 50% token budget halves latency.

Training adds activation memory and scheduling effects. Activation checkpointing trade-offs can change which intermediates are saved, while zero-bubble pipeline parallelism may expose gaps that sparse work can or cannot fill. Reprofile the composed system. A routing win measured on one isolated layer should not be advertised as end-to-end throughput before communication, batching, and fallback are included.

Ship a dense escape path and a complete receipt

Define failure behavior for invalid scores, NaNs, capacity mismatch, unsupported shapes, missing fused kernels, and distributed disagreement. A dense path is the safest fallback when the checkpoint supports it; otherwise refuse the configuration rather than silently inventing a new route. Monitor selected counts and route-score distributions, but keep high-cardinality token content out of general telemetry.

Publish checkpoint identity, training capacity, inference capacity, masks, mandatory tokens, rounding, tie policy, causal policy, router dtype, kernel version, parallel layout, quality slices, and latency curves. Include the deterministic artifact and dense masked equivalence test. Mixture of depths then becomes a reproducible compute budget rather than a diagram of tokens taking shortcuts.

The final decision should name the workload. A capacity that benefits long batched prefill may do nothing for single-token decode. A quality threshold acceptable for summarization may be unacceptable for constrained tool calls.

Keep the route ledger beside the deployment profile and refresh it whenever batcher, compiler, hardware, checkpoint, or sequence mix changes. The model is spending depth position by position; the release evidence should be equally specific about where those savings came from and what they cost.