更快的查询-键学习锐化自注意力模型中的注意力

# 更快的查询-键学习锐化自注意力模型中的注意力

来源：https://arxiv.org/html/2608.06776

###### 摘要

标准自注意力层由两个相互作用的回路组成：控制注意力分配的查询-键回路，以及将注意力表示映射到预测的输出-值回路。查询-键回路和输出-值回路的折叠参数化与分解参数化会产生性质不同的注意力模式。特别地，某些参数化能够在训练损失相近的情况下，将注意力更聚焦于与任务相关的词元。我们分析了这些参数化如何影响在下一词元预测上训练的单层自注意力模型中的参数轨迹。通过梯度流分析，我们表明分解会隐式地重新缩放两个回路的学习率。我们推导出闭式动力学，表明输出-值和查询-键参数沿一条直线移动，其相对速度由它们的学习率决定。因此，相对于输出-值学习，更快的查询-键学习会产生更锐利的注意力，因为模型通过增加对相关词元的注意力质量来补偿较慢的输出-值学习。实验表明，两个回路相对学习率的差异控制着注意力集中度。这在保持可比较预测性能的同时，改善了注意力可解释性代理指标。

机器学习，ICML

## 1 引言

Transformer模型（Vaswani等人，2017 (https://arxiv.org/html/2608.06776#bib.bib32)）现在已成为语言（Brown等人，2020 (https://arxiv.org/html/2608.06776#bib.bib7)）、视觉（Dosovitskiy等人，2020 (https://arxiv.org/html/2608.06776#bib.bib12)）和语音（Latif等人，2023 (https://arxiv.org/html/2608.06776#bib.bib18)）领域的标准。它们的成功源于自注意力机制，该机制产生上下文表示，并通常使用下一词元预测进行训练。实证研究表明，预测性能相似的已训练模型可能表现出各种各样的注意力模式（Jain & Wallace,2019 (https://arxiv.org/html/2608.06776#bib.bib16); Wiegreffe & Pinter,2019 (https://arxiv.org/html/2608.06776#bib.bib33); Serrano & Smith,2019 (https://arxiv.org/html/2608.06776#bib.bib29)）。先前的工作也研究了注意力权重可能与基于注意力的解释一致或不一致的正式设置（Pandey等人，2023 (https://arxiv.org/html/2608.06776#bib.bib25)）。然而，架构和参数化选择在这种注意力模式的出现中所起的作用仍不清楚。

现有的理论工作已经证明Transformer具有高度表达能力，包括在合适的参数设置下的普适性和图灵完备性（Yun等人，2019 (https://arxiv.org/html/2608.06776#bib.bib34); Bhattamishra等人，2020a (https://arxiv.org/html/2608.06776#bib.bib4),b (https://arxiv.org/html/2608.06776#bib.bib5); Dehghani等人，2018 (https://arxiv.org/html/2608.06776#bib.bib8); Pérez等人，2021 (https://arxiv.org/html/2608.06776#bib.bib26)）。然而，这些结果依赖于理想化的参数构造，因此无法解释在标准基于梯度的优化下注意力结构如何出现。因此，仅靠表达能力对注意力机制和预测层在训练过程中如何共同演化提供的见解有限。

在这项工作中，我们分析了为下一词元预测而训练的单层Transformer的训练动力学（Radford等人，2018 (https://arxiv.org/html/2608.06776#bib.bib27)）。该模型自然分解为两个相互作用的线性回路（Elhage等人，2021 (https://arxiv.org/html/2608.06776#bib.bib14); Olsson等人，2022 (https://arxiv.org/html/2608.06776#bib.bib24); Bietti等人，2023 (https://arxiv.org/html/2608.06776#bib.bib6)）：

- **查询-键回路** \(W_K W_Q^\top\)，决定注意力模式，以及
- **输出-值回路** \(W_O W_V\)，将注意力表示映射到预测。

参见说明(a) 单层自注意力模型示意图
参见说明(b) 合成数据示意图

图1：总体设置：(a) 序列 \(\mathbf{X}\)（包含上下文词元与查询词元 \(\{x_1,x_2,\ldots,x_{T-1},x_T\}\)）输入到1层Transformer中，以预测下一词元 \(y=x_{T+1}\)；(b) 合成数据示意图。对于每个序列类别，存在一组该类别特有的词元，称为独特词元，以及一组所有序列类别共有的词元，称为常见词元。例如，‘the’、‘sat’、‘ate’和‘on’是常见词元，而‘cat’、‘mat’是类别1特有的词元。

我们表明，这两个回路的相对优化速度在塑造注意力结构方面起着核心作用。改变查询-键和输出-值参数的学习率会产生性质不同的注意力模式，即使预测性能保持相似。特别是，查询-键回路的更快学习会导致对任务相关词元的注意力更锐利。

本文的一个关键焦点是参数化的作用。我们比较查询-键和输出-值回路的分解与折叠参数化，并表明分解会在相应的折叠参数空间中引起梯度更新的状态相关重新缩放。因此，即使分解模型和折叠模型取得可比较的损失，它们也可能遵循不同的优化轨迹。

然后，我们在受控的合成环境中刻画折叠参数化下的训练动力学。在正交性和对称性假设下，我们推导出两个回路的闭式总体梯度流动力学，并表明它们的有效参数以不同速率演化。特别是，输出-值标量按 \(\log(t)\) 量级增长，而查询-键标量按 \(\log^2(t)\) 量级增长，其常数由学习率之比 \(\eta_{QK}/\eta_{OV}\) 控制。由于查询-键回路通过饱和非线性影响注意力，这种增长差异为相对学习速度如何在训练过程中塑造注意力锐化提供了机制性解释。

我们用合成数据和带有词元级相关性标注的真实世界数据集（包括HateXplain（Mathew等人，2020 (https://arxiv.org/html/2608.06776#bib.bib23)）、主谓一致（Linzen等人，2016 (https://arxiv.org/html/2608.06776#bib.bib21)）和SQuAD（Rajpurkar等人，2016 (https://arxiv.org/html/2608.06776#bib.bib28)））的实验来支持该理论。在各种设置下，我们观察到相对优化速度和参数化的变化以理论预测的方式改变注意力结构，而预测性能基本保持不变。

## 2 相关工作

参见说明(a) FAFO
参见说明(b) FACO
参见说明(c) CAFO
参见说明(d) CACO
参见说明(e) FAFO，QK学习率增加10×
参见说明(f) FAFO，QK学习率线性增加1-20×
参见说明(g) CACO，QK学习率增加5×
参见说明(h) CACO，QK学习率线性增加1-10×

图2：在SGD下，不同参数化在测试数据上的独特词元注意力-预测热力图，使用相同学习率（上一行）和更快查询-键学习（下一行）。所有模型均训练到负对数似然接近零（≤0.001）。此处，我们训练模型5个随机种子并报告平均值，热力图标准差参见附录。

许多近期工作研究了注意力和基于自注意力的模型的优化动力学，尤其是在简化或单层设置中（Tian等人，2023 (https://arxiv.org/html/2608.06776#bib.bib30); Deora等人，2024a (https://arxiv.org/html/2608.06776#bib.bib9); Lu等人，2021 (https://arxiv.org/html/2608.06776#bib.bib22); Vashisht & Ramaswamy,2023 (https://arxiv.org/html/2608.06776#bib.bib31)）。Tian等人（2023 (https://arxiv.org/html/2608.06776#bib.bib30)）在输出-值参数比查询-键参数学习更快的假设下分析注意力动力学。Deora等人（2024a (https://arxiv.org/html/2608.06776#bib.bib9)）在数据的可实现性假设下研究多头自注意力模型的优化和泛化性质。在这些分析中，优化动力学通常在关于注意力组件的相对学习行为的固定或隐式假设下进行检验。相比之下，我们改变查询-键和输出-值回路的相对学习率，并表明查询-键参数的更快学习会锐化注意力。虽然先前的工作通常为了分析可处理性而关注单样本或批大小为1的设置，但我们分析的全批次训练动力学更贴近现代实践。此外，我们将分析从注意力机制扩展到最终预测层，为相互作用回路如何塑造训练轨迹提供了统一视角。

一个相关的工作方向对Transformer模型采用机制性视角。Elhage等人（2021 (https://arxiv.org/html/2608.06776#bib.bib14)）的Transformer回路框架提供了一种将注意力层系统地分解为相互作用组件的方法，这激发了后续对训练模型中结构如何出现的研究，包括注意力模式中的聚类行为（Geshkovski等人，2023 (https://arxiv.org/html/2608.06776#bib.bib15); Karagodin等人，2024 (https://arxiv.org/html/2608.06776#bib.bib17)）、类记忆机制（Bietti等人，2023 (https://arxiv.org/html/2608.06776#bib.bib6)）以及主题结构的涌现（Li等人，2023 (https://arxiv.org/html/2608.06776#bib.bib19)）。这些工作为训练后注意力头和所学表示的功能角色提供了宝贵见解。然而，它们主要事后刻画注意力结构，并未分析优化动力学如何在训练过程中塑造相互作用回路的相对演化。我们的工作则侧重于查询-键和输出-值回路的训练轨迹如何在基于梯度的优化下共同决定注意力锐化。

## 3 问题设置

我们研究一个为下一词元预测训练的单层自注意力模型。设 \(\mathbf{X}=[x_1,\ldots,x_T]^\top \in \mathbb{R}^{T\times d}\) 表示输入序列，其中 \(x_T\) 是查询词元。模型参数为 \(W_Q,W_K,W_V\in\mathbb{R}^{d\times d}\) 和 \(W_O\in\mathbb{R}^{M\times d}\)，其中 \(M\) 是词汇表大小。

给定序列 \(\mathbf{X}\)，模型输出为

\[
\mathbf{f}(\mathbf{X}) = W_O W_V \mathbf{X}^\top \mathbb{S}\!\left(\mathbf{X} W_K W_Q^\top x_T\right), \tag{1}
\]

其中 \(\mathbb{S}(\cdot)\) 表示softmax算子。图1(a) (https://arxiv.org/html/2608.06776#S1.F1.sf1) 展示了单层自注意力模型的计算过程。

模型自然分解为两个线性回路。查询-键回路 \(W_K W_Q^\top\) 决定词元上的注意力分布，而输出-值回路 \(W_O W_V\) 将注意力表示映射到预测logits。我们研究四种参数化，它们在回路是分解还是折叠方面有所不同：

- **FAFO**：分解注意力与分解输出-值。
- **CAFO**：折叠注意力（\(W_{QK}=W_K W_Q^\top\)）与分解输出-值。
- **FACO**：分解注意力与折叠输出-值（\(W_{OV}=W_O W_V\)）。
- **CACO**：折叠注意力与折叠输出-值。

设 \(\mathbf{E}=[e_1,\ldots,e_M]^\top\in\mathbb{R}^{M\times d}\) 表示嵌入矩阵。给定一个带有下一词元标签 \(y\in[M]\) 的序列 \(\mathbf{X}\)，训练目标是最小化负对数似然

\[
\mathcal{L}=-\log \mathbb{S}_y(\mathbf{f}(\mathbf{X})). \tag{2}
\]

为了在受控环境中分析注意力行为，我们采用基于Tian等人（2023 (https://arxiv.org/html/2608.06776#bib.bib30)）的合成数据构造。每个序列类别与固定的查询和下一词元对 \((q,n)\) 相关联，其中下一词元 \(n\) 唯一标识该类别。查询之前的上下文词元从类别特定的*独特*词元与类别无关的*常见*词元的混合中采样，如图1(b) (https://arxiv.org/html/2608.06776#S1.F1.sf2) 所示。

例如，考虑三个序列类别：

- 类别1：“cat sat on the mat”、“cat ate on the mat”，
- 类别2：“dog sat on the rug”、“dog ate on the rug”，
- 类别3：“bird sat on the branch”、“bird ate on the branch”。

这里，“sat”、“ate”和“on”等词元出现在所有类别中，属于常见词元，而“cat”、“mat”、“dog”和“rug”等词元是类别特有的，属于独特词元。

形式上，定义 \(\Omega(l)=\{n:\mathbb{P}(l\mid n)>0\}\)。满足 \(|\Omega(l)|=1\) 的词元是独特词元，而满足 \(|\Omega(l)|>1\) 的词元是常见词元。上下文词元根据以下分布采样：

\[
\mathbb{P}(l\mid n)=
\begin{cases}
\gamma/C, & |\Omega(l)|=1,\\
(1-\gamma)/D, & |\Omega(l)|>1,\\
0, & \text{otherwise},
\end{cases}
\]

其中 \(C\) 和 \(D\) 分别表示独特和常见词元的数量，\(\gamma\in(0,1)\) 控制混合权重。这种构造将信息性词元与非信息性词元分离，并能够精确分析在不同参数化和优化速度下注意力和预测如何演化。

## 4 不同参数化下的自注意力分析

我们观察到一个启发我们理论的实证现象。即使是在相同的优化设置下训练并取得可比较的预测性能，不同的自注意力参数化也会表现出性质不同的注意力模式。

我们考虑一个包含四个序列类别的合成下一词元预测任务。每个序列包含 \(m\) 个类别特定词元（对预测有信息性）、\(n\) 个常见词元（没有类别信息）以及一个查询词元。

为了评估注意力与预测之间的一致性，我们引入了独特词元注意力-预测（DTAP）热力图。DTAP是一个二维直方图，其中x轴表示分配给类别特定词元的注意力比例，y轴表示正确类别的预测概率。每个格子报告落入相应注意力和置信度范围的示例比例。理想的模型将大部分质量集中在右上角，表明对信息性词元有强注意力且预测置信度高。

我们比较四种参数化：