拆解 Transformer：从几何与线性代数视角看透注意力机制的数学本质

Posted on 2026-05-26 In AI 技术

自 2017 年 Google 发布那篇著名的《Attention is All You Need》论文以来，Transformer 架构便以摧枯拉朽之势席卷了整个深度学习领域。从自然语言处理的 BERT、GPT 系列，到计算机视觉的 ViT，再到如今横扫千军的大语言模型（LLM），Transformer 已经成为了现代人工智能的基石。

然而，许多开发者在初学 Transformer 时，往往会被公式中五花八门的矩阵乘法、缩放因子、多头机制弄得晕头转向。我们习惯了将其当成一个黑盒：输入一段文本，输出一堆向量。但如果不能深刻理解其背后的数学本质，我们在进行模型微调、架构改进或排查长文本丢失问题时，就会犹如盲人摸象。

今天，我们将暂时抛开繁杂的工程实现细节，剥开代码的外衣，从线性代数与高维几何的视角，深度剖析 Transformer 注意力机制的数学本质。

一、宏观视角：Q、K、V 到底在做什么？

要理解注意力机制，首先要回答一个根本问题：什么是“注意力”？

在人类认知中，注意力意味着在处理当前信息时，能够根据上下文动态地将焦点集中在某些关键部分。在 Transformer 中，这个过程被巧妙地转化为了一个信息检索与融合的数学过程。

注意力机制的核心公式如下：

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

公式中的 $Q$ (Query)、 $K$ (Key)、 $V$ (Value) 借鉴了数据库检索的概念：

Query (查询)：当前正在处理的词（或 Token），它想要了解自己，并在句中寻找上下文。
Key (索引/键)：上下文中其他词（也包括自己）的“身份标签”，用于等待 Query 的匹配。
Value (值)：上下文中其他词的实际“内容/语义信息”。

如果用通俗的语言来解释：每一个词都在通过 $Q$ 发出提问（“我是谁，谁跟我有关？”），然后拿这个问题去和所有词的 $K$ （标签）进行比对。比对得分越高的词，其 $V$ （实际内容）就会被分配更大的权重，最终融合进当前词的新表示中。

但这依然太宏观了。接下来，我们进入微观的数学世界。

二、微观拆解：一步步看透矩阵的几何变换

让我们把公式 $\text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$ 拆解成四个独立的数学步骤。

1. 线性投影：升维与子空间切割

在进入公式之前，输入序列首先会被映射为 $Q, K, V$ 。假设输入矩阵为 $X \in \mathbb{R}^{N \times d}$ ，其中 $N$ 是序列长度， $d$ 是词向量维度。

Q = X W^Q, \quad K = X W^K, \quad V = X W^V

这里的 $W^Q, W^K, W^V$ 是三个维度为 $d \times d_k$ 的参数矩阵。
数学本质：这是一个标准的线性变换（仿射变换忽略偏置）。在几何上，这相当于将原本在高维空间中的词向量，投影到了三个不同的特征子空间中。
为什么这么做？因为一个词的“提问特征”、“索引特征”和“语义特征”通常不在同一个方向上。通过三个不同的矩阵，模型允许词在三个不同的几何空间中表达不同的角色。

2. 点积的奥秘： $Q K^T$ 的本质是内积与夹角

得到了 $Q$ 和 $K$ 之后，最核心的操作是计算 $Q K^T$ 。

假设 $Q \in \mathbb{R}^{N \times d_k}$ ， $K \in \mathbb{R}^{N \times d_k}$ ，那么 $Q K^T$ 将生成一个 $N \times N$ 的矩阵，我们称之为注意力分数矩阵（Attention Score Matrix）。

数学本质：矩阵乘法 $Q K^T$ 实际上是在计算 $Q$ 的每一行（当前词的 Query 向量）与 $K^T$ 的每一列（上下文词的 Key 向量）的点积。

在线性代数中，两个向量的点积定义如下：

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

