Technical Summary: Attention Drift in Autoregressive Speculative Decoding

Problem Statement

Speculative decoding accelerates Large Language Model (LLM) inference by using a lightweight "drafter" model to predict future tokens, which are then verified by a larger "target" model. While effective, current drafter models (specifically EAGLE-3 and Multi-Token Prediction heads) exhibit significant fragility under deployment variations, such as template perturbations and long-context inputs. Practitioners often must retrain drafters for specific inference engines or prompts to maintain performance, suggesting an underlying architectural instability.

The paper identifies a previously unreported phenomenon called Attention Drift. As a drafter generates a sequence of speculative tokens, its attention mechanism progressively shifts away from the "attention sink" (a stable anchor token, typically near the start of the prompt) and concentrates increasingly on its own recently generated tokens. This drift correlates with a degradation in acceptance rates, particularly in long-context scenarios and when input templates deviate from training distributions.

Methodology and Analysis

The authors investigate the root causes of attention drift through a combination of empirical observation and architectural analysis across multiple model families (Llama 3.1, Qwen 3/3.5, and GPT-oss).

1. Phenomenological Observation

Using attention heatmaps and quantitative metrics, the authors demonstrate that:

• Sink Migration: In models with strong attention sinks (e.g., Llama, Qwen), the drafter's attention to the sink token diminishes as the speculation chain deepens, while attention to the most recently drafted token increases.
• Universality: This drift is observed in both EAGLE-3 drafters (trained post-hoc on frozen targets) and MTP heads (jointly trained with the target), suggesting it is a property of autoregressive designs that consume their own outputs.
• Sink Absence: Even in models designed to suppress sinks (e.g., GPT-oss with gated attention or learnable biases), a form of drift occurs where attention concentrates on recurring template tokens rather than the prompt's beginning.

2. Root Cause Analysis: Hidden State Magnitude

The paper traces attention drift to the unnormalized residual path between speculation steps in standard pre-norm architectures:

• Monotonic Growth: In pre-norm transformers, residual connections accumulate magnitude without normalization. The authors show that the drafter's hidden state magnitude ($\|h^*_k\|$) grows monotonically with the speculation depth ($k$).
• Layer Stacking vs. Autoregression: This growth causes the drafter to behave less like a standalone autoregressive predictor and more like a sequence of additional transformer layers stacked on top of the target. At step $k$, the drafter effectively approximates the target model at depth $N+k$, rather than predicting the next token based on a stable distribution.
• Imbalanced Inputs: The fused hidden states from the target model ($h_{FC}$) are often dominated by deeper layers with larger magnitudes, further exacerbating the scale mismatch when fed into the drafter.

3. Proposed Interventions

To counteract magnitude accumulation and stabilize the drafting process, the authors propose two architectural modifications:

1. Post-Norm on Drafter Hidden States: Applying RMSNorm to the drafter's hidden states after the residual connection (between speculation steps) prevents the monotonic growth of magnitude. This forces the drafter to learn a depth-invariant autoregressive rule rather than depth-specific refinements.
2. Per-Hidden-State RMSNorm: Adding RMSNorm layers before fusing the target's hidden states ($h_{low}, h_{mid}, h_{high}$) into the drafter input to balance the magnitude of the incoming signal.

The authors also evaluate Gated Attention (suppressing reliance on a fixed sink) as a separate intervention. They find that while gating eliminates the sink, it does not prevent magnitude growth, and combining gating with post-norm can lead to over-regularization (entropy collapse).

Key Results

The proposed post-norm architecture was evaluated against the current state-of-the-art pre-norm EAGLE-3 across various benchmarks and conditions:

• Template Perturbation: Under template perturbations (e.g., removing chat templates or beginning-of-sequence tokens), the post-norm drafter improved acceptance length by up to 2× compared to pre-norm. Pre-norm models suffered catastrophic drops (up to 52% loss), while post-norm models remained resilient (max 5% loss).
• Long-Context Tasks: On long-context tasks (LongBench, multi-turn chat with sliding window attention), post-norm improved acceptance length by 1.18×. In full-attention long-context settings, pre-norm models collapsed to near-zero acceptance, while post-norm maintained usable performance.
• Standard Benchmarks: Across seven standard benchmarks (multi-turn chat, math, coding), post-norm showed a consistent 1.10× improvement in acceptance length.
• Training Efficiency: Post-norm drafters generalize to inference depths beyond their training horizon. This allows training with shorter "train-time-test" (TTT) depths (e.g., reducing from 8 to 4 steps), cutting training time by roughly one-third without sacrificing performance.
• Noise Robustness: Post-norm architectures tolerate significantly more noise in hidden states (58% acceptance retention at $\alpha=0.5$ noise vs. 5% for pre-norm), suggesting better robustness to quantization and distribution shifts.

Significance and Claims

The paper claims that Attention Drift is a critical, previously unreported failure mode in speculative decoding that stems from the interaction between unnormalized residual paths and autoregressive self-feedback.

• Diagnosis: The authors argue that attention drift is a symptom of a deeper issue: the accumulation of hidden-state magnitude, which causes the drafter to learn depth-dependent transformations rather than stable token prediction.
• Solution: They assert that simply fixing the attention sink (e.g., via gating) is insufficient; the underlying magnitude dynamics must be addressed. The proposed post-norm architecture is presented as a simple, effective fix that stabilizes the drafting process.
• Generalization: The findings suggest that post-norm should be considered a default architectural choice for speculative decoding drafters, particularly for production systems requiring robustness against varying prompts and long contexts.

The paper concludes that naming and characterizing "attention drift" provides a useful diagnostic for drafter robustness, and that architectural interventions targeting magnitude accumulation are essential for reliable speculative decoding.