告别玄学！从数学本质彻底搞懂 Transformer 的注意力机制

Posted on 2026-05-23 In AI 技术

提到 Transformer，很多人脑海中第一时间浮现的往往是那句名言：“Attention is all you need”。作为大语言模型（LLM）、Vision Transformer（ViT）等当今最前沿 AI 技术的基石，注意力机制已经成为每个算法工程师乃至 AI 爱好者必须跨越的门槛。

然而，面对层出不穷的教程，很多人依然停留在“玄学”的理解层面：知道它是用来计算权重的，知道它涉及 Q、K、V 三个矩阵，但若是深究一句：“为什么是 Q、K、V？为什么 Attention 公式里要除以 $\sqrt{d_k}$ ？点积为什么能代表相似度？” 很多朋友可能就支支吾吾了。

今天，我们将抛开干瘪的代码调用，从底层的线性代数与信息检索视角出发，用最硬核的数学语言，带你彻底剥开 Transformer 注意力机制的外衣。

一、从信息检索说起：Q、K、V 到底是什么？

理解注意力机制最快的方法，是类比我们在数据库中的查询过程。

想象一下你在视频网站上搜索一部电影：

你在搜索框输入的关键词，就是查询。
网站上每个视频自带的标题、标签、简介，就是键。
视频实际播放的内容，就是值。

在传统数据库中， $Q$ 和 $K$ 往往是精确匹配的（比如你搜“盗梦空间”，只有标题完全包含这个词的才会返回）。但在自然语言处理中，我们面临的是模糊匹配的需求。比如你输入“动作片”，系统不仅应该返回标签为“动作片”的视频，还可能返回标签为“科幻”、“漫威”的视频，且赋予不同的权重。

注意力机制的本质，就是一个基于向量空间模糊匹配的动态加权检索系统。

在 Transformer 中，输入是一个矩阵 $X$ （包含序列中每个 Token 的 Embedding）。我们需要用三个线性变换矩阵 $W^Q, W^K, W^V$ 将 $X$ 投影到三个不同的子空间中，得到 $Q, K, V$ ：

Q = X W^Q

K = X W^K

V = X W^V

数学直觉： 为什么需要三个矩阵而不是直接用 $X$ ？
线性变换 $W$ 的作用是空间旋转与拉伸。通过将原始向量投影到不同的表示子空间，模型能够在一个更适合计算相似度的空间（ $Q$ 和 $K$ 的空间）中衡量关系，并在另一个包含丰富特征的空间（ $V$ 的空间）中提取信息。

二、核心公式拆解：点积与 Softmax

得到 $Q, K, V$ 之后，Transformer 论文中给出的缩放点积注意力公式如下：

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

这个公式看似简单，实则字字珠玑。我们将其拆解为三个步骤。

2.1 为什么是 $Q K^T$ ？（点积的几何意义）

首先计算 $Q$ 和 $K$ 的点积。假设 $Q$ 的维度是 $(L_q, d_k)$ ， $K$ 的维度是 $(L_k, d_k)$ （在自注意力中 $L_q = L_k = L$ ）。它们的乘积 $Q K^T$ 是一个 $(L, L)$ 的方阵。

这个方阵里的每一个元素 $S_{i,j}$ ，代表第 $i$ 个 Token 的 Query 向量与第 $j$ 个 Token 的 Key 向量的点积。

为什么点积能衡量相似度？从几何上看，两个向量 $\vec{a}$ 和 $\vec{b}$ 的点积公式为：

\vec{a} \cdot \vec{b} = |\vec{a}| |\vec{b}| \cos(\theta)

当向量长度固定时，点积的大小完全取决于两个向量之间的夹角 $\theta$ 。夹角越小（方向越一致），点积越大。因此，点积本质上是在衡量两个向量在超空间中的方向一致性，即语义相关性。

2.2 致命的 $\sqrt{d_k}$ ：梯度消失的终结者

公式中非常刺眼的是除以 $\sqrt{d_k}$ 。很多人只知道这是“缩放”，却不知道为何要缩放。

假设 $Q$ 和 $K$ 的各个维度是均值为 0，方差为 1 的独立同分布随机变量。那么它们点积的数学期望为：

E(q \cdot k) = \sum_{i=1}^{d_k} E(q_i k_i) = 0

但方差却会随着维度累加：

Var(q \cdot k) = \sum_{i=1}^{d_k} Var(q_i k_i) = d_k \times 1 = d_k