抛开向量长度不看，点积的核心意义在于衡量两个向量在方向上的相似度（夹角 $\theta$ 的大小）。

夹角越小（方向越一致），点积越大，说明 Query 和 Key 越相关。
夹角接近 90 度（正交），点积趋近于 0，说明两者毫无关系。

因此， $Q K^T$ 实际上是构建了一个全局相似度图。矩阵中第 $i$ 行第 $j$ 列的元素，精确地量化了序列中第 $i$ 个词对第 $j$ 个词的“关注程度”。

3. 除以 $\sqrt{d_k}$ ：防止梯度消失的缩放因子

这是公式中经常让人疑惑的一步：为什么要除以 $\sqrt{d_k}$ （Key 向量的维度的平方根）？

数学本质：假设 $q$ 和 $k$ 的每个分量都是均值为 0、方差为 1 的独立随机变量，那么它们的点积 $q \cdot k = \sum_{i=1}^{d_k} q_i k_i$ 的数学期望为 0，但方差会随着维度 $d_k$ 的增大而线性增大（即为 $d_k$ ）。

当模型维度 $d_k$ 很大时（例如 64 或 128），点积的结果会变得非常大，导致不同词之间的得分差异极其悬殊。

如果你直接将这些极大的数值输入到 Softmax 函数中：

\text{softmax}(z_i) = \frac{e^{z_i}}{\sum_j e^{z_j}}

指数函数 $e^x$ 会进一步放大这种差异，使得 Softmax 的输出趋近于 One-Hot 分布（即最大值逼近 1，其余逼近 0）。
在反向传播时，Softmax 在输入值极端不平衡的情况下，梯度会变得非常趋近于 0（梯度消失），导致模型无法有效更新参数。

除以 $\sqrt{d_k}$ ，本质上是做了一次方差归一化，强行将点积的方差拉回 1，保证了无论特征维度多高，Softmax 都能工作在一个梯度稳定的非线性敏感区。

4. Softmax 与乘 $V$ ：概率下的加权动态重构

经过缩放和 Softmax 之后，我们得到了一个完美的 $N \times N$ 概率转移矩阵 $A$ 。其中 $A_{i,j}$ 表示第 $i$ 个词融合第 $j$ 个词信息的比例。

最后一步： $A V$ 。

数学本质：这依然是矩阵乘法。对于输出矩阵的每一行（即某个词经过注意力层后的新表示），其数学计算为：

\text{Output}_i = \sum_{j=1}^{N} A_{i,j} V_j

这是一个向量的凸组合（因为 $A_{i,j}$ 加起来等于 1）。
从几何上看，输出向量 $\text{Output}_i$ 实际上是所有上下文词向量 $V_j$ 在高维空间中根据相关性大小进行的一次**“加权重心计算”**。

在这个过程中，原本静态的、孤立的词向量，通过借用周围词的 Value 向量，完成了语境化的动态重构。比如“苹果”这个词，如果周围有“吃”字，它的输出向量就会向“水果”的语义空间偏移；如果周围有“发布会”字眼，它的输出向量就会向“科技公司”的语义空间偏移。这就是注意力机制的魔力所在。

三、维度升级：多头注意力的几何意义

理解了单头注意力后，我们来看看 Transformer 真正的杀手锏：多头注意力机制。

论文中并没有让整个词向量只做一次注意力计算，而是将其切分为多个头：

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

\text{where } \text{head}_i = \text{Attention}(Q W_i^Q, K W_i^K, V W_i^V)

为什么要这么做？其背后的数学直觉是什么？

如果我们只用一个单头注意力，模型虽然能融合上下文，但它只能学到一种关系模式（比如语法上的主谓关系）。

自然语言是极其复杂的，词与词之间的关系包含：语法依赖、指代关系、语义相似度、时态对应等。
多头注意力的数学本质，是将高维语义空间正交分解为多个低秩子空间。

想象一下，如果模型有 8 个头，就相当于模型生成了 8 个平行的“独立观察者”。每个观察者都被强制投影到一个特定的低维子空间中。在这个子空间里：

