arXiv:2608.06776cs.LG2026-08中稿 · ICML

通过加速查询键学习,让注意力更聚焦关键信息。

Faster Query-Key Learning Sharpens Attention in Self-Attention Models

论文配图:Faster Query-Key Learning Sharpens Attention in Self-Attention Models
图 1 · 摘自论文原文
  • 让查询键部分比输出值部分学得更快,可增强注意力聚焦能力。
  • 实验显示,相对学习率差异决定注意力集中程度,性能基本不变。
  • 适合关注模型可解释性与注意力机制优化的研究者。

标准自注意力层包含两个交互模块:负责注意力分配的查询-键电路,以及将注意力后的表示映射到预测的输出-值电路。对这两个电路进行压缩和因子化参数化会带来不同的注意力模式。某些参数化方式能在相近训练损失下使注意力更聚焦于任务相关标记。通过梯度流分析,我们发现因子化会隐式调整两个电路的学习率比例。推导出闭式动态表明,输出-值与查询-键参数沿直线移动,相对速度由学习率决定。当查询-键学习速度高于输出-值时,模型会通过增加对关键标记的关注来补偿输出-值学习较慢的问题,从而产生更尖锐的注意力。实验验证了两电路相对学习率差异主导注意力集中度,提升了注意力可解释性代理指标,同时保持相近预测性能。

原文摘要 · Abstract (English)

A standard self-attention layer consists of two interacting circuits: the query-key circuit that governs attention allocation, and the output-value circuit that maps attended representations to predictions. Collapsed and factorized parameterizations of the query-key and output-value circuits lead to qualitatively different attention patterns. In particular, some parameterizations give sharper attention to task-relevant tokens, at a similar training loss. We analyze how the parameterizations of these circuits shape the parameter trajectories in single-layer self-attention models trained for next-token prediction. Through gradient-flow analysis, we show that factorization induces implicit rescaling of the two circuits' learning rates. We derive closed-form dynamics showing that output-value and query-key parameters move along a line, with relative speeds determined by their learning rates. Faster query-key learning relative to output-value learning thus produces sharper attention, as the model compensates for slower output-value learning by increasing attention mass on relevant tokens. Experiments show that differences in the relative learning rates of the two circuits govern attention concentration. This improves attention interpretability proxies while maintaining comparable predictive performance.

自注意力注意力机制模型优化

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