AI Research

04. Self-Attention：Transformer 的核心计算

04. Self-Attention：Transformer 的核心计算

Self-Attention 不是一个”技巧”，而是 Transformer 之所以为 Transformer 的根本原因。 理解它，不能只记住 QKTVQK^TV 的计算步骤——你需要理解它为什么被设计成这样。

一、序列建模的本质困难

语言是一个序列。序列中的每个词，其含义都取决于上下文：

“苹果发布了新产品” vs “他吃了一个苹果”

同一个”苹果”，意思完全不同。所以模型不能孤立地处理每个词——它必须让词与词之间交换信息。

问题是：怎么交换？

RNN 的方案：逐步传递

RNN 把信息像接力棒一样逐步传递：词 1 → 词 2 → 词 3 → …

这有两个致命问题：

  1. 长距离衰减：词 1 的信息传到词 100 时，已经被多次压缩、覆盖，几乎消失
  2. 无法并行：必须按顺序计算，GPU 的并行能力完全浪费

CNN 的方案：局部窗口

CNN 用固定大小的卷积核扫描序列，每次只看局部几个词。

问题：要捕捉远距离关系，就得堆很多层——第 1 层看 3 个词，第 2 层看 5 个词……效率低，且远距离信号仍然被间接传递。

Self-Attention 的方案：全局直连

Self-Attention 的核心思想极其简单：

让序列中的每个位置，直接与所有其他位置交互，一步到位。

不需要逐步传递，不需要堆叠多层。词 1 和词 100 之间的关系，在一次计算中就能建立。

这就是 Transformer 论文标题 “Attention Is All You Need” 的含义——不需要 RNN 的递归，不需要 CNN 的卷积，仅靠注意力机制就能完成序列建模。

二、Self-Attention 在做什么：一个精确的描述

在进入公式之前，先用自然语言精确描述 Self-Attention 的计算：

输入：nn 个向量（对应 nn 个 token），每个向量 dd 维

输出：nn 个新向量，每个向量仍然 dd 维

每个输出向量是怎么算的？

对于位置 ii 的输出：

  1. 位置 ii 向所有位置（包括自己）“提问”：你跟我有多相关？
  2. 得到一组相关度分数
  3. 把这些分数归一化成概率分布（加起来等于 1）
  4. 用这个概率分布，对所有位置的信息做加权平均

结果：位置 ii 的输出向量，是整个序列信息的加权混合，权重由”相关度”决定。

这就是”注意力”的含义：不是平等地看所有词，而是有选择地关注。

三、Q / K / V：不是三种数据，是三个”角色”

这是最容易产生困惑的地方。Q、K、V 不是三种不同的输入数据——它们来自同一个输入，通过三个不同的线性变换产生：

Q=XWQ,K=XWK,V=XWVQ = XW_Q, \quad K = XW_K, \quad V = XW_V

为什么需要三个？因为”计算相关度”和”提取信息”是两个不同的任务，用同一个表示来做这两件事会产生冲突。

类比：搜索引擎

想象你在用搜索引擎：

角色含义对应
Query你输入的搜索词”我在找什么”
Key每个网页的索引关键词”这个网页关于什么”
Value网页的实际内容”这个网页能提供什么”

搜索过程：

  1. 用你的搜索词（Q）去匹配每个网页的关键词（K）→ 得到相关度排名
  2. 按相关度排名，提取网页内容（V）

为什么不能 Q = K？

如果 Query 空间和 Key 空间是同一个，模型就被迫用同一个表示来同时完成”表达我在找什么”和”表达我能提供什么”。这两个任务的需求不同：

  • “我在找什么”需要表达需求（比如”他”这个代词需要找到它的指代对象）
  • “我能提供什么”需要表达身份（比如”小明”需要表达自己是一个人名）

分开投影，让模型有更大的自由度来学习这两种不同的表示。

为什么不能直接用 K 代替 V？

K 是为”匹配”优化的表示——它被训练来回答”你跟我有多相关”。 V 是为”信息传递”优化的表示——它被训练来回答”我应该给你什么信息”。

一个好的匹配标签不一定是好的信息载体。比如一本书的标题（K）可以帮你判断相关性，但你真正需要的是书的内容（V）。

四、完整公式：逐步拆解

Attention(Q,K,V)=softmax ⁣(QK⊤dk)V\text{Attention}(Q, K, V) = \text{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right) V

把这个公式拆成 4 个阶段：