这意味着，当模型的维度 $d_k$ 很大时（例如 Transformer 中常设为 64 或 512），点积的结果在数值上会变得非常大，方差极大。

这会引发什么灾难？接下来的一步是 $\text{softmax}$ 。 $\text{softmax}$ 函数在输入值差异过大时，会表现出强烈的“赢者通吃”特性——最大值被映射为 1，其他所有值被压榨成无限接近 0。

一旦输出无限接近 0，在反向传播时，梯度的乘积就会迅速趋近于 0，导致梯度消失，模型几乎无法学习。

除以 $\sqrt{d_k}$ 的本质，就是对点积结果进行标准化（Standardization），强行将其方差拉回 1。 这样无论维度多高，进入 softmax 的数值都保持在合理的范围，保证了梯度的稳定流动。

2.3 矩阵乘以 $V$ ：动态信息聚合

经过 $\text{softmax}$ 之后，我们得到了一个 $(L, L)$ 的注意力权重矩阵 $A$ 。矩阵中的每一行代表当前 Token 对所有其他 Token 的关注度分布（且和为 1）。

最后一步，乘以 $V$ ：

\text{Output} = A V

这一步的几何意义是加权求和。输出的向量序列中，每一个 Token 的新表示，都是其他所有 Token 的 Value 向量的线性组合。权重正是刚刚算出的相关性。

核心感悟： 自注意力机制实际上是一种信息融合。它打破了 RNN 只能依赖历史信息的局限，让序列中的每一个词都能“一眼看穿”全局，根据语义的相关性，自主吸收有用的信息，重新塑造自身的表征。

三、跃升至多维：多头注意力的几何解释

如果单头注意力已经足够强大，为什么 Transformer 还需要多头注意力？

让我们回到数学公式。单头注意力实际上是将 $Q, K, V$ 限制在了一个固定的 $d_k$ 维子空间中计算。这就像是你只用一种“标准”（比如语法关系）去衡量词语的关联。

但语言是复杂的，“苹果”和“吃”是动宾关系，“苹果”和“红色”是属性关系。我们希望模型能同时从不同的子空间中捕捉不同的关系模式。

多头注意力就是引入 $h$ 个头，每个头拥有独立的线性变换矩阵 $W_i^Q, W_i^K, W_i^V$ ，将输入分别映射到不同的 $d_k = d_{model}/h$ 维子空间中，计算 $h$ 次注意力，然后将结果拼接：

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

从线性代数的角度看，多头注意力本质上是对高维特征空间进行了一次降维打击和正交分解。
它强制模型在不同的低维子空间中寻找独立的相关性特征，不仅极大地增强了模型的表达能力，最后拼接并乘以 $W^O$ 矩阵的操作，也巧妙地利用了块对角矩阵的特性，降低了整体的计算复杂度。

四、纯手工实现：用 PyTorch 写一个极简 Self-Attention

理论讲完了，动手时间到。如果你只调包 torch.nn.MultiheadAttention，你可能永远无法体会矩阵计算的优雅。下面我们从零实现一个单头自注意力层。

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class SelfAttention(nn.Module):
    def __init__(self, embed_dim):
        super(SelfAttention, self).__init__()
        self.embed_dim = embed_dim
        
        # 初始化 Q, K, V 的线性映射层
        # 我们不设置 bias，这是 Transformer 的常见做法
        self.W_q = nn.Linear(embed_dim, embed_dim, bias=False)
        self.W_k = nn.Linear(embed_dim, embed_dim, bias=False)
        self.W_v = nn.Linear(embed_dim, embed_dim, bias=False)
        
    def forward(self, x):
        # x 的 shape 为 (batch_size, seq_len, embed_dim)
        batch_size, seq_len, _ = x.size()
        
        # 1. 线性变换得到 Q, K, V
        Q = self.W_q(x)  # (batch_size, seq_len, embed_dim)
        K = self.W_k(x)  # (batch_size, seq_len, embed_dim)
        V = self.W_v(x)  # (batch_size, seq_len, embed_dim)
        
        # 2. 计算 Q 和 K 的点积，即未缩放的注意力分数
        # torch.bmm 是批量矩阵乘法，适用于 3D tensors
        # Q 的形状是 (..., seq_len, embed_dim)
        # K^T 的形状是 (..., embed_dim, seq_len)
        # scores 的形状为 (batch_size, seq_len, seq_len)
        scores = torch.bmm(Q, K.transpose(1, 2))
        
        # 3. 缩放：除以根号下 d_k，防止梯度消失
        scores = scores / math.sqrt(self.embed_dim)
        
        # 4. Softmax 归一化，得到权重概率分布 (每一行和为1)
        # dim=-1 表示在最后一个维度（即 K 的序列长度维度）上进行 softmax
        attn_weights = F.softmax(scores, dim=-1)
        
        # 5. 将注意力权重与 V 进行矩阵相乘，进行信息聚合
        # attn_weights: (batch_size, seq_len, seq_len)
        # V: (batch_size, seq_len, embed_dim)
        # out: (batch_size, seq_len, embed_dim)
        out = torch.bmm(attn_weights, V)
        
        return out

