forgo.cloud
Sign in
Repo workspace

forkjoin-ai/gnosis

Post-mortem: layer early-exit does not pay on this model family

distributed-inference/docs/layer-early-exit-postmortem.md
forkjoin-ai/gnosis

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:

  1. 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.
  2. 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 the 2·τ bound).

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 + cheap relative_l2_delta saturation probe.
  • src/bin/bench-early-exit.rs — warmup, determinism self-check, cost-model dump, ε/θ sweep with knee detection, and --profile-only to 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_one is 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,20

A 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.