Anticipate what breaks when you bolt naive early exit onto a decoder at inference

MCQ·Medium·4.0 · 0·~1 min·Asked atPalantir Perplexity Two Sigma

Attempt it

**Early exit** lets a decoder stop computing at some intermediate layer when the prediction at that layer is already confident enough. It promises a speedup proportional to the layers skipped. **The catch in autoregressive decoding.** Every token's K and V vectors are stored in the KV cache for every layer, so future tokens can attend to them. If token at position 5 exits at layer 12 of a 32-layer model, layers 13-32 never compute its K and V. When token at position 6 runs all 32 layers, it tries to attend to position 5 at layers 13-32 and finds nothing. **What this means for output.** Each future token's attention softmax at the affected layers runs over a broken distribution, missing entries it expects. The model's deeper-layer computations operate on garbage, and quality collapses. **How real systems fix it.** Copy the exited token's hidden state forward into the higher layers' KV slots (cheap but approximate). Recompute the missing layers lazily for tokens that later turn out to need them. Train the model with layer dropout (LayerSkip) so intermediate hidden states remain head-compatible and KV-substitutable. The other options miss the real failure mode: the head can read intermediate states when trained for it; sampled token ids stay valid; weights do not corrupt at inference.

**The KV cache is the bottleneck.** In autoregressive decoding, every produced token's K and V vectors at every layer are stored so subsequent tokens can attend back. Early exit cuts the forward pass short, which means the exited token never produces K and V at the skipped layers. Subsequent tokens that run the full stack expect to find those entries and do not. **Concrete failure.** Token 5 exits at layer 12 of 32. Layers 13-32 of its forward pass are not executed, so `K_5` and `V_5` at those layers do not exist. Token 6 enters layer 13 and computes attention over positions 1 through 5. The K and V entries for tokens 1-4 (which ran the full stack) are present; the entry for token 5 is missing. The softmax either renormalizes over a smaller key set or hits an undefined slot, depending on the implementation. Either way, the per-head attention distribution at layers 13-32 is wrong for every token following an early exit. **Practical mitigations.** SkipDecode and CALM copy the hidden state at the exit layer up the rest of the stack and write synthetic K and V from it; the substitution is approximate but cheap. LayerSkip (Meta) trains with layer dropout so intermediate representations are explicitly head-compatible, then runs self-speculative decoding to verify exited tokens against the full stack. Some schemes recompute the skipped layers lazily on cache misses.

TL;DR

Tokens that exit early skip producing K and V at the higher layers, so later tokens attending to them at those layers see holes in the KV cache and softmax over a broken distribution.

Memory aid

Easy to grasp

Picture a relay race where every runner has to hand a baton to every later runner who passes their position. If one runner leaves the track halfway through, future runners who reach the back half find no baton waiting at that spot. The race does not gracefully degrade; it stops or produces nonsense. Early exit in a decoder is the same: when a token is allowed to stop early, the upper layers never store anything for that token, and later tokens that try to attend to it at those upper layers find empty slots. The fix is to either fill in the missing slots from what the exited token did produce, or to train the model so partial layers are okay.

Key concepts

Transformer Architecture

Concept explanation~2 min read

Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example. Click to expand.

Early exit is a tempting optimization. If a decoder can produce a confident prediction at layer 12 of 32, why run the remaining 20 layers? The promise is a meaningful per-token speedup proportional to the depth saved. The reality, especially in autoregressive generation, is that early exit collides with the KV cache in a way that breaks subsequent tokens' attention computations.

This question asks what specifically breaks when you bolt naive early exit onto a decoder at inference. The answer is the KV cache: when a token exits at layer L, it never computes K and V at layers L+1 through N, and any later token that runs the full stack expects to find those entries when it attends back. The cache develops holes, attention softmaxes operate over the wrong key set, and quality collapses within a few tokens.

The deep dive walks through why the cache stores per-layer state, the concrete failure mode, three classes of mitigation that real systems use, and why production decoders mostly do not ship token-level early exit despite the speedup potential.

Why the KV cache stores per-layer K and V for every token

Autoregressive attention is causal. When token t+1 runs the forward pass, its attention at every layer reads keys and values for tokens 1 through t from the same layer. The cache stores, for each layer L and each produced token, the per-head K_t^L and V_t^L. Without the cache, generating each token would re-run the full prefix through every layer, scaling quadratically with sequence length.

The cache is part of the model's runtime state, not a serving optimization. The model's correctness depends on the cache being consistent with what a full forward pass would have produced. Any optimization that alters the forward pass per token has to maintain that consistency or accept the quality consequences.

Memory math. For a model with L layers, H_kv KV heads, head dim d_h, sequence length T, batch size b, and 2 bytes per element, the cache costs 2 * L * H_kv * d_h * T * b bytes. For Llama 3 70B at 8k tokens and batch 1, that is roughly 2.6 GB. The size is part of why long context is expensive at inference.

The implication for early exit. Every entry the cache expects to find has to be produced by some computation. If the forward pass skips layers, the entries are not there. The model's per-layer expectation is structural, not soft.

The concrete failure mode of naive early exit

Three classes of mitigation that real systems use

Why production decoders mostly do not ship token-level early exit

Situations where this technique stops working.

2–4 min · Everything important, quickly.

python

# Early-exit breaks the KV cache invariant: later tokens may need
# cached K/V from layers that never ran for earlier tokens.
for t in range(T):
    x = embed(tokens[t])
    for L, layer in enumerate(layers):
        # cache write is conditional on reaching layer L
        k, v = layer.kv(x)
        kv_cache[L].append(k, v)  # MISSING for early-exited tokens
        x = layer.attn(x, kv_cache[L]) + x
        x = layer.ffn(x) + x
        if exit_classifier(x).confident():
            break  # later tokens lose access to layers > L for token t

Real products, models, and research that use this idea.

Meta's LayerSkip (2024) trains models with layer dropout plus self-speculative decoding, mitigating the KV hole problem by design rather than at inference time
CALM (confident adaptive language modeling) by Google was an early proposal for token-level early exit; later analysis highlighted the KV cache complication

What an interviewer would ask next. Try answering before peeking at the approach.

QIf you propagate the exit-layer hidden state up to fill KV entries, what subtle quality cost do you pay?

The synthesized K and V are derived from a residual state that the upper layers would have transformed further. Future tokens attending to those entries see content that is shifted from what a full forward pass would have produced. The drift is small per token but compounds in long generations.

1 more follow-up an interviewer would ask next. Sign in to reveal them.

Red flags & common mistakes

The phrases that signal junior thinking. Click to expand.

Most common mistake

Assuming the unembedding head cannot read intermediate states, when the real obstacle is the KV cache that future tokens depend on at the skipped layers.

60 second bullets to scan on the way to the call.

Why the KV cache stores per-layer K and V for every produced token
What concretely breaks when an exited token leaves holes at deeper layers

Primary sources. Browse if you want the original framing.