阶段 1：计算相关度矩阵 QK⊤QK^\top

S=QK⊤shape: [n,n]S = QK^\top \quad \text{shape: } [n, n]

SijS_{ij} 表示位置 ii 的 Query 和位置 jj 的 Key 之间的点积相似度。

点积为什么能衡量相似度？因为两个向量的点积 = 它们的模之积 × 夹角的余弦。方向越一致，点积越大。

SS 矩阵的每一行，是一个 token 对所有 token 的相关度打分。

阶段 2：缩放 ÷dk\div \sqrt{d_k}

S′=SdkS' = \frac{S}{\sqrt{d_k}}

为什么需要这一步？ 这不是可有可无的细节，而是数值稳定性的关键。

点积的期望值会随维度 dkd_k 线性增长。假设 QQ 和 KK 的每个分量都是均值 0、方差 1 的独立随机变量，则：

E[q⋅k]=0,Var[q⋅k]=dk\mathbb{E}[q \cdot k] = 0, \quad \text{Var}[q \cdot k] = d_k

当 dk=64d_k = 64 时，点积的标准差就是 8，容易出现绝对值很大的分数。

大的分数进入 softmax 后，会产生接近 one-hot 的分布——几乎所有权重集中在一个位置上。这导致：

  • 梯度消失（softmax 饱和区梯度趋近于零）
  • 模型失去”综合多个位置信息”的能力

除以 dk\sqrt{d_k} 后，方差回到 1，softmax 的输入在合理范围内，分布既不太尖也不太平。

你可以把 dk\sqrt{d_k} 理解为一个自适应温度系数：维度越高，自动降温越多。

阶段 3：Softmax 归一化

A=softmax(S′)（按行做）A = \text{softmax}(S') \quad \text{（按行做）}

Softmax 做两件事：

  1. 把分数变成正数（通过 exe^x）
  2. 把每行归一化为概率分布（加和为 1）

AijA_{ij} 的含义：位置 ii 分配给位置 jj 的注意力权重。

每行加和为 1，意味着每个 token 的注意力预算是固定的——给了 A 更多注意力，就必须从 B 那里减少。这是一种竞争机制。

阶段 4：加权求和

Output=AVshape: [n,dk]\text{Output} = AV \quad \text{shape: } [n, d_k]

每个输出向量：

outi=∑j=1nAij⋅Vj\text{out}_i = \sum_{j=1}^{n} A_{ij} \cdot V_j

这是一个加权平均：把所有位置的 Value 向量按注意力权重混合在一起。

核心直觉：Self-Attention 的输出，是对整个序列做了一次软性信息检索。不是硬性地选出一个位置，而是按相关度把所有位置的信息融合起来。

五、维度变化全览

矩阵Shape含义
XX[n, d_model]输入序列
WQ,WK,WVW_Q, W_K, W_V[d_model, d_k]投影权重
Q,K,VQ, K, V[n, d_k]投影后的表示
QK⊤QK^\top[n, n]相关度矩阵
softmax(⋅)\text{softmax}(\cdot)[n, n]注意力权重
Output[n, d_k]融合后的表示

关键观察：中间产生了一个 [n, n] 的矩阵。这就是 Self-Attention O(n2)O(n^2) 复杂度的来源——序列长度翻倍，计算量翻四倍。

六、手算示例：3 个 token，2 维

用最小的数字，保证你能追踪每一步。

输入

3 个 token，每个 2 维：

X=[100111]X = \begin{bmatrix} 1 & 0 \\ 0 & 1 \\ 1 & 1 \end{bmatrix}

为简化，令 WQ=WK=WV=IW_Q = W_K = W_V = I（单位矩阵），所以 Q=K=V=XQ = K = V = X。

Step 1: QK⊤QK^\top

QK⊤=[100111][101011]=[101011112]QK^\top = \begin{bmatrix} 1 & 0 \\ 0 & 1 \\ 1 & 1 \end{bmatrix} \begin{bmatrix} 1 & 0 & 1 \\ 0 & 1 & 1 \end{bmatrix} = \begin{bmatrix} 1 & 0 & 1 \\ 0 & 1 & 1 \\ 1 & 1 & 2 \end{bmatrix}

读法：token C（第 3 行）跟自己最相关（分数 2），跟 A 和 B 同等相关（分数 1）。

Step 2: 缩放

dk=2d_k = 2，所以除以 2≈1.414\sqrt{2} \approx 1.414：

S′=[0.70700.70700.7070.7070.7070.7071.414]S' = \begin{bmatrix} 0.707 & 0 & 0.707 \\ 0 & 0.707 & 0.707 \\ 0.707 & 0.707 & 1.414 \end{bmatrix}

Step 3: Softmax（按行）

以第 3 行 [0.707, 0.707, 1.414] 为例：

e0.707≈2.028,e0.707≈2.028,e1.414≈4.113e^{0.707} \approx 2.028, \quad e^{0.707} \approx 2.028, \quad e^{1.414} \approx 4.113 sum=8.169\text{sum} = 8.169 weights=[0.248,  0.248,  0.504]\text{weights} = [0.248, \; 0.248, \; 0.504]

含义：token C 把约 50% 的注意力放在自己身上，各约 25% 给 A 和 B。

Step 4: 加权求和

outC=0.248×[1,0]+0.248×[0,1]+0.504×[1,1]=[0.752,  0.752]\text{out}_C = 0.248 \times [1, 0] + 0.248 \times [0, 1] + 0.504 \times [1, 1] = [0.752, \; 0.752]

这个输出向量不再是”纯粹的 C”，而是融合了 A 和 B 信息的 C。

七、Causal Mask：让模型”不能偷看未来”

在自回归生成（GPT 式模型）中，预测第 tt 个词时只能看到前 t−1t-1 个词。但标准 Self-Attention 会让每个位置看到所有位置——包括未来。

解决方案：在 softmax 之前，把未来位置的分数设为 −∞-\infty：

mask=[0−∞−∞00−∞000]\text{mask} = \begin{bmatrix} 0 & -\infty & -\infty \\ 0 & 0 & -\infty \\ 0 & 0 & 0 \end{bmatrix}

e−∞=0e^{-\infty} = 0，所以 softmax 后这些位置的权重为 0，信息完全被屏蔽。

注意顺序：必须先加 mask 再 softmax。如果先 softmax 再置零，权重之和就不再等于 1，概率分布被破坏。

加上 mask 后，完整公式变为：

Attention(Q,K,V)=softmax ⁣(QK⊤dk+mask)V\text{Attention}(Q, K, V) = \text{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}} + \text{mask}\right) V

