arXiv:2605.09992cs.LGcs.AI2026-05被引 2

发现自回归推测解码中注意力逐渐脱离提示,转向自身生成内容。

Attention Drift: What Autoregressive Speculative Decoding Models Learn

论文配图:Attention Drift: What Autoregressive Speculative Decoding Models Learn
图 1 · 摘自论文原文
  • 通过分析发现,推测模型生成时注意力会随序列推进逐渐从输入提示移向自身生成的文本。
  • 改进架构后,在模板扰动下接受长度提升2倍,长上下文任务提升1.18倍。
  • 适合关注推理加速与模型稳定性的研究人员,尤其在复杂提示场景下有效。

推测解码通过小模型预生成未来词元来加速大语言模型推理,但其在模板扰动和长上下文输入下性能显著下降。我们发现一种此前未被报道的现象:注意力漂移——随着推测链中连续生成词元,注意力逐渐从原始提示转移至自身生成的最近词元。该现象在EAGLE3 drafter和MTP heads中均被观察到,表明其为推测模型设计的共性特征。我们追溯其根源在于链式步骤间的未归一化残差路径:推理解码器隐藏状态幅值随链深度单调增长,表现出类似在目标模型上叠加更多预归一化层而非独立自回归预测器的动力学行为。为限制增长,我们提出两种架构改进:对解码器隐藏状态采用后归一化,以及在捕获目标隐藏状态后对每个隐藏状态应用RMSNorm。实验显示,这些干预使接受长度相较当前领先模型pre-norm EAGLE3提升最高达2倍(模板扰动下)、1.18倍(长上下文任务),并在七个标准基准(涵盖多轮对话、数学、编程)上平均提升1.10倍。此外,较短训练时测试深度的模型也能泛化到更长的推测序列。

原文摘要 · Abstract (English)

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.

推理加速注意力机制自回归模型模型优化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。