2017年，Google Research 团队发表论文《Attention Is All You Need》，提出了一种名为 Transformer 的神经网络架构。论文标题是一个明确的判断——在此之前的神经机器翻译模型依赖于循环神经网络（RNN）和卷积神经网络（CNN）的组合，而 Transformer 仅使用注意力机制就达到了当时的最优性能。

八年后的今天，Transformer 已经成为现代语言模型的核心架构。从 BERT 到 GPT-4，从 LLaMA 到 Claude，理解注意力机制是理解这些模型工作原理的基础。

从信息检索理解注意力的本质

考虑一个信息检索的场景：在图书馆中查找一本关于"机器学习"的书。检索过程包含三个步骤：首先，形成查询需求（Query）——“机器学习”；其次，扫描书架上每本书的标签或书名（Key）；最后，根据匹配程度获取目标书籍的内容（Value）。

注意力机制遵循相似的逻辑：用查询向量匹配键向量，根据匹配程度加权求和值向量。在 Transformer 的自注意力中，序列中的每个词同时扮演这三个角色——它既是查询者，也是被查询者，还是被提取的信息本身。

序列建模的信息瓶颈问题

在 Transformer 出现之前，序列到序列模型面临一个根本性挑战。编码器需要将整个输入序列压缩为一个固定长度的向量，解码器只能从这个向量中提取信息。对于长序列，这种压缩导致早期信息丢失。

Bahdanau 等人在 2014 年提出的注意力机制部分解决了这个问题。解码器在生成每个输出时，可以动态关注编码器的不同位置。Transformer 进一步发展了这一思想：序列中的每个位置都能直接关注其他所有位置，不再依赖顺序处理。这种全局连接是 Transformer 并行化处理的关键。

缩放点积注意力的计算过程

Transformer 使用缩放点积注意力（Scaled Dot-Product Attention），其数学形式为：

$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$

其中：

  • $Q \in \mathbb{R}^{n \times d_k}$：查询矩阵，$n$ 是序列长度，$d_k$ 是查询向量维度
  • $K \in \mathbb{R}^{n \times d_k}$：键矩阵
  • $V \in \mathbb{R}^{n \times d_v}$：值矩阵

这个公式包含四个步骤：计算点积、缩放、Softmax 归一化、加权求和。下面逐一分析每个步骤的数学原理。

第一步：点积计算注意力分数

矩阵乘法 $QK^T$ 计算查询向量与键向量的点积。对于查询向量 $q_i$ 和键向量 $k_j$，点积定义为：

$$\text{score}_{ij} = q_i \cdot k_j^T = \sum_{m=1}^{d_k} q_{im} \cdot k_{jm}$$

点积衡量两个向量的相似程度。当 $q_i$ 和 $k_j$ 方向相似时，点积为较大的正值；当方向相反时，点积为负值。

假设序列长度为 4（词 a, b, c, D），$d_k = 2$。计算 $QK^T$ 得到一个 $4 \times 4$ 的分数矩阵：

$$QK^T = \begin{bmatrix} q_a \cdot k_a & q_a \cdot k_b & q_a \cdot k_c & q_a \cdot k_D \\ q_b \cdot k_a & q_b \cdot k_b & q_b \cdot k_c & q_b \cdot k_D \\ q_c \cdot k_a & q_c \cdot k_b & q_c \cdot k_c & q_c \cdot k_D \\ q_D \cdot k_a & q_D \cdot k_b & q_D \cdot k_c & q_D \cdot k_D \end{bmatrix}$$

每个元素 $(i, j)$ 表示第 $i$ 个词对第 $j$ 个词的"关注程度"。

第二步：缩放因子的必要性

除以 $\sqrt{d_k}$ 是缩放点积注意力的关键设计。这个缩放因子的作用是防止 Softmax 函数进入饱和区。

论文《Attention Is All You Need》对此的解释是：

“We suspect that for large values of $d_k$, the dot products grow large in magnitude, pushing the softmax function into regions where it has extremely small gradients.”

问题的根源在于 Softmax 函数的特性。Softmax 将向量转换为概率分布：

$$\text{softmax}(x_i) = \frac{e^{x_i}}{\sum_j e^{x_j}}$$

当输入向量的元素差异较大时，Softmax 输出会趋向于 one-hot 向量——最大元素接近 1，其他元素接近 0。例如，输入 [1, 1, 1, 5] 经过 Softmax 后得到 [0.017, 0.017, 0.017, 0.948]。

为什么点积结果可能很大？假设 $q$ 和 $k$ 的元素是均值为 0、方差为 1 的独立随机变量。点积的方差为：

