自从 2017 年 Google Brain 的那篇石破天惊的论文《Attention Is All You Need》发表以来，Transformer 架构便彻底重塑了整个深度学习的版图。从自然语言处理（NLP）的 BERT、GPT 系列，到如今席卷全球的多模态大模型（如 Sora、GPT-4o），其底层核心无一例外都是 Transformer。
而在 Transformer 架构中，最核心的灵魂莫过于自注意力机制 。许多初学者对它的认知往往停留在直觉层面——“它让模型关注到句子中重要的词”。但这种粗糙的理解无法解释为什么 Transformer 能够拥有如此强大的表征能力，更无法指导我们在显存优化、长文本处理（如 FlashAttention）等方向上的工程实践。
今天，我们将剥开 API 调用的黑盒，从纯粹的数学视角 出发，深入剖析注意力机制的代数本质、几何意义，并辅以优雅的 PyTorch 实现。
在深入复杂的矩阵微积分之前，我们可以先用一个通俗的“信息检索”模型来理解注意力机制。
假设你在图书馆找书：
查询 ：你脑海中想找的书的特征（例如“关于深度学习的书籍”）。键 ：图书馆里每本书贴的标签（例如“机器学习”、“悬疑小说”）。值 ：书的实际内容。
注意力机制的本质就是：计算你的 Query 与图书馆里所有 Key 的匹配程度（打分），然后根据分数的高低，按比例将对应的 Value 加权求和，最终得到你需要的信息。
在自注意力中，Q , K , V Q, K, V Q , K , V X X X W Q , W K , W V W^Q, W^K, W^V W Q , W K , W V
注意力机制的标准数学表达式如下：
Attention ( Q , K , V ) = softmax ( Q K T d k ) V \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V
Attention ( Q , K , V ) = softmax ( d k Q K T ) V
这个公式看似简单，却蕴含着极其精妙的数学设计。我们将其拆解为四个核心步骤。
Q Q Q K T K^T K T 假设我们有一个长度为 n n n d k d_k d k Q Q Q ( n , d k ) (n, d_k) ( n , d k ) K K K ( n , d k ) (n, d_k) ( n , d k )
当我们计算 Q K T QK^T Q K T ( n , n ) (n, n) ( n , n ) S = Q K T S = QK^T S = Q K T
方阵 S S S i i i j j j S i j S_{ij} S i j Q Q Q i i i q i q_i q i K T K^T K T j j j k j k_j k j 点积 ：
S i j = q i ⋅ k j = ∑ k = 1 d k q i k k j k S_{ij} = q_i \cdot k_j = \sum_{k=1}^{d_k} q_{ik} k_{jk}
S i j = q i ⋅ k j = k = 1 ∑ d k q i k k j k
数学本质 ：在代数中，两个向量的点积代表它们在彼此方向上的投影长度，也就是相似度 。因此，S i j S_{ij} S i j i i i j j j Q K T QK^T Q K T
d k \sqrt{d_k} d k 在经过 Softmax 函数之前，我们需要将 S S S d k \sqrt{d_k} d k
核心推导：为什么是 d k \sqrt{d_k} d k
假设向量 q q q k k k
因为 S i j = ∑ k = 1 d k q i k k j k S_{ij} = \sum_{k=1}^{d_k} q_{ik} k_{jk} S i j = ∑ k = 1 d k q i k k j k
V a r ( S i j ) = ∑ k = 1 d k V a r ( q i k k j k ) Var(S_{ij}) = \sum_{k=1}^{d_k} Var(q_{ik} k_{jk})
V a r ( S i j ) = k = 1 ∑ d k V a r ( q i k k j k )
因为 q i k q_{ik} q i k k j k k_{jk} k j k
V a r ( q i k k j k ) = E [ ( q i k k j k ) 2 ] − ( E [ q i k k j k ] ) 2 Var(q_{ik} k_{jk}) = E[(q_{ik} k_{jk})^2] - (E[q_{ik} k_{jk}])^2
V a r ( q i k k j k ) = E [ ( q i k k j k ) 2 ] − ( E [ q i k k j k ] ) 2
= E [ q i k 2 ] E [ k j k 2 ] − 0 = E[q_{ik}^2]E[k_{jk}^2] - 0
= E [ q i k 2 ] E [ k j k 2 ] − 0
= V a r ( q i k ) V a r ( k j k ) = 1 × 1 = 1 = Var(q_{ik})Var(k_{jk}) = 1 \times 1 = 1
= V a r ( q i k ) V a r ( k j k ) = 1 × 1 = 1
因此，点积结果 S i j S_{ij} S i j
V a r ( S i j ) = d k × 1 = d k Var(S_{ij}) = d_k \times 1 = d_k
V a r ( S i j ) = d k × 1 = d k
几何后果 ：当维度 d k d_k d k S S S
如果我们把这样极其悬殊的数值直接送入 Softmax 函数：
softmax ( z i ) = e z i ∑ j e z j \text{softmax}(z_i) = \frac{e^{z_i}}{\sum_j e^{z_j}}
softmax ( z i ) = ∑ j e z j e z i
由于指数函数 e x e^x e x 饱和区 ，导致梯度极度变小（梯度消失），模型在反向传播时几乎无法更新权重。
解决方案 ：除以 d k \sqrt{d_k} d k
V a r ( S i j d k ) = 1 d k V a r ( S i j ) = d k d k = 1 Var\left(\frac{S_{ij}}{\sqrt{d_k}}\right) = \frac{1}{d_k} Var(S_{ij}) = \frac{d_k}{d_k} = 1
V a r ( d k S i j ) = d k 1 V a r ( S i j ) = d k d k = 1
通过除以 d k \sqrt{d_k} d k
得到缩放后的分数后，我们对其按行进行 Softmax 操作：
A = softmax ( Q K T d k ) A = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)
A = softmax ( d k Q K T )
得到的矩阵 A A A ( n , n ) (n, n) ( n , n ) 离散的概率分布 。
V V V 最后一步，我们将概率分布矩阵 A A A V V V
Output = A V \text{Output} = A V
Output = A V
输出的维度是 ( n , n ) × ( n , d v ) = ( n , d v ) (n, n) \times (n, d_v) = (n, d_v) ( n , n ) × ( n , d v ) = ( n , d v )
对于序列中的第 i i i 整个序列中所有词的 Value 向量的加权期望（凸组合） 。权重就是刚才算出的注意力概率。
几何解释 ：这一步实际上是在由 Value 向量张成的向量空间中，根据注意力权重寻找一个特定的坐标点。它实现了全局信息的融合，让每一个词都拥有了上下文语境。
如果仅仅使用单头注意力，模型很容易将注意力集中在单一的上下文特征上。Transformer 引入了多头注意力 。
MultiHead ( Q , K , V ) = Concat ( head 1 , . . . , head h ) W O \text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, ..., \text{head}_h)W^O
MultiHead ( Q , K , V ) = Concat ( head 1 , . . . , head h ) W O
where head i = Attention ( Q W i Q , K W i K , V W i V ) \text{where } \text{head}_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)
where head i = Attention ( Q W i Q , K W i K , V W i V )
从数学角度看，多头注意力本质上是一种高维空间的低秩分解与特征集成 。
降维打击 ：假设词嵌入维度 d m o d e l = 512 d_{model} = 512 d m o d e l = 5 1 2 h = 8 h=8 h = 8 Q , K , V Q, K, V Q , K , V d k = d v = d m o d e l / h = 64 d_k = d_v = d_{model} / h = 64 d k = d v = d m o d e l / h = 6 4 独立子空间 ：在这 8 个独立的 64 维子空间中，模型可以并行地学习不同类型的依赖关系。例如，第 1 个头可能关注语法主谓关系，第 2 个头可能关注时态，第 3 个头关注代词指代。空间融合 ：最后将 8 个头的输出拼接起来，并通过一个线性映射矩阵 W O W^O W O
如果我们抛开代数运算，从图论的角度来看待 Transformer，会有更深刻的理解。
实际上，自注意力机制完全等价于图神经网络（GNN）中的消息传递机制 。
图的构建 ：输入序列构成了一个全连接图。节点就是序列中的 Token。消息传递 ：
Q Q Q K K K A A A 带权邻接矩阵 。与普通 GNN 不同的是，这个邻接矩阵是动态计算的，且是密集的。V V V 特征消息 。A V AV A V
这种视角的转换，解释了为什么 Transformer 能够处理极其复杂的全局依赖，因为它的底层逻辑就是一个全连接图上的特征聚合过程。
理解了数学原理后，我们用 PyTorch 从零实现一个标准的多头注意力模块。这比直接调用 nn.MultiheadAttention 更能让人看清内部的张量流动。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 import torchimport torch.nn as nnimport torch.nn.functional as Fimport mathclass MultiHeadAttention (nn.Module): def __init__ (self, d_model, num_heads ): super (MultiHeadAttention, self ).__init__() assert d_model % num_heads == 0 , "d_model 必须能被 num_heads 整除" self .d_model = d_model self .num_heads = num_heads self .d_k = d_model // num_heads self .W_q = nn.Linear(d_model, d_model) self .W_k = nn.Linear(d_model, d_model) self .W_v = nn.Linear(d_model, d_model) self .W_o = nn.Linear(d_model, d_model) def scaled_dot_product_attention (self, Q, K, V, mask=None ): """ 缩放点积注意力 Q, K, V 的维度: (batch_size, num_heads, seq_len, d_k) """ scores = torch.matmul(Q, K.transpose(-2 , -1 )) / math.sqrt(self .d_k) if mask is not None : scores = scores.masked_fill(mask == 0 , -1e9 ) attn_weights = F.softmax(scores, dim=-1 ) output = torch.matmul(attn_weights, V) return output, attn_weights def forward (self, query, key, value, mask=None ): batch_size = query.size(0 ) Q = self .W_q(query).view(batch_size, -1 , self .num_heads, self .d_k).transpose(1 , 2 ) K = self .W_k(key).view(batch_size, -1 , self .num_heads, self .d_k).transpose(1 , 2 ) V = self .W_v(value).view(batch_size, -1 , self .num_heads, self .d_k).transpose(1 , 2 ) attn_output, attn_weights = self .scaled_dot_product_attention(Q, K, V, mask) attn_output = attn_output.transpose(1 , 2 ).contiguous().view(batch_size, -1 , self .d_model) final_output = self .W_o(attn_output) return final_output, attn_weights if __name__ == "__main__" : d_model = 512 num_heads = 8 seq_len = 10 batch_size = 4 mha = MultiHeadAttention(d_model, num_heads) x = torch.randn(batch_size, seq_len, d_model) output, weights = mha(x, x, x) print (f"输入形状: {x.shape} " ) print (f"输出形状: {output.shape} " ) print (f"注意力权重形状: {weights.shape} " )
在这段代码中，有几个极其精妙的张量操作细节值得品味：
维度的重排 (transpose(1, 2)) ：在拆分多头后，我们并没有将 seq_len 和 num_heads 连续存放，而是将 num_heads 移到了 batch_size 之后。这样做的目的是使得在后续计算 QK^T 时，不同头在物理内存上完全独立，从而实现真正的并行矩阵乘法。掩码机制 (masked_fill) ：我们在 scores 上加上了一个极大的负数（-1e9）。在数学上，经过 Softmax 的指数运算后，e − 1 e 9 ≈ 0 e^{-1e9} \approx 0 e − 1 e 9 ≈ 0
回顾全文，Transformer 注意力机制的数学本质并不晦涩，但其组合却极其强大：
相似度度量 ：利用矩阵乘法 Q K T QK^T Q K T 方差控制 ：引入 d k \sqrt{d_k} d k 信息聚合 ：通过 Softmax 加权求和，实现图拓扑结构上的消息传递与特征融合。空间变换 ：利用多头机制在多个低维正交子空间中捕获异构的依赖关系。
然而，天下没有免费的午餐。注意力机制的 Q K T QK^T Q K T O ( N 2 ) O(N^2) O ( N 2 ) N N N N N N O ( N 2 ) O(N^2) O ( N 2 )
这也正是当前 AI 基础研究的热点方向所在：从 FlashAttention（通过硬件感知的分块计算减少 HBM 访问）到线性注意力（Linear Attention，通过核函数变换将 Softmax 解耦，将复杂度降至 O ( N ) O(N) O ( N )
理解了这些底层的数学逻辑，当你下次面对大模型长长的上下文窗口报错，或者试图优化推理速度时，脑海中浮现的将不再是神秘的黑盒，而是一张张清晰的矩阵乘法图。