PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Tianyu Liu, Yuhao Shen, Xinyi Hu, Baolin Zhang, Hengxin Zhang, Jun Dai, Jun Zhang, Shuang Ge, Lei Chen, Yue Li, MingCheng Wan

arXiv:2604.26412v2cs.CL

TL;DR

Long-range draft accuracy declines as speculative steps increase, and this decay persists even in TTT-trained drafters. The paper uses KVShot to compare hidden-only, KV-only, and hybrid reuse, finding improved long-range acceptance but only marginal end-to-end gains under current pipelines.

  • Problem

    Hidden-state-based drafters lose accuracy at longer speculative steps, and the decay persists despite TTT addressing train-inference mismatch.

    Source: §1 · pp. 1–2

  • Method

    KVShot compares hidden-only, KV-only, and hybrid reuse in a systematic study of representation reuse for speculative drafting.

    Source: §1 · pp. 2–3, §2 · p. 3

  • Results

    KV-only drafters become relatively more competitive at longer steps, while the hybrid drafter improves draft-acceptance MAT over EAGLE-3 (2.54 vs. 2.37); HF-measured MAT rises only from 5.01 to 5.04 as drafting latency increases by 5–10%.

    Source: §1 · p. 2, §6 · p. 13

  • Takeaways & Limitations

    The identified bottlenecks motivate direct evaluation of block-wise training, which may better fit KV-aware drafting than sequential autoregressive TTT.

    Source: §6 · p. 14

    Autoregressive TTT faces query-estimation difficulty, sparse gradients for draft-side KV projections, and gate-induced gradient starvation in warm-started hybrid drafters.

    Source: §6 · p. 14

Abstract

from arXiv · show

Speculative decoding accelerates LLM inference, but SOTA hidden-state-based drafters suffer from long-range decay: draft accuracy degrades as the speculative step increases. Existing work attributes this decay to train-inference mismatch and proposes test-time training (TTT) as a remedy, yet we observe that long-range decay persists even in TTT-trained drafters. We revisit long-range decay from the perspective of context information preservation. In hidden-state reuse, we argue the target hidden state acts as a biased context compression: it aggregates historical token information according to the attention query at the current position, yielding a compact representation optimized for immediate next-token prediction. This compression can suppress information less relevant to the current query but important for later speculative steps. In contrast, the target model's KV cache serves as an explicit context, retaining the complete set of token-wise KV representations. We therefore posit the KV-Reuse Hypothesis: allowing the draft model to reuse the target KV cache can provide richer signals for long-horizon drafting. To test this hypothesis, we introduce KVShot, a diagnostic framework that compares three reuse paradigms: hidden-only, KV-only, and hybrid. Extensive evaluations on Qwen3-8B show that KV-Reuse improves long-range acceptance, although end-to-end speedups remain marginal under current training pipelines. Our analysis identifies two key structural bottlenecks: shallow drafters struggle to estimate target queries accurately, and draft-side KV projections receive sparse gradient signals. These findings suggest that realizing the full potential of KV-aware decoding requires moving beyond TTT toward block-wise training paradigms. By exposing these bottlenecks, KVShot provides a foundational diagnostic testbed and a clear roadmap for designing next-generation inference architectures.

1. Introduction

Hidden-state drafters lose accuracy at longer speculative steps, and this decay persists even with TTT. The paper proposes that KV reuse better preserves context, tests that idea with KVShot, and finds improved long-range acceptance but no significant end-to-end speedup under current pipelines.

