Post-mortem: layer early-exit does not pay on this model family
Status: negative result, double-confirmed (lossless + lossy). Kept on purpose — the tooling is here, default-OFF, so the conclusion can be re-tested on new models rather than re-derived. We learn from failure; we don't hide it.
Origin: assessment of NextLat (arXiv 2511.05963).
The companion design note is ../../distributed-inference-host/docs/speculative-decoding-plan.md.
The win from that assessment was batched-GEMM prefill (21× lossless, landed). This
doc is about the part that didn't work.
The hypothesis
Skip transformer layers per decode token: run layers 0..k, then final_norm + lm_head
on the shallow residual, when the output is provably (or empirically) unchanged. Two
gates were tried:
- Lossless / conformal — exit when the proven
2·τ ≤ γenergy bound certifies the argmax is locked (tukey_bowl.rs/conformal_gate.rs). τ = early-exit→full-depth sup-norm logit residual, calibrated on prefill; γ = top1−top2 logit margin. - Lossy — accept a bounded flip rate:
- Saturation gate: exit when
‖h_l − h_{l-1}‖₂ / ‖h_l‖₂ < ε(cheap, O(hidden_dim), no lm_head probe — the intended unlock for high-vocab models). - Margin gate: exit when
γ ≥ θfor a free θ (decoupled from the2·τbound).
- Saturation gate: exit when
The result: it doesn't pay
Lossless (conformal): zero tokens admitted on Loki-8B. The early-exit→full-depth
residual is ~11–14 logits, so the gate needs γ ≥ ~23–28 to fire; greedy decode margins
mid-sentence are almost never that large. Correct "miss, not lie" behavior — but no speedup.
Lossy saturation: the residual stream never converges before the final layers. Measured per-layer relative delta during decode on Loki-8B:
back half (l16..31): 0.25 – 0.66 final layers l30,l31: ~0.40 – 0.52
min back-half delta (l24..28): ~0.19 – 0.27 ← still far above any quality-preserving εSo any ε small enough to preserve quality never fires; any ε large enough to fire (≥0.25–0.30) flips the first token and the error cascades through the KV cache. The qwen2.5-0.5b sweep confirms it: no ε reaches ≥98% agreement with any speedup (speedup only appears at ε≥0.3 where agreement is ~0%).
Lossy margin: preserves 100% agreement but the lm_head probe costs more than the rare exits save → net slowdown (0.5–0.6×).
Cost model (warm, measured)
| model | layers | vocab | per-layer µs | lm_head µs | lm_head / layer |
|---|---|---|---|---|---|
| Loki-8B | 32 | 128k | 283,457 | 636,521 | 2.25× |
| qwen2.5-0.5b | 24 | 152k | 3,753 | 13,410 | 3.57× |
A margin probe costs 2.25–3.57 layers. So a margin-gated exit must fire and save ≥4 layers on a large fraction of tokens to break even — it doesn't. The saturation probe is genuinely free (the premise was sound); the premise that fails is "the residual saturates."
Why this is a real conclusion, not a tuning miss
Both an exact certificate and a lossy knob were swept across q / ε / θ, on the operator's real model and a 5×-cheaper-lm_head model, with a fixed determinism self-check (which caught and removed a real confound — see below). The signal is structural: this model family keeps moving the residual ~25–50% per layer right up to the output, so there is nothing to safely shortcut at the single-token greedy-argmax level. Depth is being used.
Bug found along the way (worth keeping)
The first bench runs were nondeterministic — the source of an earlier "1.47× speedup"
that was actually a cold-cache artifact. Root cause: the FFN resonance cache
(resonance_threshold = 0.8) survives reset_context_and_kv, so consecutive runs
diverged. The bench now constructs with resonance = 1.0 (off) and asserts a
two-baseline determinism check before trusting any number. Any future A/B benchmark in
this crate should do the same.
How to re-test on a new model
The tooling is in-tree, default-OFF:
src/early_exit.rs— gates + cheaprelative_l2_deltasaturation probe.src/bin/bench-early-exit.rs— warmup, determinism self-check, cost-model dump, ε/θ sweep with knee detection, and--profile-onlyto cheaply dump the per-layer saturation curve on a slow model.- Flags:
GNOSIS_EARLY_EXIT,GNOSIS_EARLY_EXIT_EPSILON,GNOSIS_EARLY_EXIT_MIN_LAYER,GNOSIS_EARLY_EXIT_THETA. When unset,forward_oneis byte-identical (parity tests).
cargo build --release --bin bench-early-exit
./target/release/bench-early-exit --knot <model.knot> --gen 64 --warmup 8 \
--eps 0.05,0.1,0.2,0.3 --theta 8,14,20A model would flip this conclusion only if its back-half per-layer delta drops below
~0.05–0.10 (genuine saturation) while decision margins stay healthy. Worth a quick
--profile-only check before assuming layer-skip helps on any future model.
Next exploration
The cheap-probe insight is right; the signal is wrong. Instead of the raw residual
delta, project h_l onto a small precomputed top-PCA-of-lm_head subspace (k ≪ vocab)
and watch that converge — the residual may not saturate, but the decision-relevant
projection (the only thing that needs to be stable) might, at a fraction of a full
lm_head's cost. Also re-profile on real tokenized prompts (the in-crate qwen2_bpe.rs)
rather than arbitrary token IDs. If both fail, the honest standing conclusion holds:
this model family does not admit single-token greedy depth-shortcutting, lossless or lossy.