Reproduce “Attention Drift: What Autoregressive Speculative Decoding Models Learn” on OpenResearch

About this paper

Doğaç Eldenk, Payal Mohapatra, Yigitcan Comlek, Kaan Oktay, Hongyang Zhang, Stephen Xia

The paper examines attention drift in EAGLE-3, where a speculative-decoding drafter’s focus shifts from an attention sink toward recently generated tokens as its chain advances. It attributes this partly to growing hidden-state magnitude and proposes post-normalization, including RMSNorm on target features. On Llama 3.1 8B, post-norm raises format-perturbed acceptance length from 1.34 to 2.59; broader drafter generality remains untested.

Key insights

• Autoregressive speculative decoding drafters consistently exhibit "attention drift," progressively shifting their attention mass from the initial prompt's sink token to their own self-generated tokens.
• The primary cause of attention drift and drafter instability is the monotonic growth of hidden-state magnitudes in pre-norm architectures, making drafters behave like stacked additional transformer layers rather than stable autoregressive predictors.
• Post-normalization effectively prevents hidden-state magnitude accumulation, enabling drafters to generalize from shallow training depths to much longer inference depths while maintaining stable performance.

Summary by alphaXiv.