1. Introduction

  • Long-range decay persists in TTT-trained drafters, so train-inference mismatch alone cannot explain the declining acceptance at later speculative steps.The figure describes decreasing acceptance as the speculative step increases for both a Qwen3.5-4B MTP drafter and EAGLE-3.Source: §1 · pp. 1–2, Figure 1 · p. 2
  • Hidden states compress context according to the current attention query, potentially suppressing information useful for later speculative steps.KV caches retain token-level key/value pairs before attention aggregation, allowing the drafter to re-attend to prefix positions using estimated future queries.Source: §1 · p. 2
  • KVShot compares hidden-only, KV-only, and hybrid reuse under the same autoregressive TTT pipeline to isolate representation choice.The comparison uses a matched hidden-only EAGLE-3 baseline on Qwen3-8B.Source: §1 · pp. 2–3
  • KV reuse degrades more gracefully than hidden-state reuse at longer draft steps, but the gains remain too small for significant end-to-end speedup.The paper attributes the gap to query-estimation difficulty and sparse gradients for draft-side KV projections.Source: §1 · pp. 2–3

2. Context Information Preservation View of Long-Range Decay

The section explains long-range decay as a context-preservation problem: hidden states compress information for the current query, while KV reuse preserves token-level context for future re-attention. This shifts the main challenge toward estimating future queries and motivates testing KV reuse, its bottlenecks, and hybrid designs.

2.1. Hidden States as Biased Context Compression

  • Hidden-state reuse compresses context according to the current attention query, suppressing weakly attended information that may be needed for later-token prediction and making recovery ill-conditioned.The analysis focuses on attention aggregation; it omits output projection, residual connections, layer normalization, and the feed-forward sub-layer.Source: §2.1 · pp. 3–4
  • Multi-layer hidden-state fusion can partially recover information suppressed at one layer, but remains a fixed-dimensional aggregation that may omit information unattended across all fused layers.Different layers can attend to different aspects of the input.Source: §2.1 · p. 4

2.2. KV Cache as Re-attention Without Prior Aggregation

  • KV reuse preserves per-position key/value pairs for fresh draft-side re-attention, shifting approximation difficulty toward estimating a useful future query rather than recovering information lost to prior aggregation.This is a simplified conceptual account, not a formal error bound.Source: §2.2 · pp. 4–5

2.3. Are Hidden States and KV Cache Equivalent?

  • Although KV caches derive from preceding hidden states, top-layer information is already aggregated, and reconstructing KV projections from hidden states adds work that competes with query estimation for shallow drafter capacity.The section argues that direct KV reuse removes the projection-reconstruction responsibility.Source: §2.3 · p. 5

2.4. Comparing the Two Error Regimes

  • Hidden-state error combines target-side compression loss, present from k=1, with recursive draft-state drift that grows with k; TTT primarily targets drift, not compression loss.The two sources are coupled because discarded information in the initial compressed state can propagate through later draft steps.Source: §2.4 · p. 5
  • Under KV reuse, compression loss is reduced in the simplified view, but realizing the benefit still requires useful future queries from limited drafter depth and training signal.This reframes the task from information recovery to query function approximation.Source: §2.4 · p. 6, §2.5 · p. 6

2.5. KV-Reuse Hypothesis

  • KV reuse should degrade more gracefully at later speculative steps, but its benefit depends on query-estimation capacity and hidden-state reuse may retain a short-range edge.The authors predict that KV reuse is most advantageous at k≥3 and that hybrid reuse may be most valuable if all three predictions hold.Source: §2.5 · p. 6

3. KVShot: Testing the KV-Reuse Hypothesis

KVShot compares hidden-only, KV-only, and hybrid drafting under a shared autoregressive TTT setup, testing whether target KV reuse improves long-range acceptance. KV reuse helps at longer speculative steps, but its step-wise gains do not yield meaningful end-to-end speedups in the evaluated pipeline.

3.1. Experimental Setup

  • Experiments use Qwen3-8B, autoregressive TTT, and a matched 1-layer EAGLE-3 baseline; initial ablations use about 70k ShareGPT samples, while end-to-end training scales to 280k samples.The larger end-to-end dataset combines ShareGPT and UltraChat with target-model-regenerated responses.Source: §3.1 · p. 6

3.2. KV-Only Reuse: Does Re-attention Help?

  • KV-only ablations show linear projection beats head concatenation (MAT 1.83 vs. 1.78), while all KV-only variants remain below EAGLE-3 (MAT 2.37).Re-applying RoPE after projection yields only a marginal gain (MAT 1.84 vs. 1.83).Source: Table 1 · p. 7, §3.2 · p. 7

