Attention Drift: What Autoregressive Speculative Decoding Models Learn

Abstract: Speculative decoding accelerates LLM inference by drafting future tokens with a small model, but drafter models degrade sharply under template perturbation and long-context inputs. We identify a previously-unreported phenomenon we call \textbf{attention drift}: as the drafter generates successive tokens within a speculation chain, attention progressively moves from the prompt onto its own recently-generated tokens. We observe this across both \emph{EAGLE3} drafters and \emph{MTP heads}, suggesting drift is a property of drafter designs. We trace this to the un-normalized residual path between chain steps: the drafter's hidden state magnitude grows monotonically with chain depth, which exhibits dynamics consistent with additional pre-norm transformer layers stacked on the target rather than as a standalone autoregressive predictor. In order to limit the growth, we propose two architectural changes: Post-norm on the drafter hidden states and per-hidden-state RMSNorm after capturing target hidden states. Our interventions improve acceptance length over the current leading model, pre-norm EAGLE3, by up to $2\times$ under template perturbation, $1.18\times$ on long-context tasks, and $1.10\times$ on seven standard benchmarks spanning multi-turn chat, math, and coding. Our changes also allow shorter train-time-test depths to generalize over longer drafting sequences.

• The paper identifies attention drift as a key factor in degrading autoregressive speculative decoding performance due to unnormalized residual pathways.
• It demonstrates through empirical analyses that attention gradually shifts from prompt sinks to recent outputs, undermining decoder stability.
• Architectural interventions like post-norm normalization yield up to 2× improvements in acceptance length and boost robustness under template shifts.

Authoritative Summary of "Attention Drift: What Autoregressive Speculative Decoding Models Learn" (2605.09992)

Introduction and Problem Statement

The paper exposes a phenomenon termed attention drift with profound implications for autoregressive speculative decoding drafters, particularly those used in LLM serving systems such as EAGLE-3 and Multi-Token Prediction (MTP) heads. Speculative decoding achieves throughput gains by having a fast drafter model predict future tokens, subsequently validated by a larger "verifier" model. However, drafter robustness under deployment shift—template variations, long-context settings, and prompt perturbations—remains suboptimal. The authors identify attention drift as a core factor behind this degradation, systematically characterize its emergence, and trace its origins to architectural aspects of the chain-residual path in transformer stacking.

Figure 1: Attention drift: as speculation proceeds, attention migrates from the sink token to recently-generated tokens, visualized across heatmap, graph, and per-token statistics.

Empirical Characterization of Attention Drift

Attention drift manifests when the drafter’s attention distribution transitions away from prompt sinks and onto its own generated tokens during speculative chains, observed across multiple model families (e.g., Llama, Qwen, GPT-oss) and architectures. Notably, in models with pronounced attention sinks, the sink weakens progressively with drafting depth, causing recent tokens to absorb increasing attention mass.

Figure 2: Attention heatmap illustrating drift during speculative decoding for Qwen3.5 35B.

Figure 3: Quantifies attention transition: left shows sink token concentration, right shows latest token acquisition for EAGLE-3 drafters.

This drift is evident not only in post-hoc EAGLE-3 drafters but also in MTP heads, where auxiliary prediction shares weights with the verifier. Drift persists even in architectures designed to suppress sinks (e.g., gated attention or learned biases), instead concentrating attention on template tokens or recurring markers.

Hidden State Pathology and Architectural Cause

The root cause is traced to monotonically increasing hidden-state magnitudes arising from unnormalized residual paths between chain steps. In standard pre-norm architectures, each speculation step adds magnitude, causing the drafter to behave like an extended stack of transformer layers atop the target model, rather than as an autoregressive predictor.

This is quantitatively shown: the drafter’s hidden state RMS grows rapidly with speculation depth, diverging from the magnitude profile seen in verifiers. Attention drift is mechanistically linked to this scale change: query-key similarity shifts, making recent drafter outputs more salient in attention computation.

Figure 4: Heatmap visualization of attention drift on Qwen3.5 9B MTP head.

Figure 5: Aggregated statistics confirming drift across sink and self-token attention for Qwen3.5 9B MTP heads.

Architectural Interventions

Two architectural modifications are proposed to address drift:

• Post-norm on drafter output: RMS normalization applied post-residual, stabilizing state magnitudes across the drafting chain.
• Per-hidden-state RMSNorm before fusion: Normalizes all target hidden states before fusion to remove imbalances.

These interventions pin hidden state magnitude and stabilize the attention pattern. Post-norm drafters generalize robustly to longer speculation depths not seen during training, a critical property for production deployment.

Figure 6: Illustrates sink and self-token attention stabilization for models incorporating the proposed fixes.

Figure 7: Demonstrates magnitude drift on Qwen3.5 9B MTP heads—corrected with interventions.

Performance and Robustness Improvements

Post-norm achieves significant acceptance length gains over pre-norm EAGLE-3: up to 2× improvement under template perturbation, 1.18× on long-context tasks, and 1.10× averaged across seven standard benchmarks. Furthermore, post-norm architectures enable shorter train-time-test depths, reducing training cost and improving inference efficiency.

Figure 8: GPT-OSS 20B acceptance length on SGLang; post-norm consistently outperforms pre-norm across settings.

Template and Long-Context Sensitivity

Drafters are acutely sensitive to template tokens and prompt formatting. Even eliminating attention sinks (e.g. via gated-attention) does not fully rectify this fragility. Only post-norm fixes deliver robust performance across template variants.

Figure 9: Template sensitivity comparison, illustrating post-norm’s resilience to template disruptions for Llama 3.1 8B.

Post-norm also demonstrates robustness to system prompt length variation, maintaining high and stable accuracy even as context grows.

Evaluation in Long-Context and Multi-Turn Scenarios

Long-context failure is severe for pre-norm drafters, especially outside trained window lengths. SWA-based serving and post-norm architecture interaction rescues performance, yielding up to 20% improvement in acceptance length over pre-norm in both synthetic and real-world benchmarks.

Figure 10: Multi-turn chat accuracy for Llama 3.1 8B drafter across long-context scenarios.

Theoretical and Practical Implications

This work recontextualizes drafter design and training. It demonstrates that normalization regime (post-norm vs pre-norm) is not a mere implementation detail but has systemic effects on hidden-state scaling, attention allocation, and robustness. Fixing attention drift (symptom) alone does not guarantee performance recovery; the architectural intervention must address magnitude accumulation (cause).

Practically, these findings inform production drafter design—post-norm should be adopted for robust speculative decoding in systems exposed to template shift and long-context. Theoretically, the results bridge transformer depth stacking and autoregressive token prediction, offering new insight into chain-residual learning dynamics. The post-norm approach may extend to other forms of iterative sequence modeling beyond speculative decoding.

Outlook and Future Directions

The analysis opens further avenues: comprehensive evaluation across Hydra, D-Flash, and alternative drafter/decoder designs; exploration at larger parameter scales; deeper study on training window effects; and investigation into long-context compositional generalization. The naming and mechanistic explanation of attention drift provides the community with a diagnostic lever for both robustness evaluation and model selection criteria.

Conclusion

Attention drift in autoregressive speculative decoding drafters originates from unnormalized chain-residual path magnitude accumulation, causing both attention and template fragility across model families and architectures. Post-norm normalization of drafter outputs and target representation fusion addresses this failure mode, yielding consistent acceptance length improvements, robustness to deployment shift, and generalization to unseen drafting depths. The architectural intervention reframes drafter learning as stable autoregressive prediction rather than depth-specific stacking, suggesting post-norm as a default for robust, scalable, and efficient speculative decoding.