$$\text{Var}(q \cdot k^T) = \text{Var}\left(\sum_{i=1}^{d_k} q_i k_i\right) = d_k$$

因为 $\text{Var}(q_i k_i) = E[q_i^2]E[k_i^2] = 1 \times 1 = 1$（假设 $q_i$ 和 $k_i$ 独立）。

当 $d_k = 64$ 时，点积的标准差为 8。在正态分布下，相当比例的点积结果会落在 $[-16, 16]$ 范围之外。这些大数值会导致 Softmax 输出接近 one-hot 向量，其梯度接近零，阻碍训练。

缩放因子 $\frac{1}{\sqrt{d_k}}$ 将点积的方差归一化为 1，使 Softmax 的输入保持在梯度有效的范围内。

图片来源: Scaled Dot-Product Attention and Masking in Transformers - CodeSignal

第三步：Softmax 归一化

对缩放后的分数矩阵每一行应用 Softmax 函数：

$$\text{Attention Weights} = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)$$

每一行的 Softmax 输出是一个概率分布，所有元素之和为 1。第 $i$ 行表示第 $i$ 个词对序列中所有词的注意力权重分布。

第四步：加权求和

最后，用注意力权重对值矩阵加权：

$$\text{Output} = \text{Attention Weights} \times V$$

对于第 $i$ 个词的输出：

$$\text{output}_i = \sum_{j=1}^{n} \alpha_{ij} v_j$$

其中 $\alpha_{ij}$ 是注意力权重。每个输出向量都是所有值向量的加权和，权重由查询-键相似度决定。

Q、K、V 向量的来源

在自注意力中，$Q$、$K$、$V$ 通过对同一输入序列 $X$ 进行线性变换得到。假设输入序列的嵌入矩阵为 $X \in \mathbb{R}^{n \times d_{model}}$：

$$Q = XW^Q, \quad K = XW^K, \quad V = XW^V$$

其中 $W^Q, W^K \in \mathbb{R}^{d_{model} \times d_k}$，$W^V \in \mathbb{R}^{d_{model} \times d_v}$ 是可学习参数。

为什么要投影而非直接使用原始嵌入？

同一个词在不同角色中需要表达不同的信息：作为查询时表达"我在寻找什么"，作为键时表达"我能提供什么匹配信息"，作为值时表达"我的内容是什么"。投影矩阵学习这些不同的表示方式。

多头注意力的设计原理

单个注意力头学习一种注意力模式。但语言中的依赖关系是多样的——主谓关系、修饰关系、指代关系等。多头注意力通过并行运行多个头捕获这些不同的关系：

$$\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, ..., \text{head}_h)W^O$$

其中：

$$\text{head}_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)$$

每个头有独立的投影矩阵 $W_i^Q, W_i^K, W_i^V$，可以学习不同的表示子空间。

图片来源: Multi-Head Attention Mechanism - GeeksforGeeks

多头是否真的必要？

2021 年的一篇论文《Multi-head or Single-head? An Empirical Comparison for Transformer》探讨了这个问题。研究发现，多头注意力在训练稳定性上优于等价的深单头注意力：

  1. 并行学习多种模式：每个头可以独立学习不同的语义关系，无需等待前一层学习完成
  2. 减少深度，提高稳定性：达到相同表达能力，多头所需的层数更少
  3. 缓解梯度消失：每个头有独立的梯度路径

原始 Transformer 使用 8 个注意力头。在现代大模型中，头数通常更多——GPT-3 有 96 个头，LLaMA-2-70B 有 64 个头。

多头注意力的实现细节

假设 $d_{model} = 512$，$h = 8$，则每个头的维度：

$$d_k = d_v = d_{model} / h = 64$$

实际实现中，并非真正创建 $h$ 个独立的注意力模块，而是将大的投影矩阵视为 $h$ 个小矩阵的拼接：

# 输入: x [batch, seq_len, d_model]
# 一次投影得到所有头的 Q, K, V
qkv = linear(x)  # [batch, seq_len, 3 * d_model]
q, k, v = qkv.chunk(3, dim=-1)

# 重塑为多头形式
q = q.view(batch, seq_len, num_heads, head_dim).transpose(1, 2)
# 形状: [batch, num_heads, seq_len, head_dim]

# 并行计算所有头的注意力
attn = softmax(q @ k.transpose(-2, -1) / sqrt(head_dim))
output = attn @ v  # [batch, num_heads, seq_len, head_dim]

# 合并多头
output = output.transpose(1, 2).reshape(batch, seq_len, d_model)