3.3. Depth Scaling: Is Query Estimation the Bottleneck?

  • At k=6, 4-layer KV-only acceptance (0.495) exceeds EAGLE-3 (0.469), while at k=0 it trails (0.614 vs. 0.638); retention reaches 80.6% versus 73.5%.The four-layer drafter’s MAT is 2.34, below the 1-layer EAGLE-3 baseline’s 2.37, and comes at higher drafting cost.Source: Table 2 · p. 8, §3.3 · p. 8

3.4. Hybrid Reuse: Can KV Complement Hidden States?

  • Because hidden states retain a short-range edge while KV reuse is more competitive at long range, the hybrid design keeps hidden-state drafting as an anchor and uses KV as a correction.The study frames this hybrid as a way to combine the two reuse pathways rather than replace hidden states.Source: §3.4 · p. 8

3.4.1. Gated Delta Rule

  • The gated delta rule merges target-KV cross-attention into the hidden representation as a correction, with a shared gate controlling its contribution across draft steps.A near-zero gate reduces the model to standard EAGLE-3; a large gate shifts the representation toward the cross-attention correction.Source: §3.4.1 · pp. 8–9

3.4.2. Results

  • The checkpoint-initialized hybrid reaches MAT 2.54 versus EAGLE-3’s 2.37, with α0 rising from 0.638 to 0.665 and α6 from 0.469 to 0.514.Random initialization reaches MAT 2.44, while Cross-only reaches 2.34–2.35, below EAGLE-3; the self-attention anchor is therefore retained in the stronger hybrid.Source: Table 3 · p. 9, §3.4.2 · p. 9

3.5. End-to-End Evaluation: Does the Gain Survive Overhead?

  • In a controlled same-checkpoint comparison, the hybrid drafter raised HF-measured MAT only from 5.01 to 5.04 (+0.6%), while adding roughly 5–10% drafting latency, so meaningful end-to-end speedup is unsupported.Source: §3.5 · p. 10
  • The end-to-end evaluation scales training to 280k regenerated samples and uses tree verification; the stronger retrained baseline and tree verification may both narrow the visible hybrid gap.The authors do not treat either factor as a complete explanation for the shrinkage from the step-wise gain.Source: §3.5 · p. 10

4. Why Does Current Autoregressive TTT Fit KV Reuse Poorly?

KV reuse’s step-wise advantages largely match the predictions, but current autoregressive TTT struggles to convert them into an end-to-end gain. The section identifies query estimation, sparse KV-projection training signals, and gate-induced gradient starvation as bottlenecks, pointing toward block-wise training.

4. Why Does Current Autoregressive TTT Fit KV Reuse Poorly?

  • KV reuse becomes more competitive at later draft steps, deeper KV-only drafters improve, and hidden-state reuse retains the short-range edge.These predicted step-wise effects still yield only a marginal end-to-end gain under the current pipeline.Source: §4 · p. 10

4.1. Query Estimation Is Harder Than It Appears

  • Moving from 1 to 2 layers raises MAT by +0.39, yet even 4 layers cannot match a 1-layer EAGLE-3 baseline.Query estimation requires capacity to approximate the target’s multilayer query computation, beyond predicting a next hidden state.Source: §4.1 · pp. 10–11

4.2. Sparse Optimization of Draft-Side KV Projections

  • Draft-side KV projections receive sparse, unbalanced gradients because only a few draft positions train them, unlike the query projection’s dense gradients.Scaling KV-projection gradients by 50× yields no meaningful improvement: MAT 2.21 vs. 2.18, suggesting signal diversity—not magnitude—is the issue.Source: §4.2 · p. 11

4.3. Gate-Induced Gradient Starvation

  • In warm-started hybrid drafters, a randomly initialized cross-attention branch can be gated down early, weakening its own learning signal and slowing recovery.The mean gate falls from about 0.5 to 0.02, then recovers to about 0.10; that small mean is not equivalent to ignoring the branch, and starvation alone does not explain the end-to-end gap.Source: §4.3 · p. 11, Figure 4 · p. 12