Encoder（如 BERT）不需要 causal mask，因为它可以双向看；Decoder（如 GPT）必须用。

八、Multi-Head Attention：多个子空间并行

单头注意力只有一组 WQ,WK,WVW_Q, W_K, W_V，只能学到一种”相关性模式”。但语言中的关系是多维度的：

  • 语法关系：主语-谓语-宾语
  • 指代关系：“他” → “小明”
  • 语义关系：“巴黎” → “法国”
  • 位置关系：相邻词之间的局部依赖

Multi-Head Attention 的做法：

  1. 把 dmodeld_{model} 维度拆成 hh 个头，每个头 dk=dmodel/hd_k = d_{model} / h 维
  2. 每个头独立做一次完整的 Self-Attention
  3. 把 hh 个头的输出拼接起来
  4. 再过一个线性变换 WOW_O，映射回 dmodeld_{model} 维
MultiHead(X)=Concat(head1,…,headh)WO\text{MultiHead}(X) = \text{Concat}(\text{head}_1, \dots, \text{head}_h) W_O headi=Attention(XWQi,XWKi,XWVi)\text{head}_i = \text{Attention}(XW_Q^i, XW_K^i, XW_V^i)

关键洞察：计算量和单头几乎相同（因为每个头的维度更小），但模型能同时关注多种不同的关系。

实际观察到的现象（可视化研究）：

  • 有些头学会了关注相邻位置（局部语法）
  • 有些头学会了长距离的指代消解
  • 有些头近似均匀分布（捕捉全局统计信息）
  • 不同层的头关注模式差异很大

九、Self-Attention 在 Transformer 中的位置

Self-Attention 不是孤立存在的。在一个 Transformer 层中，它和其他组件配合工作：

输入 x
  │
  ├──→ Self-Attention ──→ + ←── 残差连接
  │                        │
  │                    LayerNorm
  │                        │
  │                        ├──→ FFN ──→ + ←── 残差连接
  │                                     │
  │                                 LayerNorm
  │                                     │
  └─────────────────────────────────→ 输出

每个组件的分工：

组件职责
Self-Attentiontoken 间的信息交换（“看别人”）
FFN单个 token 内部的非线性变换（“消化信息”）
残差连接保证梯度流通 + 保留原始信息
LayerNorm稳定训练、加速收敛

