理解 Self-Attention

什么是 Self-Attention？

Self-Attention（自注意力机制）的核心思想可以概括为一句话：对序列中的每一个 Token，都让它观察序列中允许访问的各个 Token，并根据相关程度，有选择地吸收这些 Token 所提供的信息。

在普通的双向 Self-Attention 中，每个 Token 可以观察整个序列；在 GPT 使用的因果 Self-Attention（Causal Self-Attention）中，每个 Token 只能观察自身以及位于它之前的 Token。Self-Attention 并不是将整个句子压缩成一个单一的向量，而是为序列中的每个 Token 分别生成一个新的上下文表示。

设输入序列由 \(n\) 个元素组成：

\[(t_1, t_2, \ldots, t_n) \]

其中，\(t_i\) 表示第 \(i\) 个词或 Token。经过词嵌入（Word Embedding）、位置编码（Positional Encoding），或者上一层 Transformer 的处理后，每个 Token 对应一个 \(d_{\text{model}}\) 维的隐藏向量：

\[x_i \in \mathbb{R}^{1 \times d_{\text{model}}} \]

将所有 Token 的隐藏表示按行排列，可得到输入矩阵 \(X\)：

\[X = \begin{bmatrix} x_1 \\ x_2 \\ \vdots \\ x_n \end{bmatrix} \in \mathbb{R}^{n \times d_{\text{model}}} \]

其中，第 \(i\) 行 \(x_i\) 表示第 \(i\) 个 Token 在当前层的输入表示。

Self-Attention 的目标是为序列中的每个 Token 生成一个新的向量，使其能够根据当前上下文，动态聚合自身以及其他可访问 Token 所提供的信息。因此，Self-Attention 接收输入矩阵 \(X\)，并输出上下文表示矩阵 \(Z\)：

\[Z = \begin{bmatrix} z_1 \\ z_2 \\ \vdots \\ z_n \end{bmatrix} \in \mathbb{R}^{n \times d_v} \]

其中：

\[z_i \in \mathbb{R}^{1 \times d_v} \]

表示第 \(i\) 个 Token 经过上下文信息聚合后得到的新表示。

从单个 Token 的角度看，Self-Attention 所完成的计算可以概括为：

\[\boxed{ z_i = \sum_{j\in\mathcal A_i}\alpha_{ij} v_j } \]

其中 \(\mathcal A_i\) 表示第 \(i\) 个 Token 允许访问的位置集合。在无掩码 Self-Attention 中，\(\mathcal A_i=\{1,\ldots,n\}\)；在因果 Self-Attention 中，\(\mathcal A_i=\{1,\ldots,i\}\)。\(v_j\) 表示第 \(j\) 个 Token 实际提供的信息；\(\alpha_{ij}\) 表示第 \(i\) 个 Token 从第 \(j\) 个 Token 获取信息的权重。

经过 Softmax 后，注意力权重均为非负数。对于第 \(i\) 个 Query，其所有允许访问位置上的注意力权重之和为 1：

\[\alpha_{ij} \geq 0, \qquad \sum_{j\in\mathcal A_i}\alpha_{ij}=1 \]

也就是说，\(z_i\) 是所有允许访问位置对应的 Value 向量的加权和，只不过不同 Token 对应的权重分布不同。由此引出了 Self-Attention 的关键问题：如何根据第 \(i\) 个 Token 与第 \(j\) 个 Token 之间的关系，计算注意力权重 \(\alpha_{ij}\)？

Query、Key 和 Value 正是为解决这一问题而引入的。

---

Query、Key 和 Value

Self-Attention 中有三个重要概念：

• Query（查询），记作 \(Q\)
• Key（键），记作 \(K\)
• Value（值），记作 \(V\)

可以将它们类比为一个搜索系统：用户输入的搜索内容相当于 Query；每篇网页的标题、标签和关键词相当于 Key；网页真正包含的内容相当于 Value。

搜索过程可以理解为：使用 Query 与所有的 Key 进行比较，判断哪些 Key 与 Query 更匹配，并根据匹配程度赋予不同的权重，最后按照这些权重读取并汇总对应的 Value。因此，可以用一句话概括三者的分工：