4.4. Implications for Future Training Pipelines

  • Autoregressive token-by-token training limits query depth, the number of KV-training positions, and early gradients to gated cross-attention.Block-wise drafting may address the first two bottlenecks through deeper query computation and diverse parallel draft-token gradients, but this paper does not verify that hypothesis in a controlled comparison.Source: §4.4 · pp. 11–12

5. Related Work

Prior speculative-decoding work explores verification strategies and drafter designs, including hidden-state and KV-cache reuse. This paper distinguishes itself by testing these representations head-to-head under a shared autoregressive TTT setup and examining the resulting training bottlenecks.

Speculative decoding

  • Speculative-decoding work spans token-tree verification, self-speculative drafting, adaptive draft lengths, retrieval-based methods, and exact parallel alternatives.Examples include SpecInfer, Sequoia, PEARL [Liu et al., 2025], and Lookahead Decoding.Source: §5 · pp. 12–13

Drafters that reuse target hidden states

Drafters that reuse target KV cache

  • GLIDE+CAPE and LongSpec use cross-attention to target KV, but focus respectively on low-overhead drafting and long-context memory efficiency.Unlike those systems, this work compares KV and hidden-state reuse, includes hybrids, and evaluates under EAGLE-3’s autoregressive TTT objective.Source: §5 · p. 13

Beyond autoregressive drafting

  • DFlash [Chen et al., 2026] uses block diffusion drafting, a direction consistent with addressing autoregressive TTT’s gradient-sparsity bottleneck.The paper’s findings suggest the training setup matters at least as much as the choice of reused representation.Source: §5 · p. 13

6. Conclusion

The paper introduces KVShot to test whether target KV reuse can address long-range decay, contrasting its context preservation with hidden-state compression. KV reuse improves long-range drafting, but end-to-end gains remain small under the evaluated pipeline, motivating further work on training setups.

6. Conclusion

  • KVShot evaluates reuse paradigms experimentally on Qwen3-8B and frames hidden-state reuse as compression versus KV reuse as access to key/value pairs through re-attention.With KV reuse, query estimation becomes the dominant difficulty in the simplified account.Source: §6 · p. 13
  • MAT is 2.54 vs. 2.37 for hybrid drafting over EAGLE-3, while HF-measured MAT rises only from 5.01 to 5.04 and drafting latency increases by 5–10%.KV-only drafters also become relatively more competitive at longer draft steps.Source: §6 · p. 13
  • The paper traces limitations under autoregressive TTT to shallow-drafter query estimation, sparse KV-projection gradients, and gate-induced starvation in warm-started hybrids.The authors propose block-wise training as a direction for future evaluation, not as a demonstrated solution.Source: §6 · p. 14
  • KVShot separates the availability of useful KV-cache information from the training setup’s ability to use it efficiently.The authors identify efficient exploitation of that signal as the remaining challenge.Source: §6 · p. 14

A. Training Details

The training details specify the drafter objective and target model, KV-layer sampling, datasets, and evaluation measures.

A. Training Details

  • All drafters use EAGLE-3’s autoregressive test-time training objective, with Qwen3-8B as the target model.The cited models are Li et al., 2025c, and Qwen Team et al., 2025.Source: A. Training Details · p. 17
  • Unless noted otherwise, KVShot samples KV cache from 3 uniformly spaced target layers, matching EAGLE-3’s fused-layer count.Varying the sampled layer count has negligible effect, according to Section D.Source: A. Training Details · p. 17
  • Ablations use about 70k ShareGPT samples for 3 epochs; end-to-end evaluation uses 280k ShareGPT and UltraChat samples with target-regenerated responses.Regeneration is used to ensure distributional alignment.Source: A. Training Details · p. 17
  • Evaluation reports held-out step-wise acceptance rates and end-to-end MAT using HuggingFace speculative decoding with a (8, 10, 60) draft tree.Source: A. Training Details · p. 17