Self-Attention 负责跨位置的通信，FFN 负责位置内部的计算。两者交替堆叠，形成 Transformer 的主体。

十、从直觉到本质：Self-Attention 是什么？

现在你已经看完了所有细节，我们可以从更高的视角来理解 Self-Attention。

它是一种动态的、数据依赖的信息路由

传统神经网络（全连接层、卷积层）的连接权重是固定的——训练好之后，权重就不变了，无论输入是什么。

Self-Attention 的”权重”（注意力矩阵 AA）是动态计算的——它取决于当前的输入是什么。不同的输入序列，会产生完全不同的注意力模式。

这意味着：模型可以根据内容自适应地决定信息流向。当”他”出现在”小明吃了苹果，他很开心”中时，模型会学到把”小明”的信息路由到”他”；而在”她吃了苹果”中，同样的”他”可能被忽略。

它是一种可微分的软性寻址

如果你有计算机体系结构的背景：Self-Attention 类似于一次内容寻址的内存读取。

  • Value 矩阵 = 内存
  • Query = 读取地址（基于内容，而非基于位置）
  • Key = 地址标签
  • 注意力权重 = 软性地址匹配结果

与硬件内存不同的是，这里的”寻址”是软性的（softmax 产生概率分布）、可微分的（可以用梯度下降训练），而且可以同时从多个位置读取信息。

十一、常见误解澄清

❌ “Self-Attention 能理解语义”

Self-Attention 本身只是一个加权求和——它不”理解”任何东西。语义理解是通过大量数据训练后，在权重中涌现出来的。

❌ “注意力权重 = 模型的推理依据”

注意力权重可视化经常被用来”解释”模型行为，但研究表明，注意力权重和模型的实际决策之间并没有可靠的因果关系。

❌ “位置 ii 的输出只取决于注意力权重高的几个位置”

不是。所有位置都有贡献（除非被 mask 屏蔽），只是贡献的大小不同。即使权重很小的位置，其 Value 也会被纳入加权和。

❌ “Self-Attention 自带位置信息”

恰恰相反——Self-Attention 本身是置换不变的（permutation invariant）。打乱输入序列的顺序，只要 Q、K、V 跟着打乱，输出（在打乱后的排列下）完全相同。这就是为什么 Transformer 需要额外添加位置编码（Positional Encoding）。

十二、计算复杂度与优化方向

Self-Attention 的时间和空间复杂度都是 O(n2)O(n^2)（nn 为序列长度）：

  • 时间：QK⊤QK^\top 需要 n×nn \times n 次点积
  • 空间：需要存储 n×nn \times n 的注意力矩阵

当 n=4096n = 4096 时，注意力矩阵有 1600 万个元素；n=128Kn = 128\text{K} 时，这个数字达到 160 亿。

这是 Transformer 处理长序列的主要瓶颈。主流优化方向：

方法思路
FlashAttention利用 GPU 内存层次结构，分块计算，避免存储完整的 n×nn \times n 矩阵
稀疏注意力只计算部分位置对的注意力（如局部窗口 + 全局 token）
线性注意力用核函数近似 softmax，把复杂度降到 O(n)O(n)
Ring Attention多设备分布式计算，每个设备只处理序列的一部分

十三、自测题

  1. 为什么 Self-Attention 需要三个独立的投影（Q、K、V），而不是只用一个？
  2. dk\sqrt{d_k} 缩放解决的是什么问题？如果不缩放会怎样？
  3. 为什么 causal mask 必须在 softmax 之前加？
  4. Self-Attention 的输出维度和输入维度有什么关系？
  5. 为什么说 Self-Attention 是”置换不变”的？这带来了什么问题？
  6. Multi-Head Attention 的计算量比单头大吗？为什么？

本章总结

Self-Attention 的本质是一种动态的、数据依赖的信息路由机制：

  1. 每个位置通过 Q 和 K 的匹配，计算与其他所有位置的相关度
  2. 通过 softmax 将相关度转化为概率分布
  3. 用这个分布对所有位置的 V 做加权求和，完成信息融合

它解决了序列建模的核心问题：如何让任意两个位置直接交换信息，且交换的方式由数据本身决定。

下一篇：05. Multi-Head Attention 与 FFN 我们会深入 Multi-Head 的实现细节，以及 FFN 为什么是必不可少的——Self-Attention 只负责”交流”，FFN 才负责”思考”。

目录