\[\boxed{ \text{Q 和 K 决定各位置的信息应当以多大权重参与聚合，V 决定被聚合的信息内容。} } \]

\(Q\)、\(K\)、\(V\) 的计算角色是固定的，但它们并不分别保存人工预先规定的固定语义。它们都是同一个 Token 的隐藏向量经过不同线性变换后得到的三种表示，其具体编码方式由模型在训练过程中自动学习形成。

Self-Attention 使用三个可训练的权重矩阵：

\[W_Q \in \mathbb{R}^{d_{\text{model}} \times d_k}, \qquad W_K \in \mathbb{R}^{d_{\text{model}} \times d_k}, \qquad W_V \in \mathbb{R}^{d_{\text{model}} \times d_v} \]

它们分别对输入矩阵 \(X\) 进行线性变换：

\[\boxed{ Q = X W_Q, \qquad K = X W_K, \qquad V = X W_V } \]

运算后的维度结果为：

\[Q \in \mathbb{R}^{n \times d_k}, \qquad K \in \mathbb{R}^{n \times d_k}, \qquad V \in \mathbb{R}^{n \times d_v} \]

这里令 Query 和 Key 具有相同的特征维度 \(d_k\)，是因为后续需要计算二者的点积。将 \(Q\)、\(K\)、\(V\) 按行展开：

\[Q = \begin{bmatrix} q_1 \\ q_2 \\ \vdots \\ q_n \end{bmatrix} = \begin{bmatrix} x_1 W_Q \\ x_2 W_Q\\ \vdots \\ x_n W_Q \end{bmatrix} = \begin{bmatrix} x_1 \\ x_2 \\ \vdots \\ x_n \end{bmatrix} W_Q = X W_Q \]

\[K = \begin{bmatrix} k_1 \\ k_2 \\ \vdots \\ k_n \end{bmatrix} = \begin{bmatrix} x_1 W_K \\ x_2 W_K\\ \vdots \\ x_n W_K \end{bmatrix} = \begin{bmatrix} x_1 \\ x_2 \\ \vdots \\ x_n \end{bmatrix} W_K = X W_K \]

\[V = \begin{bmatrix} v_1 \\ v_2 \\ \vdots \\ v_n \end{bmatrix} = \begin{bmatrix} x_1 W_V \\ x_2 W_V\\ \vdots \\ x_n W_V \end{bmatrix} = \begin{bmatrix} x_1 \\ x_2 \\ \vdots \\ x_n \end{bmatrix} W_V = X W_V \]

对于第 \(i\) 个 Token，有：

• \(q_i = x_i W_Q \in \mathbb{R}^{1 \times d_k}\) （第 \(i\) 个 Token 的 Query 表示）
• \(k_i = x_i W_K \in \mathbb{R}^{1 \times d_k}\) （第 \(i\) 个 Token 的 Key 表示）
• \(v_i = x_i W_V \in \mathbb{R}^{1 \times d_v}\) （第 \(i\) 个 Token 的 Value 表示）

从功能上看，可以将三者直观地理解为：

• \(q_i\)：第 \(i\) 个 Token 当前需要匹配什么信息。
• \(k_i\)：第 \(i\) 个 Token 提供哪些可用于匹配的特征。
• \(v_i\)：第 \(i\) 个 Token 实际能够传递什么内容。

需要注意的是，真实模型中的语义通常采用分布式编码，不能简单认为某一个维度固定表示“主语”或“时态”等明确概念。训练开始时，\(W_Q, W_K, W_V\) 通常采用随机初始化。随着模型在大量样本上不断进行前向计算、误差反向传播和参数更新：

• \(W_Q\) 逐渐学会生成有用的“查询需求”；
• \(W_K\) 逐渐学会生成便于匹配的“索引特征”；
• \(W_V\) 逐渐学会生成真正值得传递的“内容表示”。

因此，将 \(W_Q, W_K, W_V\) 理解为可训练的特征投影矩阵，而将 \(Q, K, V\) 理解为针对当前输入动态计算得到的中间表示，是更为准确的视角。