B. Gated KV Fusion: Architecture Details

The section details gated KV fusion within a single draft layer, keeping the EAGLE-3 input pathway while adding cross-attention over target and draft KV representations. A learned gate applies cross-attention as a correction to the self-attention output.

Input representation

  • The input representation remains unchanged from EAGLE-3, with all changes confined to attention in the single draft layer.Each token input concatenates the input embedding and target hidden states, then projects them.Source: B. Gated KV Fusion: Architecture Details · p. 17

Cross-attention memory

  • Cross-attention uses copied target KV pairs for verified prefix tokens and learned draft-side KV projections for subsequent positions.The prefix KV is available without estimation, whereas draft-generated keys and values must be learned.Source: B. Gated KV Fusion: Architecture Details · p. 17

Why a gated delta, not a direct sum

  • The gated delta treats cross-attention as a correction, reducing the risk that it overrides the self-attention path’s short-range accuracy.It projects the difference between cross- and self-attention outputs, then scales it element-wise with a learned gate.Source: B. Gated KV Fusion: Architecture Details · p. 17

Shared gate

  • The gate is shared across draft steps to test the mechanism without step-specific gating and to reduce parameters and stabilize training.The stated stabilization motivation is especially relevant because the cross-attention branch is randomly initialized.Source: B. Gated KV Fusion: Architecture Details · p. 17

Initialization from EAGLE-3 checkpoint

  • EAGLE-3 warm-starts retain the self-attention and shared-module weights, while cross-attention, gate, and delta parameters are randomly initialized; the Cross-only variant omits self-attention output.The Cross-only setup isolates the hidden-state pathway’s contribution from the initialization choice.Source: B. Gated KV Fusion: Architecture Details · p. 17

C. Full Step-wise Acceptance Rates

The complete seven-step acceptance curves show that Gated KV (ckpt) stays above the EAGLE-3 baseline at every step. They also show that Head concat falls behind Linear projection early in drafting.

Full curves

  • Gated KV (ckpt) remains above the EAGLE-3 baseline at every step, so its improvement is not confined to one speculative step.Table 5 provides the complete 7-step acceptance rates, α0 through α6, for the main-table experiments.Source: C. Full Step-wise Acceptance Rates · p. 18
  • At α1, Head concat reaches 0.393 versus 0.425 for Linear projection, showing faster early-step decay for Head concat.The passage attributes the early gap to difficulty attending across disjoint KV spaces.Source: C. Full Step-wise Acceptance Rates · p. 18
  • The full curves indicate that the Head concat difficulty appears early rather than emerging only at long range.Source: C. Full Step-wise Acceptance Rates · p. 18

D. Additional Ablations

The additional ablations find limited gains from adding a sampled target layer and lower MAT with direct KV reuse and online TTT than with the tested alternatives. Gradient scaling and QK-Norm do not resolve the reported training bottleneck.

Target KV configuration

  • Increasing sampled target layers from 3 to 4 yields only a marginal MAT gain, from 2.23 to 2.24.The passage identifies query estimation, rather than KV-input richness, as the drafter’s bottleneck.Source: D. Additional Ablations · p. 18

Projector architecture

  • MAT is 2.23 with linear projection, versus 2.20 and 2.19 with two-layer MLP variants and 2.04 when projecting hidden states into KV space.The passage says the pre-computed KV representation matters and links MLP optimization difficulty to sparse gradients.Source: D. Additional Ablations · p. 18

Training regime

  • Offline training reduces MAT from 2.23 to 2.18, while 50× gradient scaling recovers only a small portion of the gap at 2.21; QK-Norm has no measurable change at 2.23.The passage identifies training-signal diversity, rather than gradient magnitude, as the core issue.Source: D. Additional Ablations · p. 18

Additional ablations

  • The ablation table covers 2-layer pure KV reuse with ShareGPT training and defines retention as the long-range ratio α6/α0.Source: Table 6 · p. 19

Related papers

Loading 2604.26412v2…