Inference Optimisation advanced 8 min read 10 flashcards

Early Exit and Adaptive Depth

Not every token needs all eighty layers, and the methods that act on this observation break batching and the KV cache in ways that decide whether the idea survives contact with a serving stack.

Predicting the token after "the capital of France is" and predicting the next step of a proof cost exactly the same amount of compute in a standard transformer: every token traverses every layer. That uniformity is an artefact of the architecture, not a property of the problem. Adaptive-depth methods try to spend compute where difficulty actually is, and the reason they are still rare in production is not that the idea is wrong but that its interactions with batching and the KV cache are brutal.

Exit when confident

The basic construction attaches a classifier head to intermediate layers, computes a confidence signal at each candidate exit, and stops as soon as confidence clears a threshold. CALM formalised the three problems this raises: what confidence measure to use (softmax margin, a trained exit head, or hidden-state saturation), how to translate a global constraint like "stay within \(\epsilon\) of the full model's output with probability \(1-\delta\)" into per-token thresholds, and what to do about the KV entries for layers that were never computed for earlier tokens (Schuster et al., 2022, Confident Adaptive Language Modeling, arXiv:2207.07061).

That third problem is the one that surprises people. If token 17 exited at layer 12, then layers 13 through 80 have no keys or values for position 17. When token 18 runs deep, its attention at layer 40 needs them. The options are to copy the layer-12 hidden state up the stack as a substitute, to recompute the missing layers, or to restrict exits so this cannot arise. Each costs some of the saving.

Train the model to be exitable

Retrofitting exits onto a model trained only to be read at the last layer works poorly, because intermediate representations were never required to be decodable. LayerSkip trains for it directly: layer dropout with rates increasing by depth, plus an early-exit loss applied at every layer through a shared output head. The same property then enables self-speculative decoding, where the model's own early layers draft tokens that its full depth verifies, reported at 1.34× to 2.16× end-to-end speedups depending on task (Elhoushi et al., 2024, LayerSkip, arXiv:2404.16710).

Self-speculation is the practical reframing of early exit. Instead of trusting a shallow prediction, you use it as a draft and let the full model verify, which gives back exactness: the output distribution is the full model's, and the saving comes from parallel verification rather than from accepting a worse prediction.

Adaptive depth as routing

Mixture-of-Depths approaches the same objective from the opposite direction. Rather than deciding per token when to stop, each layer runs a router that selects a fixed top-\(k\) fraction of the sequence's tokens to process, and the rest bypass that block through the residual connection. Because \(k\) is fixed in advance, the compute graph is static and the tensor shapes are known at compile time, which is what makes it batch and compile like a normal model (Raposo et al., 2024, arXiv:2404.02258). The cost is that selection is competitive across the sequence, so it is not a pure per-token difficulty measure, and causal routing during autoregressive decoding needs care because a token's selection cannot depend on tokens that do not exist yet.

When it breaks

  • Batching destroys the saving. In a batch of 64 sequences, the batch cannot leave layer 12 until every member has exited. Per-sequence savings become max-over-batch savings, and with continuous batching the fast requests simply wait. This, more than accuracy, is why adaptive depth is rare in high-throughput serving and more attractive in single-stream, on-device settings.
  • Confidence is not correctness. Exit thresholds calibrated on one distribution drift on another, and an overconfident shallow exit produces a wrong token that the rest of the generation then conditions on.
  • The savings profile is uneven. Easy tokens are common in boilerplate and rare in the reasoning traces where compute actually hurts, so measured end-to-end wins are usually well below the average-layers-saved figure.
  • It competes with simpler options. Quantisation, speculative decoding with a small dedicated draft model, and routing easy queries to a smaller model all deliver similar gains with far less invasive changes to the serving stack. Adaptive depth wins when a single deployment must serve wildly heterogeneous difficulty and cannot host a second model.
Check yourself

10 flashcards for this concept

Click a card to reveal the answer.

Drill the whole track