Why architecture matters here
The case for early exit rests on a measurable fact about how transformers compute: token difficulty is enormously uneven. In natural text, a large fraction of tokens are near-deterministic given context — punctuation, function words, the completion of a partially typed word, the predictable structure of code and markup. A model's internal representation of such a token often stabilizes to the correct answer within the first half of the layers, and the remaining layers change nothing but cost the same. Uniform-depth inference pays the maximum price for the minimum-difficulty token, which in aggregate is where most of the wasted compute lives. Early exit reallocates compute toward the tokens that actually need it.
Why this matters architecturally, rather than as a nice-to-have, is that it changes the shape of the latency-quality trade-off. The conventional levers — quantization, distillation, pruning — reduce cost uniformly: every token gets cheaper by the same factor, and quality degrades globally. Early exit is conditional: it spends full capacity exactly where difficulty demands it and saves everywhere else, so at a given average latency it can preserve quality on hard tokens that uniform methods would have degraded. It composes with the uniform methods too — an early-exit model can also be quantized — because it operates on an orthogonal axis, depth-per-token rather than cost-per-layer.
But the same conditionality is what makes the architecture demanding. The moment different tokens run different numbers of layers, three things that were trivial become hard. First, the exit decision must be cheap and reliable: a confidence estimate that costs as much as the layers it saves, or that is wrong, destroys the benefit. Second, the KV cache — the stored keys and values every future token attends to — assumes every layer was computed for every past token; if a token exited early, the deeper layers' KV entries for it were never produced, and later tokens attending to those layers find holes. Third, batching, the workhorse of efficient serving, assumes all tokens march through layers in lockstep; ragged exits break that lockstep. A design that ignores any of these three delivers FLOP savings on paper and no speedup, or worse, wrong answers, in practice. The architecture is precisely the set of choices that makes conditional depth actually pay off.
It helps to situate early exit among its relatives, because the adaptive-computation design space is broader than a single mechanism. Layer-skipping and adaptive-depth methods drop or bypass layers conditionally; mixture-of-depths routes each token to a learned subset of layers; and speculative decoding uses a small draft model to propose tokens a large model verifies. Early exit is the member of this family that keeps a single model and a single forward direction, simply stopping when confident — which makes it the easiest to reason about and to retrofit onto an existing SLM, at the cost of the KV and batching complexities described here. Understanding it well is a good foundation for the rest of the family, which shares the same core tension between saving compute and keeping the cache and the batch coherent.
The architecture: every piece explained
Top row: the exit decision path. An input token is embedded and flows through the early layers that every token shares. After a designated layer, an exit classifier produces a confidence signal about the next-token prediction available at this depth. The signal can be the softmax entropy or max-probability of the intermediate LM head's output, the agreement between consecutive layers' predictions, or a small learned confidence head. The exit threshold is then applied: if confidence clears it, the token exits — the prediction from this layer is emitted — otherwise it continues. Thresholds are not global constants; they are calibrated per exit point and per task, because the reliability of a layer-6 prediction differs from a layer-18 one and differs across domains.
Middle row: what makes exits usable and consistent. The deeper layers run only for tokens that didn't exit. A shared LM head (or a small head per exit) maps the hidden state at each exit point to vocabulary logits; sharing the head keeps parameters and training cost down, and is why intermediate representations must be trained to be head-compatible. The critical piece is KV cache handling: when a token exits at layer k, the keys and values for layers k+1..N were never computed, yet a future token running to full depth will attend to those layers for this position. The two standard resolutions are to copy the layer-k KV state forward to fill the skipped layers (state propagation), or to run the skipped layers lazily only if a later token actually needs them; either way the cache must present a complete tensor to attention. Calibrated thresholds tie exit rate to a target quality budget.
Bottom rows: training and serving realities. Training the exits requires auxiliary losses at each intermediate head so that early predictions are genuinely accurate rather than half-formed — the model is trained to be usable at multiple depths, which is a different objective than standard training and slightly regularizes the whole network. Batch handling is the serving crux: within a batch, tokens want to exit at different layers, producing ragged computation; implementations either process the batch to full depth but mask exited tokens (simple, limited savings), or dynamically shrink the active batch as tokens exit (more savings, more complex memory management), or group tokens by predicted difficulty. The ops strip is the control surface: the quality-versus-latency curve you actually operate on, live monitoring of exit rates by layer (a sudden drop means the model stopped trusting early exits — a data-shift signal), and periodic threshold re-tuning.
The choice of confidence signal is a design axis worth making explicit, because it directly sets the cost-benefit of every exit check. The cheapest signals read the intermediate LM head directly — the maximum softmax probability or the entropy of the next-token distribution at this layer — and exit when the distribution is peaked enough. Slightly more robust is saturation: exit when consecutive layers agree on the top prediction, which resists the overconfidence a single layer can exhibit. Most robust, and most expensive, is a small learned confidence head trained to predict whether the early answer will match the full-depth answer. The trade-off is direct: richer signals exit more accurately but cost more per check, and since every token pays for the check whether or not it exits, an over-engineered signal can erase the very savings it was meant to protect.
End-to-end flow
Walk a short generation on an on-device SLM assistant completing a sentence. The prompt is processed (prefill), and the model begins generating tokens one at a time. The first content token is genuinely ambiguous — the model must choose the subject of the sentence — so at the layer-6 exit its confidence is low (high entropy across several plausible words), it fails the threshold, continues through the deeper layers, and only at layer 20 does its prediction sharpen enough to exit near the top of the stack. That token used almost the full network, correctly, because it was hard.
The next several tokens are easy: the model is now completing a common phrase. At the layer-6 exit, the intermediate head's prediction is already sharply peaked on the obvious next word; confidence clears the calibrated threshold; the token is emitted immediately and layers 7..N are skipped. Here the KV-cache machinery does its quiet work: because a future token will attend to all layers at this position, the serving engine propagates the layer-6 key/value state into the skipped layers' cache slots (or marks them for lazy computation), so attention downstream sees a complete cache and produces correct results. From the user's perspective these easy tokens appeared almost instantly, because they cost a third of the compute of the hard one.
Now scale to a server batching eight requests. At each layer, some tokens across the batch want to exit and others don't. A naive implementation runs the whole batch to full depth and simply ignores the layers past each token's exit — correct, but it captures little of the theoretical saving because the batch still pays for the deepest token at every layer. A better serving engine compacts the active set: as tokens exit, they are removed from the working batch, so deeper layers process a shrinking tensor, and the batch's total compute tracks the sum of per-token depths rather than the maximum. This is where the FLOP saving becomes wall-clock saving — and where the engineering is hard, because it requires dynamic tensor shapes and careful memory reuse.
Consider the stress case that reveals the design's limits. A distribution shift arrives — the user switches to a technical domain the exit heads were poorly calibrated on. The confidence signal becomes unreliable: early heads are overconfident on tokens they get wrong, so tokens exit early and quality drops, or they are underconfident and nothing exits, so latency reverts to full depth with the overhead of the exit checks on top. The monitoring catches both: the exit-rate-by-layer dashboard shows either an anomalous spike in early exits paired with a quality regression, or a collapse in exit rate. The response is to re-calibrate thresholds on in-domain data, or fall back to a conservative threshold that only exits on the highest-confidence tokens — trading some speed for guaranteed quality. The whole path — per-token exit layer, confidence, and the resulting quality metric — is logged, so the trade-off is observable rather than a black box, which is what lets you operate early exit at a chosen point on the quality-latency curve instead of hoping.
A subtlety the batch case exposes is the interaction between compaction and the KV cache, which is where the two hard problems meet. When exited tokens are removed from the active batch to shrink deeper-layer compute, their positions still exist in the sequence and future tokens will attend to them, so the engine must have already materialized — or lazily arranged to materialize — the skipped layers' KV for those positions before they leave the batch. An implementation that compacts the batch for speed but forgets to backfill the cache produces fast, wrong output; one that backfills eagerly for every exit gives back much of the compute it saved. The workable designs sit between: backfill lazily, and only when a later token actually attends to a skipped slot, so the cost is paid only where it is genuinely needed.