这种"逻辑分头"的实现方式让计算可以通过矩阵乘法高度并行化，效率远高于 $h$ 次独立计算。

自注意力与交叉注意力的区别

Transformer 中有两种注意力形式：

自注意力（Self-Attention）

$Q$、$K$、$V$ 来自同一序列。用于捕获序列内部的依赖关系。

在编码器中，自注意力用于理解输入序列；在解码器中，自注意力用于理解已生成的输出序列。

交叉注意力（Cross-Attention）

$Q$ 来自一个序列，$K$ 和 $V$ 来自另一个序列。用于两个序列之间的信息传递。

在翻译任务中：

  • 解码器的 $Q$：来自目标语言的已生成部分
  • 编码器的 $K$、$V$：来自源语言的完整表示

交叉注意力是编码器-解码器架构中信息传递的桥梁。

图片来源: Self-Attention vs Cross-Attention: From Fundamentals to Applications - Medium

注意力掩码的作用

掩码用于控制注意力范围。通过将特定位置的注意力分数设为负无穷，Softmax 后这些位置的权重变为 0。

填充掩码（Padding Mask）

在批处理中，不同序列长度不同，需要用 <PAD> 标记填充到相同长度。填充掩码防止模型关注填充标记：

$$M_{padding} = \begin{cases} 0 & \text{if position is real token} \\ -\infty & \text{if position is padding} \end{cases}$$

掩码添加在 Softmax 之前：

$$\text{Attention Weights} = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}} + M\right)$$

因果掩码（Causal Mask）

在自回归生成中，模型不能"看到未来"。因果掩码是一个下三角矩阵：

$$M_{causal} = \begin{bmatrix} 0 & -\infty & -\infty & -\infty \\ 0 & 0 & -\infty & -\infty \\ 0 & 0 & 0 & -\infty \\ 0 & 0 & 0 & 0 \end{bmatrix}$$

这确保位置 $i$ 只能关注位置 $1$ 到 $i$，不能关注 $i+1$ 及之后的位置。

为什么因果掩码只在"列"上应用（键的方向），而不是同时在"行"上（查询的方向）？

如果同时掩码行和列，被掩码位置的输出将是一个所有值向量均匀加权和，这会丢失该位置本身的信息。只掩码列，可以让被掩码位置仍然保留自己的表示，只是不影响其他位置。

图片来源: How Do Self-Attention Masks Work? - Gabriel Mongaras

计算复杂度分析

自注意力的计算复杂度是 $O(n^2 d)$，其中 $n$ 是序列长度，$d$ 是隐藏维度。瓶颈在于 $QK^T$ 的矩阵乘法，产生一个 $n \times n$ 的注意力矩阵。

模型类型 每层复杂度 序列操作数 最大路径长度
Self-Attention $O(n^2 \cdot d)$ $O(1)$ $O(1)$
Recurrent $O(n \cdot d^2)$ $O(n)$ $O(n)$
Convolutional $O(k \cdot n \cdot d^2)$ $O(1)$ $O(\log_k n)$

当序列长度 $n$ 小于隐藏维度 $d$ 时，自注意力比 RNN 更高效；但当处理超长序列时，$O(n^2)$ 的复杂度成为瓶颈。这也是原始 Transformer 的上下文长度受限（通常 512-2048 tokens），以及后续出现 FlashAttention、Ring Attention 等优化方案的原因。

注意力权重的局限性

注意力权重常被解释为模型的"关注焦点"，但这种解释需要审慎。一篇题为《Attention is not Explanation》的论文指出：

  1. 多个不同的注意力分布可能产生相同的预测结果
  2. 人为改变注意力权重，模型的预测可能不变
  3. 注意力权重只是模型的内部状态，不必然反映决策依据

注意力权重确实提供了模型行为的线索，但不是完整的解释。理解 Transformer 的"思考过程"需要更深入的分析。

参考文献

  1. Vaswani, A., et al. (2017). Attention Is All You Need. NeurIPS 2017.
  2. Bahdanau, D., Cho, K., & Bengio, Y. (2014). Neural Machine Translation by Jointly Learning to Align and Translate. arXiv:1409.0473.
  3. Li, Z., et al. (2021). Multi-head or Single-head? An Empirical Comparison for Transformer. arXiv:2106.09650.
  4. Jain, S., & Wallace, B. C. (2019). Attention is not Explanation. NAACL 2019.
  5. Dao, T., et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS 2022.
  6. Wiegreffe, S., & Pinter, Y. (2019). Attention is not not Explanation. EMNLP 2019.