Head 1 可能专门计算主谓关系。
Head 2 可能专门计算形容词修饰名词的关系。
Head 3 可能专注于寻找句号，以便断句。

由于各个头的 $W$ 矩阵初始化不同，梯度下降会驱使它们收敛到不同的表征子空间，从而捕捉到不同维度的特征。最后通过 $W^O$ 矩阵再次拼接并投影回原维度，实现了信息的无损多视角融合。

从计算图的角度看，多头注意力不仅没有增加计算复杂度（总维度不变），反而因为降低了每个头的维度（ $d_k = d_{model} / h$ ），极大提升了模型的参数效率和非线性表达能力。

四、代码实战：用 PyTorch 撸一个纯粹的 Self-Attention

理论说了一万遍，不如亲手写一遍代码。下面我们剥离一切无关的模块（如 Mask 机制、Dropout），用纯粹的 PyTorch 实现一个 Scaled Dot-Product Attention，并复现其几何过程。

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

class PureSelfAttention(nn.Module):
    def __init__(self, embed_dim, num_heads):
        super(PureSelfAttention, self).__init__()
        assert embed_dim % num_heads == 0, "Embedding dim must be divisible by num_heads"
        
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        
        # 定义 Q, K, V 的线性投影矩阵
        # 严格来说，这里是多头的融合矩阵，计算时再切分
        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)
        
        # 输出投影矩阵
        self.W_o = nn.Linear(embed_dim, embed_dim, bias=False)
        
    def forward(self, x):
        """
        x shape: (batch_size, seq_len, embed_dim)
        """
        B, N, D = x.size()
        
        # 1. 线性投影：映射到 Q, K, V 空间
        # 此时维度依然是 (B, N, embed_dim)
        Q = self.W_q(x)
        K = self.W_k(x)
        V = self.W_v(x)
        
        # 2. 切分多头：将 embed_dim 拆分为 (num_heads, head_dim)
        # 目标维度：(batch_size, num_heads, seq_len, head_dim)
        Q = Q.view(B, N, self.num_heads, self.head_dim).transpose(1, 2)
        K = K.view(B, N, self.num_heads, self.head_dim).transpose(1, 2)
        V = V.view(B, N, self.num_heads, self.head_dim).transpose(1, 2)
        
        # 3. 计算注意力分数：Q 和 K 的点积
        # Q @ K.transpose(-2, -1) 维度：(B, H, N, head_dim) @ (B, H, head_dim, N) -> (B, H, N, N)
        scores = torch.matmul(Q, K.transpose(-2, -1))
        
        # 4. 缩放：除以 sqrt(head_dim)
        # 这是防止梯度消失的关键数学步骤
        scores = scores / math.sqrt(self.head_dim)
        
        # 5. Softmax 归一化，得到概率分布
        # 沿着最后一维（Key 的序列长度维度）进行归一化
        attn_weights = F.softmax(scores, dim=-1)
        
        # 6. 将注意力权重应用到 Value 上
        # (B, H, N, N) @ (B, H, N, head_dim) -> (B, H, N, head_dim)
        context = torch.matmul(attn_weights, V)
        
        # 7. 拼接多头结果并恢复原始维度
        # (B, H, N, D/H) -> (B, N, H, D/H) -> (B, N, D)
        context = context.transpose(1, 2).contiguous().view(B, N, D)
        
        # 8. 最终的线性变换
        output = self.W_o(context)
        
        return output, attn_weights

# 测试代码
if __name__ == "__main__":
    batch_size = 2
    seq_len = 10
    embed_dim = 512
    num_heads = 8
    
    # 模拟一个输入序列 (例如：经过 Word Embedding 后的向量)
    dummy_input = torch.randn(batch_size, seq_len, embed_dim)
    
    # 初始化纯注意力模块
    attention_layer = PureSelfAttention(embed_dim, num_heads)
    
    # 前向传播
    output, weights = attention_layer(dummy_input)
    
    print(f"Input Shape: {dummy_input.shape}")
    print(f"Output Shape: {output.shape}")
    print(f"Attention Weights Shape (B, H, N, N): {weights.shape}")
    
    # 验证概率分布：对于任意一个词，它对所有词的注意力权重之和应该为 1
    print(f"Sum of attention weights for token 0: {weights[0, 0, 0, :].sum().item()}")

