Attention Drift Discovered

It was confirmed that as a speculative decoding drafter continues generating subsequent tokens, its attention shifts from the prompt toward self-generated tokens, a phenomenon called attention drift.

This phenomenon appeared in both EAGLE3 drafters and MTP heads, and was presented as the cause of the sharp drop in acceptance under template perturbation and long-context inputs.

The cause was a structure in which the unnormalized residual path at each chain step continuously increases hidden state magnitude. The authors interpreted this as dynamics similar to adding a pre-norm transformer layer on top of the target.

Two improvements were proposed.

• Applying Post-norm to the drafter hidden state
• Applying RMSNorm to each hidden state after capturing the target hidden state

The effects were also presented numerically.

• Up to 2x acceptance length under template perturbation
• 1.18x on long-context tasks
• 1.10x across 7 standard benchmarks

Also, even when trained with a short train-time test depth, it showed a tendency to generalize to longer drafting sequences.