Attention Drift: What Autoregressive Speculative Decoding Models Learn

TL;DR

This paper investigates attention drift in autoregressive speculative decoding models, identifies its cause related to residual paths, and proposes architectural modifications that significantly improve model robustness and performance.

Contribution

It uncovers the phenomenon of attention drift in speculative decoding, traces its cause to residual path dynamics, and introduces architectural changes that enhance model robustness.

Findings

Attention drift causes attention to shift from prompt to generated tokens.

Architectural modifications like post-norm and RMSNorm mitigate attention drift.

Proposed changes double acceptance length and improve performance on benchmarks.

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.