通过这段代码，我们可以清晰地看到张量在各个阶段的维度流转。特别是 scores = torch.matmul(Q, K.transpose(-2, -1)) 这一行，完美对应了我们上文讨论的 $Q K^T$ 的计算。而 view 和 transpose 操作，正是多头机制在代码层面“切分高维空间”的具体体现。

五、终极思考：注意力机制的更抽象数学本质

如果你是一个对数学极为痴迷的研究者，你可能已经察觉到，注意力机制的本质并不局限于自然语言处理。从更宏观的数学视角来看，Transformer 的自注意力其实属于以下两大范畴：

1. 动态图上的信息传递

如果把序列中的每一个 Token 看作图中的一个节点，那么 $Q K^T$ 就是在动态计算节点之间的边权重。
不同于传统的图神经网络（GNN）依靠固定的拓扑图进行消息传递，Transformer 构建的是一个全连接图，并且图上的权重是随着输入样本的不同而实时计算出来的。Self-Attention 本质上完成了一次“动态邻居聚合”。

2. 核平滑与非参数化回归

在统计学中，有一种经典的回归方法叫 Nadaraya-Watson 核回归估计，其公式为：

\hat{y}(x) = \frac{\sum_{i=1}^N K(x, x_i) y_i}{\sum_{j=1}^N K(x, x_j)}

其中 $K(\cdot, \cdot)$ 是衡量距离的核函数。

对比我们的注意力公式：

\text{Output} = \sum_{i=1}^N \text{softmax}(q \cdot k_i) v_i

你会发现，Transformer 注意力机制在数学本质上就是一个可学习的核平滑方法。只不过它没有使用固定的欧式距离或高斯核，而是通过 $W^Q$ 和 $W^K$ 矩阵在数据中自适应地学习出了一个“语义相似度核函数”。

六、总结

在深度学习框架高度封装的今天，调用一个 nn.MultiheadAttention 只需要一行代码。但真正的技术壁垒，永远建立在对底层原理的深刻理解之上。

回顾本文，我们通过拆解 Transformer 的核心公式，揭示了以下几点数学本质：

线性投影：让词在多个正交子空间中扮演不同角色。
点积 $Q K^T$ ：本质是在计算高维向量之间的夹角余弦，衡量语义相似度。
缩放因子 $\sqrt{d_k}$ ：控制点积方差，挽救 Softmax 函数的梯度消失。
Softmax 乘 $V$ ：基于全局概率分布，完成向量的凸组合，实现语境化重构。
多头机制：类似多个独立观察者的集体智慧，捕捉多粒度的复杂关系。

当我们把代码、公式和几何图像在脑海中融为一体时，原本冷冰冰的矩阵运算便充满了生命力。理解了注意力机制的数学本质，我们再去学习 Mamba、RWKV、Linear Attention 等下一代非 Transformer 架构时，就能一眼看透它们到底在优化 Transformer 的哪一个数学瓶颈。

数学，永远是揭开 AI 黑盒最锐利的那把手术刀。

一、 宏观视角：Q、K、V 到底在做什么？

二、 微观拆解：一步步看透矩阵的几何变换

1. 线性投影：升维与子空间切割

2. 点积的奥秘：QKTQ K^TQKT 的本质是内积与夹角

3. 除以 dk\sqrt{d_k}dk​​：防止梯度消失的缩放因子

4. Softmax 与乘 VVV：概率下的加权动态重构

三、 维度升级：多头注意力的几何意义

四、 代码实战：用 PyTorch 撸一个纯粹的 Self-Attention

五、 终极思考：注意力机制的更抽象数学本质