---

注意力分数（Attention Scores）

得到 \(Q\) 和 \(K\) 后，模型需要衡量每个 Query 与所有 Key 之间的匹配程度。对于第 \(i\) 个 Token 的 Query（\(q_i\)）和第 \(j\) 个 Token 的 Key（\(k_j\)），使用点积计算匹配分数：

\[q_i k_j^T \]

因为 \(q_i \in \mathbb{R}^{1 \times d_k}\)，且 \(k_j^T \in \mathbb{R}^{d_k \times 1}\)，所以点积结果是一个标量。将其展开可得：

\[q_i k_j^T = \sum_{r=1}^{d_k} Q_{ir} K_{jr} \]

该标量表示第 \(i\) 个 Token 的 Query 与第 \(j\) 个 Token 的 Key 在当前注意力空间中的匹配程度。如果一次性以矩阵形式计算所有 Token 两两之间的匹配分数，则有：

\[S = Q K^T \]

由于 \(Q \in \mathbb{R}^{n \times d_k}\) 且 \(K^T \in \mathbb{R}^{d_k \times n}\)，因此注意力分数矩阵 \(S \in \mathbb{R}^{n \times n}\)，可展开写为：

\[S = \begin{bmatrix} q_1 k_1^T & q_1 k_2^T & \cdots & q_1 k_n^T \\ q_2 k_1^T & q_2 k_2^T & \cdots & q_2 k_n^T \\ \vdots & \vdots & \ddots & \vdots \\ q_n k_1^T & q_n k_2^T & \cdots & q_n k_n^T \end{bmatrix} \]

其中：

\[\boxed{ S_{ij} = (QK^T)_{ij} = q_i k_j^T } \]

该元素表示第 \(i\) 个 Token 对第 \(j\) 个 Token 计算得到的原始注意力分数。

注意：

• \(S\) 的第 \(i\) 行对应第 \(i\) 个 Query；
• \(S\) 的第 \(j\) 列对应第 \(j\) 个 Key；
• \(S_{ij}\)表示第 \(i\) 个 Token 对第 \(j\) 个 Token 的匹配分数。

该分数衡量的是模型在当前 Query-Key 表示空间中学习到的匹配程度，不一定等同于人类直觉中的词义相似度。例如，代词与其指代对象、动词与其宾语可能在字面上并不相似，但模型仍可能为它们赋予较高的注意力匹配分数。

---

缩放点积（Scaled Dot-Product）

实际的 Self-Attention 并不直接使用 \(QK^T\) 作为分数，而是将其除以 \(\sqrt{d_k}\) 进行缩放：

\[\widetilde{S} = \frac{QK^T}{\sqrt{d_k}} \]

对于单个元素，即：

\[\widetilde{S}_{ij} = \frac{q_i k_j^T}{\sqrt{d_k}} \]

缩放的原因在于：
在常用的简化假设下，若 Query 和 Key 各维度的分量相互独立，且均值为 \(0\)、方差为 \(1\)，则点积 \(q_i k_j^T = \sum_{r=1}^{d_k} Q_{ir} K_{jr}\) 的方差为 \(d_k\)。这意味着，点积结果的数值尺度会随着特征维度 \(d_k\) 的增大而显著放大。当 \(d_k\) 较大时，绝对值偏大的注意力分数易使 Softmax 函数的输出过度集中，即极少数位置的权重趋近于 \(1\)，而其余位置趋近于 \(0\)。这将导致 Softmax 函数进入梯度饱和区，使 Softmax 对大部分位置的梯度变得很小。将点积结果除以 \(\sqrt{d_k}\) 后，在上述简化假设下，缩放后分数的方差约为 1。这使得注意力分数的数值尺度不再随 \(d_k\) 的增大而显著膨胀，从而缓解了 Softmax 输出分布过度集中的问题，并改善训练过程中的数值稳定性。正因引入了这一操作，该机制被称为缩放点积注意力（Scaled Dot-Product Attention）。

---

注意力权重（Attention Weights）