# 测试一下
embed_dim = 64
seq_len = 10
batch_size = 4

# 模拟输入数据
dummy_input = torch.randn(batch_size, seq_len, embed_dim)
self_attn = SelfAttention(embed_dim)
output = self_attn(dummy_input)

print(f"Input shape: {dummy_input.shape}")
print(f"Output shape: {output.shape}")

看，短短十几行代码，就浓缩了我们在第二部分讲解的全部数学精华。transpose 实现转置，bmm 实现矩阵乘法，最后一步 out = bmm(...) 完美诠释了向量的动态加权求和。

五、隐藏的致命缺陷：排列不变性与位置编码

掌握了上述原理后，如果你足够敏锐，会发现一个恐怖的事实：在 $QK^T$ 的矩阵乘法中，模型根本不知道单词的顺序！

假设输入是“狗咬人”和“人咬狗”，对于纯 Self-Attention 机制来说，由于 $QK^T$ 计算的全是无序的两两相似度，只要这两个句子的词完全一样，它们计算出的表征在理论上是可以完全一样的。

数学上，这叫排列不变性。这对于要求严格顺序的自然语言来说是不可接受的。

因此，Transformer 的作者在输入矩阵 $X$ 送入注意力机制之前，强行注入了位置编码：

X_{final} = X_{embedding} + P_{positional}

技术前沿：
最初的 Transformer 使用的是由正弦和余弦函数构成的固定绝对位置编码。
如今的顶级大模型（如 LLaMA、ChatGLM）则普遍采用了旋转位置编码。RoPE 的数学高明之处在于，它没有在 Token 的 Value 上加位置信息，而是通过在复数域对 $Q$ 和 $K$ 进行旋转，巧妙地使得两个向量点积的结果中，自然蕴含了它们之间的绝对位置差（ $m - n$ ）。这完全契合了注意力机制的底层数学逻辑，将相对位置信息天衣无缝地融入到了 $QK^T$ 的计算中。

六、总结

回顾本文，我们并没有停留在对公式的死记硬背上，而是深入到了 Transformer 的肌理之中。我们来做一个最后的总结：

查询与匹配： 注意力机制本质是一个基于向量的信息检索系统，通过 $Q$ 和 $K$ 的点积来量化高维空间中的语义相关性。
方差与缩放： 除以 $\sqrt{d_k}$ 绝不是拍脑袋决定的，而是基于严格的统计学推导，其目的是约束方差，防止 Softmax 函数进入梯度消失的“平缓区”。
空间与特征： 多头注意力是对高维特征空间的一种切分与正交化，使模型具备了“多视角”观察句子的能力。
加权与聚合： 最终乘以 $V$ 矩阵，完成了基于全局权重的动态特征重组，这是 Transformer 能够取代 RNN 的最核心武器。

当我们把神经网络中的玄学剥离，剩下的就是极致的数学之美。理解这些底层数学本质，不仅能让我们在调参排错时更加得心应手，更能为我们未来改进模型架构（如设计更高效的线性注意力机制、探索新的位置编码）提供坚实的理论基础。

希望这篇文章，能成为你彻底搞懂 Transformer 之路上的一个坚实台阶。

一、 从信息检索说起：Q、K、V 到底是什么？

二、 核心公式拆解：点积与 Softmax

2.1 为什么是 QKTQ K^TQKT？（点积的几何意义）

2.2 致命的 dk\sqrt{d_k}dk​​：梯度消失的终结者

2.3 矩阵乘以 VVV：动态信息聚合

三、 跃升至多维：多头注意力的几何解释

四、 纯手工实现：用 PyTorch 写一个极简 Self-Attention

五、 隐藏的致命缺陷：排列不变性与位置编码

六、 总结