When Hidden States Drift: Can KV Caches Rescue Long-Range Speculative Decoding?

1. Problem Statement

Long-Range Decay in Speculative Decoding:
Speculative decoding accelerates Large Language Model (LLM) inference by using a lightweight "drafter" model to propose multiple tokens, which a larger "target" model verifies. Current state-of-the-art drafters (e.g., EAGLE-3, MTP) rely on hidden-state reuse, feeding the target model's internal hidden states into the drafter.

• The Issue: These systems suffer from long-range decay: the accuracy of the draft drops progressively as the speculative step $k$ increases.
• Existing Explanation: The community attributes this to train-inference mismatch (the drafter conditions on its own drifting hidden states during inference, unlike training). While Test-Time Training (TTT) mitigates this, decay persists.
• The Gap: The paper argues that train-inference mismatch is not the sole cause. The fundamental issue lies in context information preservation. Hidden states act as a "biased context compression" optimized for the immediate next token, discarding information relevant to future speculative steps.

2. Core Hypothesis & Theoretical Analysis

The authors propose the KV-Reuse Hypothesis: Reusing the target model's KV Cache (Key-Value pairs) rather than just hidden states can preserve richer context for long-range drafting.

• Hidden-State Reuse (Biased Compression):
  • Target hidden states are the result of attention aggregation: $h_t = \sum \alpha_i v_i$.
  • The weights $\alpha_i$ are optimized for predicting the current token $x_{t+1}$.
  • Tokens weakly relevant to $x_{t+1}$ receive near-zero weights and are effectively discarded in $h_t$.
  • Consequence: For later speculative steps ($x_{t+k}$), the drafter must recover information that was suppressed by the target's attention mechanism, an ill-conditioned "information recovery" problem.
• KV-Reuse (Explicit Context):
  • The KV cache retains the full set of $(k_i, v_i)$ pairs before aggregation.
  • The drafter can perform re-attention using its own estimated future queries against the full prefix KV cache.
  • Consequence: No information is discarded prior to the drafter's processing. The difficulty shifts from "recovering lost info" to "accurately estimating future queries" (a function approximation problem).

3. Methodology: The KVShot Framework

To test this hypothesis, the authors introduce KVShot, a diagnostic framework comparing three reuse paradigms under the same Autoregressive Test-Time Training (TTT) pipeline (used by EAGLE-3):

1. Hidden-Only Reuse: Standard EAGLE-3/MTP approach (concatenating target hidden states).
2. KV-Only Reuse: The drafter receives the target KV cache via cross-attention, replacing hidden states entirely.
3. Hybrid Reuse: A gated fusion mechanism where the hidden state provides the main anchor, and the KV cache provides a "gated delta correction" for long-range steps.

Experimental Setup:

• Target Model: Qwen3-8B.
• Datasets: ShareGPT and UltraChat (scaled from 70k to 280k samples).
• Metrics: Step-wise acceptance rates ($\alpha_k$), Mean Accepted Tokens (MAT), and End-to-End (E2E) speedup proxies.

4. Key Results

A. Validation of Predictions (Step-wise Analysis)

The experiments confirmed three theoretical predictions:

1. Long-Range Advantage: KV-Only reuse degrades more gracefully than Hidden-Only reuse as $k$ increases. At step $k=6$, a 4-layer KV-only drafter surpassed the EAGLE-3 baseline ($\alpha_6 = 0.495$ vs. $0.469$).
2. Query Estimation Bottleneck: KV-Only performance is heavily dependent on drafter depth. A 1-layer drafter performed poorly (MAT 1.84) because it cannot approximate complex target queries. Increasing depth to 2 layers yielded a massive jump (MAT 2.23), but 4 layers offered diminishing returns, still trailing the 1-layer EAGLE-3 baseline (MAT 2.37) in overall efficiency.
3. Short-Range Disadvantage: At $k=0$ or $1$, Hidden-State reuse outperforms KV-Only reuse because hidden states carry richer semantic content than raw embeddings used for KV re-attention.

B. Hybrid Performance

The Hybrid Drafter (combining both signals) achieved the best step-wise results:

• MAT: 2.54 (vs. 2.37 for EAGLE-3 baseline).
• Retention: Improved long-range retention ratio from 73.5% to 77.3%.
• Mechanism: The gated delta rule successfully allowed the KV path to correct the hidden-state path at longer horizons without disrupting short-range accuracy.

C. The End-to-End Reality Check

Despite step-wise improvements, the End-to-End (E2E) speedup was marginal:

• HF-Measured MAT: Increased only from 5.01 (EAGLE-3) to 5.04 (Hybrid).
• Latency: The additional cross-attention path introduced 5–10% extra drafting latency.
• Conclusion: The current autoregressive TTT pipeline cannot efficiently exploit the KV signal to overcome the added computational cost.

5. Structural Bottlenecks Identified

The paper identifies three specific reasons why the current pipeline fails to translate step-wise gains into E2E speedups:

1. Query Estimation Difficulty: Estimating future target queries requires replicating the multi-layer nonlinear transformations of the target model. Shallow drafters (1-2 layers) lack the capacity to approximate these queries accurately.
2. Sparse Gradient Signals: In autoregressive TTT, the draft-side KV projections ($W_K, W_V$) are only updated by the few draft-generated tokens (e.g., 7 tokens) per step, while the prefix KV is copied. This results in extremely sparse and unbalanced gradients for the KV pathway.
3. Gate-Induced Gradient Starvation: In the hybrid model, the gate controlling the KV correction tends to collapse to near-zero early in training (because the random-initialized KV branch starts as noise). Once closed, the gate blocks gradients to the KV branch, preventing it from learning to be useful. This creates a "self-reinforcing low-injection state."

6. Significance and Future Directions

• Conceptual Shift: The paper reframes long-range decay not just as a train-inference mismatch, but as an information preservation problem. It proves that hidden states are a lossy compression for long-horizon tasks.
• Diagnostic Tool: KVShot provides a foundational testbed to isolate representation effects from training artifacts.
• Roadmap for Next-Gen Architectures: The findings suggest that Autoregressive TTT is insufficient for KV-aware decoding. To realize the potential of KV reuse, the field must move toward Block-wise Training Paradigms (e.g., DFlash).
  • Block-wise training generates many draft tokens in parallel, providing dense gradient signals for KV projections.
  • It allows for deeper draft models capable of accurate query estimation without sequential latency penalties.

Summary: While KV caches theoretically offer a superior signal for long-range speculative decoding by preserving context, current autoregressive training pipelines fail to leverage this due to query estimation limits and gradient starvation. The path forward lies in block-wise, parallel training architectures.