缩放后的注意力分数仍可能为任意实数（有正有负），且每一行的总和通常不等于 1。为了将注意力分数转化为可用于信息聚合的归一化权重，需要沿矩阵的每一行执行 Softmax：

\[A = \operatorname{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right) \]

展开权重矩阵 \(A \in \mathbb{R}^{n \times n}\)：

\[A = \begin{bmatrix} \alpha_{11} & \alpha_{12} & \cdots & \alpha_{1n} \\ \alpha_{21} & \alpha_{22} & \cdots & \alpha_{2n} \\ \vdots & \vdots & \ddots & \vdots \\ \alpha_{n1} & \alpha_{n2} & \cdots & \alpha_{nn} \end{bmatrix} \]

为了清晰说明注意力分数、Softmax 和 Value 聚合的计算过程，以下首先讨论不含掩码的 Self-Attention，即对于任意位置 \(i\)，都有 \(\mathcal A_i=\{1,\ldots,n\}\)。对于固定的第 \(i\) 个 Query，它对第 \(j\) 个 Key 的注意力权重计算公式为：

\[\boxed{ \alpha_{ij} = \frac{\exp\left(q_i k_j^T / \sqrt{d_k}\right)}{\displaystyle\sum_{t=1}^{n} \exp\left(q_i k_t^T / \sqrt{d_k}\right)} } \]

且：

\[\alpha_{ij} \geq 0 \quad \text{且} \quad \sum_{j = 1}^{n}\alpha_{ij} = 1 \]

这里的 \(\alpha_{ij}\) 具体表示：第 \(i\) 个 Token 在更新自身表示时，从第 \(j\) 个 Token 的 Value 中吸收信息的相对比重。 Softmax 沿 Key 所在的维度计算，即对于每一个固定的 Query，将其与所有 Key 的匹配分数归一化为注意力权重。

---

对 Value 进行加权聚合

以下首先讨论无掩码 Self-Attention。得到注意力权重矩阵 \(A\) 后，模型根据这些权重对 Value 进行加权求和：

\[Z = A V \]

由于 \(A \in \mathbb{R}^{n \times n}\)，\(V \in \mathbb{R}^{n \times d_v}\)，结果矩阵的维度为 \(Z \in \mathbb{R}^{n \times d_v}\)。其中，\(Z\) 的第 \(i\) 个 Token 的输出行向量为：

\[\boxed{ z_i = \sum_{j = 1}^{n}\alpha_{ij} v_j } \]

将 \(Z\) 按行展开：

\[\boxed{ Z = \begin{bmatrix} z_1 \\ z_2 \\ \vdots \\ z_n \end{bmatrix} = \begin{bmatrix} \displaystyle\sum_{j=1}^{n}\alpha_{1j}v_j \\ \displaystyle\sum_{j=1}^{n}\alpha_{2j}v_j \\ \vdots \\ \displaystyle\sum_{j=1}^{n}\alpha_{nj}v_j \end{bmatrix} = \begin{bmatrix} \alpha_{11}&\alpha_{12}&\cdots&\alpha_{1n} \\ \alpha_{21}&\alpha_{22}&\cdots&\alpha_{2n} \\ \vdots&\vdots&\ddots&\vdots \\ \alpha_{n1}&\alpha_{n2}&\cdots&\alpha_{nn} \end{bmatrix} \begin{bmatrix} v_1 \\ v_2 \\ \vdots \\ v_n \end{bmatrix} = AV } \]

这意味着，第 \(i\) 个 Token 使用自身对应的一组注意力权重，对所有的 Value 向量执行加权求和，从而得到了融合上下文信息的新表示 \(z_i\)。

如果精细到矩阵内部的单个元素层面：

\[\boxed{ Z_{ir} = \sum_{j = 1}^{n}\alpha_{ij} V_{jr} } \]

其中：

• \(i\) 表示当前输出 Token 的位置；
• \(r\) 表示 Value 向量或输出向量的特征维度；
• \(j\) 表示被聚合的上下文 Token 位置。

综上所述，单头 Self-Attention 的完整矩阵表达式为：

