Attention Drift: Why Your Speculative Drafter Fails and How Post-Norm Fixes It

TL;DR

Speculative decoding is the "free lunch" of LLM inference, but it often breaks when you change a chat template or extend the context. This paper uncovers Attention Drift: a failure mode where drafters "forget" the prompt and obsess over their own recent tokens. By switching from a Pre-norm to a Post-norm architecture, the authors stabilized the model's energy, doubled acceptance rates in messy environments, and cut training time by 30%.

The Hidden Fragility of Speculative Decoding

Speculative decoding (e.g., EAGLE, Medusa) uses a tiny "drafter" model to guess upcoming tokens, which a giant "target" model then verifies. It’s fast, lossless, and efficient. However, in production, these drafters are notoriously "touchy." Change the system prompt slightly or move to a long conversation, and the acceptance rate (the number of guessed tokens that actually pass verification) collapses.

The authors discovered that during a speculation chain, the drafter's attention progressively migrates away from the "Attention Sink" (the stable anchor tokens at the start of a prompt) and sticks to its own recently generated outputs. They dubbed this Attention Drift.

The "Why": Layer Stacking vs. Autoregression

The researchers found that the root cause is unnormalized residual growth. In standard Pre-norm drafters, the magnitude of the hidden state grows monotonically with each step in the speculation chain.

Instead of learning to be a stable autoregressive predictor (where step 1 and step 10 look the same), the drafter thinks it is a stack of layers. Speculation step 5 looks like "Layer $N+5$" of the target model rather than a new token prediction.

Figure: The standard Pre-norm (Left) fails because it accumulates scale, while the proposed Post-norm (Right) resets the magnitude at each step.

The Fix: Stabilizing the Dynamics

The solution is deceptively simple: Post-normalization.

1. Balanced Fusion: Apply RMSNorm to the target's internal states before they are fused into the drafter. This prevents one layer from drowning out the others.
2. Post-Norm Chain: Applying normalization after the drafter's computation at each step.

This ensures that the "energy" (magnitude) of the hidden state remains constant across the drafting chain. The drafter is forced to learn a depth-invariant rule—making it far more robust to inputs it didn't see during training.

Experimental Results: Robustness Unleashed

The results across Llama 3.1, Qwen, and GPT-OSS families are striking:

• Context Resilience: Under "Template Perturbation" (e.g., removing the BoS token or changing formatting), Pre-norm models lost up to 52% accuracy. Post-norm models lost less than 5%.
• Long Context: On LongBench, the Post-norm architecture outperformed the state-of-the-art EAGLE-3 by 20-25%.
• Training Efficiency: Post-norm drafters generalize so well that they can be trained with shorter sequences (lower Train-Time-Test depth) and still perform perfectly on long sequences during inference.

Figure: Heatmaps showing how attention "drifts" away from the sink (left side of plots) in traditional models vs. the stable pattern in the new architecture.

Critical Insight: Drift is a Symptom, Not the Cause

Interestingly, while "fixing the sink" (via Gated Attention) removes the visual drift, it does not fix the performance collapse. This proves that Attention Drift is a symptom of magnitude growth. If you only fix the attention pattern but let the hidden state scale spiral out of control, the model still fails.

Conclusion

This work shifts the focus of speculative decoding research from "better attention" to "better dynamics." For practitioners, the message is clear: if you are building or deploying speculative drafters (like EAGLE-3), Post-norm is the new default. It delivers a more robust, cheaper-to-train, and faster inference engine that doesn't break the moment the user changes their chat template.