\[\boxed{ Z = \operatorname{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V } \]

将线性投影展开代入后可得：

\[\boxed{ Z = \operatorname{softmax}\left(\frac{(XW_Q)(XW_K)^T}{\sqrt{d_k}}\right)XW_V } \]

---

因果掩码（Causal Mask）

在 GPT 等自回归语言模型（Autoregressive Models）中，模型需要根据前面的 Token 预测后续 Token，因此第 \(i\) 个 Token 在计算时，不能“看到”位置 \(i\) 之后的 Token（即未来信息）。

例如，对于长度为 4 的序列，允许访问的位置关系（行表示 Query，列表示 Key）应当是下三角矩阵的形式：

\[\begin{bmatrix} \checkmark & \times & \times & \times \\ \checkmark & \checkmark & \times & \times \\ \checkmark & \checkmark & \checkmark & \times \\ \checkmark & \checkmark & \checkmark & \checkmark \end{bmatrix} \]

为了在并行矩阵运算中实现这一限制，需要引入掩码矩阵 \(M\)：

\[M_{ij} = \begin{cases} 0, & j \leq i \\ -\infty, & j > i \end{cases} \]

其中

\[M \in \mathbb{R}^{n \times n} \]

然后在进行 Softmax 之前，将掩码矩阵加到注意力分数上：

\[A = \operatorname{softmax}\left(\frac{QK^T}{\sqrt{d_k}} + M\right) \]

由于 \(e^{-\infty} = 0\)，所有对应“未来位置”的注意力权重在 Softmax 之后都会被强制置为 0。此时，第 \(i\) 个 Token 的输出聚合范围便被严格限制：

\[z_i = \sum_{j=1}^{i} \alpha_{ij} v_j \]

从而确保了因果性：第 \(i\) 个 Token 只能聚合自身以及之前 Token 的信息。

---

多头注意力（Multi-Head Attention）

单头 Self-Attention（自注意力机制）只使用一组投影矩阵 \(W_Q, W_K, W_V\)，因此所有 Token 都在同一个 Query-Key 表示空间中计算注意力权重，并按照同一种 Value 表示方式聚合上下文信息。

然而，自然语言中 Token 之间往往同时存在多种不同的关系。例如，在句子：

“小明把昨天买的书送给了正在复习的小红。”

对于“送”这个 Token，模型可能需要同时识别：

• 谁执行了“送”这一动作，即“小明”；
• 被送的对象是什么，即“书”；
• 物品被送给了谁，即“小红”；
• 句子中与该动作相关的时间、修饰和位置关系。

单个注意力头虽然可以同时关注多个位置，但它只能在一个 Query-Key 表示空间中生成一组注意力权重，不同类型的关系需要在同一注意力分布和同一组 Value 表示中共同建模。为了使模型能够从多个不同的表示子空间中并行提取信息，Transformer 引入了多头注意力机制（Multi-Head Attention）。

假设注意力头的数量为 \(H\)。第 \(h\) 个注意力头拥有一组独立的投影矩阵：

\[W_Q^{(h)}, \qquad W_K^{(h)}, \qquad W_V^{(h)} \]

在常见的维度设置下：

\[W_Q^{(h)}, W_K^{(h)}, W_V^{(h)} \in \mathbb{R}^{d_{\text{model}} \times d_h} \]

其中，\(d_h\) 表示单个注意力头的特征维度，在标准 Transformer 中，通常令每个注意力头的维度为：

\[d_h=\frac{d_{\text{model}}}{H} \]

因此 \(Hd_h=d_{\text{model}}\)。第 \(h\) 个注意力头首先对输入矩阵 \(X\) 进行三种独立的线性投影：

\[Q^{(h)} = XW_Q^{(h)}, \qquad K^{(h)} = XW_K^{(h)}, \qquad V^{(h)} = XW_V^{(h)} \]

因此：

\[Q^{(h)}, K^{(h)}, V^{(h)} \in \mathbb{R}^{n \times d_h} \]

在无掩码情况下第 \(h\) 个注意力头独立执行一次缩放点积注意力：

\[\boxed{ \operatorname{head}_h = \operatorname{softmax}\left(\frac{Q^{(h)}\left(K^{(h)}\right)^T}{\sqrt{d_h}}\right)V^{(h)} } \]

如果采用因果 Self-Attention，则需要在每个注意力头的 Softmax 之前加入相同的因果掩码 \(M\)：

\[\boxed{ \operatorname{head}_h = \operatorname{softmax}\left(\frac{Q^{(h)}\left(K^{(h)}\right)^T}{\sqrt{d_h}} + M \right)V^{(h)} } \]

其中：

\[\operatorname{head}_h = \begin{bmatrix} z_1^{(h)} \\ z_2^{(h)} \\ \vdots \\ z_n^{(h)} \end{bmatrix} \in\mathbb{R}^{n\times d_h}, z_i^{(h)}\in\mathbb{R}^{1\times d_h} \]

由于不同注意力头使用相互独立的投影参数，它们可以将相同的输入映射到不同的表示子空间，并学习不同或互补的注意力模式。例如，某些注意力头可能更关注局部相邻关系，某些注意力头可能更擅长捕捉长距离依赖，另一些注意力头则可能对主谓关系、动宾关系或指代关系更加敏感。

需要注意的是，这种功能分工并不是人工预先指定的，而是模型在训练过程中根据任务目标自动形成的。不同注意力头之间也可能存在一定程度的功能重叠，并不一定能够为每个头赋予明确且唯一的语义解释。

当所有 \(H\) 个注意力头计算完成后，将它们的输出沿特征维度进行拼接（Concatenation）：

\[C = \operatorname{Concat}\left(\operatorname{head}_1, \operatorname{head}_2, \ldots, \operatorname{head}_H\right) \]

假设每个注意力头的输出维度均为 \(d_h\)，则拼接后的矩阵维度为：

\[C \in \mathbb{R}^{n \times Hd_h} \]

从单个 Token 的视角来看，对于第 \(i\) 个 Token，其拼接后的上下文表示为：

\[c_i = \begin{bmatrix} z_i^{(1)} & z_i^{(2)} & \cdots & z_i^{(H)} \end{bmatrix} \]

拼接操作只是将各个注意力头的输出排列在不同的特征区间中，尚未对不同头的信息进行线性混合。为了使模型能够学习如何融合不同注意力头提取的特征，多头注意力进一步引入了一个可训练的输出投影矩阵：

\[W_O \in \mathbb{R}^{Hd_h \times d_{\text{model}}} \]

并将拼接结果投影到模型的输出表示空间：

\[O = CW_O \]

即：

\[\boxed{ O = \operatorname{Concat}\left(\operatorname{head}_1, \ldots, \operatorname{head}_H\right)W_O } \]

\(W_O\) 与 \(W_Q, W_K, W_V\) 一样，都是模型在训练过程中通过反向传播自动学习得到的参数矩阵。它一方面对不同注意力头提取的特征进行线性混合，另一方面将拼接后的表示映射到模型所需的 \(d_{\text{model}}\) 维输出空间。

因此，多头注意力的完整表达式为：

\[\boxed{ \operatorname{MultiHead}(X) = \operatorname{Concat}\left(\operatorname{head}_1, \operatorname{head}_2, \ldots, \operatorname{head}_H\right)W_O } \]

在无掩码情况下，单个注意力头的计算为：

\[\operatorname{head}_h = \operatorname{softmax}\left(\frac{(XW_Q^{(h)})(XW_K^{(h)})^T}{\sqrt{d_h}}\right)XW_V^{(h)} \]

在标准维度设置下，多头注意力的最终输出 \(O\) 满足：

\[O \in \mathbb{R}^{n \times d_{\text{model}}} \]

这与输入矩阵 \(X \in \mathbb{R}^{n \times d_{\text{model}}}\) 具有完全相同的形状，从而能够顺畅地进一步执行残差连接（Residual Connection）：

\[Y = X + O \]

核心作用总结：多头注意力通过使用多组独立的投影参数，使模型能够在多个不同的表示子空间中，并行学习 Token 之间不同或互补的关联模式。各个注意力头提取的上下文信息在经过输出投影后被重新组合，最终形成更加丰富的 Token 上下文表示。