土法炼钢兴趣小组的算法知识备份

【Transformer 与注意力机制】14|Self-Attention:让序列自己看自己

文章导航

分类入口
transformer
标签入口
#transformer#attention#self-attention#permutation-equivariance

目录

第 13 篇把 Q/K/V 三件套立起来了。这一篇继续往前走:当 Q、K、V 都来自同一序列时,attention 退化成 self-attention——这是 Transformer 的真正主力。

但「同一个序列」这个限制看似简单,会带来一连串很深的后果:

这一篇会把这四点逐一拆开。


一、从 cross 到 self:一个退化路径

1.1 cross-attention 是「两个序列对话」

第 12 篇的 Bahdanau 是 cross-attention:

两个不同的序列通过 attention 连起来:「decoder 在生成下一个 target 词时,回头看 source 序列的哪些位置」。

Bahdanau 的「对齐」就是这种跨序列对话的产物。

到了 Transformer,encoder-decoder 模型里仍然保留了 cross-attention,连接 encoder 输出和 decoder 输入。

但 Transformer 引入了一个全新的东西:让一个序列自己看自己——self-attention。

1.2 self-attention 的形式

self-attention 把 Q、K、V 都从同一序列 X 投影出来:

Q = X · W_Q
K = X · W_K
V = X · W_V

公式与 cross-attention 完全一样:

Attention(Q, K, V) = softmax(QKᵀ / √d_k) · V

唯一区别:Q 和 K/V 的「源」是同一个 X。

但这一个改动的后果非常深远。

Cross vs Self

1.3 cross-attention 与 self-attention 在公式上的对照

cross-attention:
  Q = X_target · W_Q     # 来自一个序列
  K = X_source · W_K     # 来自另一个序列
  V = X_source · W_V

self-attention:
  Q = X · W_Q            # 都来自同一序列
  K = X · W_K
  V = X · W_V

形式上只差一个「下标 source/target」。

但语义上,self-attention 让序列内部的所有 token 互相沟通,cross-attention 让两个序列之间的 token 互相沟通。

这是两件不同的事,Transformer 同时用两种 attention:

到第二十篇 Transformer 整体架构会把这三种摆在一起讲。


二、self-attention 的 O(1) 跳数

2.1 RNN 的距离问题

RNN 的工作机制:从左到右处理 token,每一步把信息塞进 hidden state,往下传。

如果你有一个长度 100 的序列,第 1 个 token 的信息要到第 100 个 token 才被使用,必须沿 99 个 RNN cell 一步步传过去。

每一步都有「信息衰减」:sigmoid/tanh 激活的非线性,加上 LSTM/GRU 的门控,让信息在传递过程中或被覆盖、或被丢弃。

理论上 LSTM 能保留长依赖,实践上长度过 50 之后准确传递特定信息已经非常困难——这是「梯度消失/信息衰减」的双重困境。

2.2 self-attention 的「直达」

self-attention 不需要逐步传递。

第 1 个 token 想知道第 100 个 token 的信息?只要它的 query 与第 100 个 token 的 key 相似度高,就直接通过 α 把 V 拿过来。

任意两个 token 之间,路径长度 = 1

这件事在论文里被称为「O(1) hops」——是 self-attention 取代 RNN 的核心论点之一。

Token 互看

2.3 但 O(1) 不是免费的

代价:

第一,计算复杂度 O(N²)。任意两 token 都要算相似度,N 个 token 就有 N² 个 (i, j) 对。

第二,显存占用 O(N²)。QKᵀ 矩阵的形状是 N×N,长序列下显存爆炸。

第三,没有递归——RNN 的信息流是天然顺序的,self-attention 必须靠 mask 来实现「不能看未来」(causal mask)。

第四,没有位置感——下一节会讲。

每个代价都引出了一个研究方向:第二代 attention 变体(FlashAttention 优化显存、Linear/Sparse Attention 降低复杂度、各种位置编码方案)。

到第十八篇 attention 复杂度那里我们会把这些放在一起讨论。


三、Permutation-Equivariance:self-attention 不知道位置

3.1 一个关键性质

性质:如果你把输入序列 X 的 token 顺序打乱(任意排列),self-attention 的输出也会以相同的方式被打乱——但不会改变内容

数学语言:

SelfAttn(P · X) = P · SelfAttn(X)

其中 P 是任意排列矩阵。

这叫做 permutation-equivariant(排列等变)。

3.2 为什么是这样

直觉证明:

self-attention 里的所有计算都是「对每对 (i, j) 做运算 + softmax 归一化 + 加权求和」。

这些运算都不依赖绝对位置——它们只依赖每对 token 的内容。

所以打乱顺序,每对 (i, j) 还在那儿,只是换了名字。

输出按相同顺序打乱,内容不变。

3.3 这件事意味着什么

意味着:「The cat sat on the mat」和「mat the on sat cat the」在 self-attention 看来是完全等价的。

但语言学上这两个序列差异巨大——前者是合法句子,后者是无意义的词袋。

self-attention 没有任何机制能区分它们。

这就是「self-attention 不知道位置」——它只看 token 内容,不看 token 在序列中的位置。

3.4 RNN 没有这个问题

RNN 沿着 token 顺序处理,第 1 个 token 进 RNN 时 hidden state 是初始值,第 2 个 token 进时 hidden state 已经被第 1 个改过了。

「先来后到」天然刻在 RNN 的状态里。

这是 RNN 的优势之一——位置感是免费的。

但代价是不能并行(必须按顺序处理)。

3.5 解药:位置编码

既然 self-attention 不知道位置,就把位置信息加到输入里

最简单的做法:给每个位置 i 一个固定向量 PE_i,然后 X_i = TokenEmbed_i + PE_i 作为 self-attention 的输入。

PE 的设计有很多种——Sinusoidal(Vaswani 2017)、Learned(BERT、GPT-2)、Relative(Shaw 2018)、RoPE(Su 2021)、ALiBi(Press 2021)。

每种都有自己的优缺点,会单独写一篇(第二十六到三十篇预告)。

但所有方案的本质都是:把「位置」变成可被 attention 看到的信号

3.6 一个微妙的点:input 和 output 都被打乱吗

permutation-equivariance 说的是「input 打乱 → output 同样打乱」。

但加了位置编码后,PE 是固定按位置的:第 1 个位置的 PE 永远是 PE_1,第 2 个永远是 PE_2。

如果你把 X 的 token 打乱但 PE 保持不变,那么打乱后的 (token, PE) 组合就变了——self-attention 的输出会真的不一样了。

这就是位置编码奏效的机制:把位置当作内容的一部分注入

理论上 self-attention 仍然 permutation-equivariant,但 input 不再是「纯 token」,而是「token + 位置」的合体。

3.7 一个更深的视角:可换性与对称性

permutation-equivariance 在数学上属于「群对称性」。

具体来说,self-attention 关于排列群 S_N 是对称的(equivariant)。

CNN 关于平移群 ℤ² 是对称的(translation-equivariant);GNN 关于图同构群是对称的;MLP 没有显式对称性。

这些「对称性」正是各种架构的归纳偏置(inductive bias)。

self-attention 的归纳偏置最弱(只有排列对称),所以它需要更多数据来弥补——但弥补之后能学到更通用的 pattern。

这是 Bronstein 等 2021 年「Geometric Deep Learning」框架的核心观点:架构的对称性决定它的归纳偏置,进而决定它的样本效率与表达力 tradeoff

self-attention 的对称性最弱、表达力最强、样本需求最大——三者一体。

3.8 排列等变 ≠ 排列不变

注意区分:

equivariant:input permute,output 同步 permute(self-attention 是这种)。

invariant:input permute,output 不变(比如对所有 token 取均值的全局 pooling 是这种)。

self-attention 不是 invariant 的——如果你 permute input,output 也会 permute(只是内容不变)。

只有再过一个 invariant 的 head(比如取所有 token 输出的均值或 [CLS] token)才能得到对集合不变的最终表示。

BERT 的 [CLS] token 就是这样:self-attention 之后,取 [CLS] 位置的输出作为整个句子的表示——配合 self-attention 的 equivariance,得到一个对 token 顺序不敏感(除了 PE 影响)的整体表示。


四、一个具体例子:It 看到 Cat

4.1 例句

The cat sat on the mat. It was tired.

人类一眼能看出 “It” 指的是 “cat”。

这是共指消解(coreference resolution)——NLP 里一个经典任务。

self-attention 能不能学到这种关系?

答案是:。Vaswani 论文 Fig 3 就展示了一个 head 学到了类似的共指模式(虽然原文用的句子稍有不同)。

It-Cat

4.2 一个 head 的 attention 模式

模型训练好之后,看 “It” 这个 query 在某个 head 里的 attention 权重:

key token α (示意)
The 0.02
cat 0.61
sat 0.11
on 0.03
the 0.02
mat 0.07
. 0.01
It 0.10
was 0.02
tired 0.01
. 0.00

注意:最高的不是 “It” 自己,而是 “cat”(约 0.61)。

这表明该 head 学到了「代词 → 先行词」的共指关系。

注:这是示意权重,真实模型每个 head 学到的模式不同;不同 head 可能 attend 不同对象。BertViz、AttentionViz 等工具可以让你看真实模型的 attention pattern。

4.3 不同 head 的不同关注

multi-head attention 的妙处在于:12 head(BERT-base)或 32 head(BERT-large)各自学一个不同的 W_Q/W_K/W_V,所以每个 head 关注不同的东西。

经验观察(来自 Clark et al. 2019、Voita et al. 2019):

「不同 head 学不同的 pattern」是 multi-head attention 性能优势的关键,第十六篇会专门讲。

4.4 数值上的「直接」

回到 Q/K/V 公式。

设 “It” 的位置是 8,“cat” 的位置是 2。

模型学到的 W_Q[head] 把 “It” 投影成 q_8;W_K[head] 把 “cat” 投影成 k_2。

q_8 · k_2 远大于 q_8 · k_其他,所以 softmax 后 α_{8, 2} 最大。

最终 o_8 = Σ α_{8, j} v_j ≈ 0.61 · v_2 + …,大量信息来自 v_2(“cat” 的 V 投影)。

这种「直接性」是 RNN 完全做不到的:RNN 要把 “cat” 的信息从位置 2 传到位置 8,必须穿过 “sat、on、the、mat、.” 五个中间步骤,沿途难免衰减。

self-attention 一步直达——这是它能在长距离依赖任务上吊打 RNN 的根本原因。

4.5 多层叠加之后的「抽象」

但这只是单层 self-attention 的故事。

实际 BERT/GPT 有 12-96 层 self-attention 叠加。

第一层学到的可能是浅层 pattern(局部 n-gram、临近 token);

第二、三层开始组合(短语级关系、依存关系);

中层学到语法(主谓、动宾、修饰);

深层学到语义(共指、指代、推理)。

这种「层级抽象」类似于 CNN 在视觉里的 pattern:浅层学边缘、中层学纹理、深层学物体。

self-attention 在 NLP 里有类似的层级——这件事到第二十一篇(Transformer 整体架构)会更详细。

但要注意:「层级抽象」是经验观察,不是定理。不同任务、不同初始化下,同一层学到什么并不固定。

4.6 为什么单 head 不够

观察 4.2 的权重表,“It” 也对自己 attend 了 0.10。

这表明该 head 不是纯粹的「共指 head」,还混入了一些「自我引用」信号。

如果只有一个 head,模型必须在「找共指」和「保留自身信息」之间妥协——做不到两者兼顾。

multi-head 的妙处就在这里:head_1 负责共指(指向 “cat”),head_2 负责保留自身(self-loop),head_3 负责局部上下文(attend “was”)……

每个 head 各司其职,最后 concat 起来过 W_O 输出。

multi-head 不是「多算几遍取平均」——它是功能分解

第十六篇会把这件事拆得更细。


五、一个要小心的事:attention ≠ 解释

5.1 漂亮的图很诱人

看到 “It” attend 到 “cat” 的权重 0.61,你可能立刻想说:「模型理解了共指!」

慢一点。

attention 权重 α 是一个数学产物,它只告诉你「在这一步加权求和时,该 head 给 cat 的 V 多少权重」。

不直接等于「模型用 cat 的信息做出预测」。

为什么?

第一,最终预测受多个 head、多层 attention、FFN 的复合影响——单个 head 的 α 只反映一个分量。

第二,attention 之外有残差连接——信息可以绕过 attention 直接传递,这条路径不被 α 反映。

第三,多层堆叠会让信号充分混合——浅层 attention 的高 α 不一定意味着深层的关键信号也来自那里。

5.2 Jain & Wallace 2019 的实验

Jain & Wallace 用一系列实验论证「attention is not explanation」:

结果令人沮丧:很多任务上,替换 α 后预测几乎不变

这说明 α 并不是模型决策的关键中介——它只是一个中间计算量。

5.3 Wiegreffe & Pinter 2019 的反驳

后来 Wiegreffe & Pinter 2019 写了一篇「attention is not not explanation」反驳。

他们论证:在更严格的实验设置下,attention 仍然在某种意义上提供了部分解释——只是不能简单地把 α 当作「重要度排序」。

这场争论的结论大致是:α 提供「相关线索」,但不是因果解释——把它当作可视化工具可以,把它当作模型解释要谨慎。

5.4 实践建议

第一,看 attention 图时,多看几个 head,不要被某一个 head 的「漂亮 pattern」骗了。

第二,用梯度归因(saliency, integrated gradients)交叉验证——梯度直接指向「影响输出的输入」,比 attention 更接近解释。

第三,做反事实测试——如果你认为模型用 “cat” 来理解 “It”,试着把 “cat” 换成 “dog”,看预测有没有变化。

这些建议在第 52 篇 attention 解释性那里会更详细展开。这里只先埋一个伏笔。

5.5 一个略显反直觉的事实

研究者用 probing 实验发现:BERT 各层学到的语言学知识非常丰富——可以从中读出词性、依存树、共指关系、甚至一些语义关系。

但这不意味着模型「在做」语言学推理。

更准确的说法是:模型在大语料上学到了一种统计上等价于语言学知识的中间表示——它在 attention pattern 里编码了类似语法树的结构,但不一定是显式地、可解释地使用它。

这种区别在哲学上很重要:模型「会用」语言≠模型「理解」语言。

到 2026 年,这个争论仍在进行——大语言模型的「理解」边界仍是一个开放问题。

self-attention 是这场争论的物理基础——它是模型「看」上下文的眼睛——但它不是争论本身的答案。


六、self-attention 与传统模型的横向对比

维度 RNN/LSTM CNN Self-Attention
路径长度(任意两 token) O(N) O(log N)(堆叠空洞 conv) O(1)
单层计算复杂度 O(N · d²) O(N · k · d²) O(N² · d)
是否可并行训练 沿时间维不可并行 可并行 可并行
是否需要位置编码 否(顺序自带位置) 否(卷积带局部位置)
是否能看到全局 难(衰减) 难(受 receptive field 限制) 容易(全连接)
长序列代价 顺序训练慢,但显存友好 中等 显存爆炸 O(N²)

每种架构有自己的强项。

self-attention 在「长依赖 + 大批训练」场景下的综合表现最优——这是 Transformer 取代 RNN/CNN 的根本原因。

但在小数据 + 短序列 + 强先验场景下,CNN/RNN 仍有优势。


七、self-attention 在不同领域的几种用法

7.1 NLP

文本编码(BERT 类)、文本生成(GPT 类)、机器翻译(T5 类)。

到 2026 年,几乎所有大语言模型都基于 self-attention(虽然加了大量优化:GQA、RoPE、SwiGLU、FlashAttention 等)。

7.2 视觉(Vision Transformer)

ViT(Dosovitskiy 2020)把图像切成 16×16 patches,把每个 patch 当 token,用 self-attention 处理。

完全抛弃 CNN——在大数据预训练下能匹敌甚至超过 ResNet。

7.3 多模态(CLIP、Flamingo、GPT-4o)

视觉 + 文本通过 self-attention 和 cross-attention 联合建模。

self-attention 用在每个模态内部,cross-attention 用在模态之间。

7.4 语音(Conformer、Whisper)

语音 frames 经过 self-attention 处理,建模长时依赖。

Whisper(OpenAI 2022)是典型例子,用纯 Transformer 做语音识别。

7.5 生物(AlphaFold 2)

蛋白质氨基酸序列经过 self-attention(在 AlphaFold 里叫 evoformer),预测 3D 结构。

这是 self-attention 在序列数据上的非语言应用最有名的例子。

7.6 图(Graph Transformer)

图节点也可以当 token,self-attention 在节点间建立全连接的「软」消息传递。

GAT 是早期版本(基于 additive),后来 Graph Transformer 用 scaled dot-product。

self-attention 的通用性是它统治了这么多领域的根本原因——只要数据是「一组带向量表示的元素」,self-attention 就能用。

7.7 强化学习(Decision Transformer)

把状态 / 动作 / 回报序列当 token,用 self-attention 建模。

Decision Transformer(Chen 2021)和 Trajectory Transformer(Janner 2021)开启了「序列建模视角看 RL」的潮流。

到 2026 年,offline RL 的 SOTA 多数是 Transformer-based。

7.8 时间序列预测

时间序列也是「一串向量」,自然适合 self-attention。

但这里有个微妙问题:时间序列的位置感非常重要(昨天 vs 一年前),所以位置编码的设计比 NLP 还关键。

PatchTST(Nie 2022)、iTransformer(Liu 2023)等专门针对时间序列优化的 Transformer 在 2024-2025 已经成主流。

7.9 通用「集合」处理

任何「集合中的元素互相 attend」的场景都适合 self-attention:

这种通用性让 self-attention 成为深度学习的「瑞士军刀」。


七之续、self-attention 在生产系统中的几个观察

7.10 不只是模型架构问题

self-attention 在生产系统中的瓶颈,常常不是「模型怎么算」,而是:

vLLM、TensorRT-LLM、SGLang 等推理引擎在 2024-2025 的核心创新都围绕这些问题。

到第二十二篇 KV cache 那里会触及一些。

7.11 不同模型对位置感的敏感度

经验观察:

这件事在做迁移学习、长度泛化时要特别小心——位置编码方案不同的模型不能直接互换。


八、几个工程上的细节

8.1 self-attention 的输入输出形状

输入:(B, L, d_model)。

经过 W_Q/W_K/W_V 投影、attention、W_O 拼接:

输出:(B, L, d_model),与输入完全一样。

这种「形状不变」让 self-attention 可以无限堆叠——每一层的输入都是前一层的输出。

8.2 Pre-norm vs Post-norm

Vaswani 原版 Transformer 用 post-norm:x + Attn(LN(x))LN(...)

后来发现 pre-norm(x + Attn(LN(x)),不加最后的 LN)训练更稳定,几乎所有现代 LLM 都用 pre-norm。

到第二十一篇 LayerNorm 那里会专门讨论这件事。

8.3 Causal mask 的实现

decoder-only 模型必须保证「当前 token 只能 attend 到自己和左侧」。

实现方式:

mask = torch.triu(torch.ones(L, L), diagonal=1).bool()
# mask[i, j] = True 当 j > i(即未来 token)

scores = Q @ K.transpose(-2, -1) / sqrt(d_k)
scores = scores.masked_fill(mask, float('-inf'))
alpha = softmax(scores, dim=-1)

未来位置的 score 被设为 -∞,softmax 后权重为 0。

到第十七篇 causal mask 会讲它的形状、为什么不能省、推理时的特殊处理(KV cache)等。

8.4 Padding mask

实际 batch 训练时,不同样本长度不同,要 pad 到统一长度。

pad token 的位置不应该被 attend——用 padding mask 把那些位置的 score 设为 -∞。

scores.masked_fill(padding_mask, -inf)

8.5 Dropout

self-attention 里有两处常见 dropout:

第一,softmax 后对 α 做 dropout(attention dropout);

第二,输出投影 W_O 后做 dropout(output dropout)。

Vaswani 原版两者都用,p=0.1。

现代 LLM 在大数据预训练时常常去掉 dropout——数据量足够大,正则化主要靠数据多样性。


九、深入:self-attention 的几何视角

9.1 投影到子空间

W_Q、W_K、W_V 都是 d_model × d_k 的矩阵(d_k < d_model)。

它们把每个 d_model 维的 token 投影到 d_k 维的子空间。

不同的 head 投到不同的子空间——h 个 head 加起来覆盖了 d_model 维空间的 h 个不同「视角」。

每个 head 的 attention 只在自己的子空间里做相似度比较——所以「不同 head 关注不同 pattern」其实是「在不同子空间里看问题」。

9.2 query 与 key 的「双射」结构

对每对 (i, j),q_i · k_j 决定 token j 对 token i 的重要度。

这是一个有向的关系——q_i · k_j 不必等于 q_j · k_i。

非对称性很关键:「我看你」和「你看我」可以用完全不同的权重。

举例:「The cat sat」中 “sat” 强烈 attend “cat”(找主语);但 “cat” 不一定强烈 attend “sat”(“cat” 当 query 时可能在找其它信息,比如冠词或修饰语)。

这种非对称性在 RNN 和 CNN 里都不存在——self-attention 是第一个原生支持「方向性关系」的通用架构。

9.3 Value 的语义

V 投影后的 v_j 是「token j 能贡献的信息」。

在 attention 公式 Σ α v 里,v 是真正被传递的内容。

这也意味着:模型对 token j 的「内容编码」隐藏在 W_V 里——而不是在 K 里。

K 决定「我能被什么 query 找到」(语义 hash),V 决定「找到我之后我能给出什么」(语义 payload)。

这种区分让模型能学到「按某个特征找你 + 给出另一个特征的信息」的复杂模式——索引和内容的双轴解耦

9.4 Self-Attention 的「平滑」效应

每次 self-attention 把每个 token 的输出变成「所有 token 的加权平均」——这是一种平滑操作。

如果不加 W_O 输出投影、不加 FFN、不加残差,单纯堆叠 self-attention 会让所有 token 的表示越来越相似——最终塌缩到一个均匀向量。

这叫做 oversmoothing,是 Graph Transformer、深 Transformer 里的真实问题。

防止 oversmoothing 的关键:

第一,残差连接——让原始 token 表示绕过 attention 直达下一层;

第二,FFN 提供非线性区分——把 attention 平滑过的输出再扩展回不同方向;

第三,多 head + 不同 W_V——不同 head 学不同 V 投影,避免所有 head 都做相同平滑。

到第二十一篇 LayerNorm 那里会把这些组件协同工作的方式讲清楚。


十、self-attention 的反向传播

10.1 公式上的梯度

o = α v 关于 v 的梯度是 α(线性,简单)。

o = α v 关于 α 的梯度是 v(线性,简单)。

α = softmax(s) 关于 s 的梯度是 softmax 的雅可比,每个 (i, j) 元素是 α_i (δ_{ij} − α_j)。

s = q · k / √d_k 关于 q 的梯度是 k / √d_k。

s = q · k / √d_k 关于 k 的梯度是 q / √d_k。

把这些链起来,就是 attention 的反向传播——每一步都是简单的矩阵乘,所以 GPU 上极快。

10.2 为什么 √d_k 出现在反向传播里

注意 √d_k 的位置:它在前向时除一次,在反向时也会出现一次(链式法则)。

如果 d_k 增大但忘了除 √d_k,前向 score 数值变大、softmax 饱和;同时反向梯度也会爆炸(因为 q·k 直接进入梯度链)。

√d_k 同时控制前向数值范围和反向梯度幅度——它不仅是「让 softmax 不饱和」的技巧,也是数值稳定的关键。

10.3 attention 的梯度饱和问题

当某个 α 趋近 1 而其它趋近 0 时,softmax 的雅可比变得接近零——因为 α (1−α) 在 α=0 或 α=1 时都是 0。

这意味着:attention 越 sharp,越难继续学

模型一旦学到一个尖锐的 attention pattern,再想调整就需要更大的梯度,但梯度本身也变小了——形成一种「自锁」。

这是为什么 attention 容易陷入早期学到的 pattern 里跳不出来——也是 attention head 在不同初始化下经常学到不同 pattern 的原因(路径依赖)。


十一、self-attention 的几个实战技巧

11.1 Warmup 学习率

attention 训练初期对学习率非常敏感。

太大:W_Q/W_K/W_V 没学好就被推到极端,attention 塌缩。

太小:训练慢,可能 stuck 在均匀 attention 阶段。

实践经验:用 warmup(线性从 0 升到 peak lr,前 1k–10k 步),再衰减。

Vaswani 原版用了一个特殊的「Noam schedule」:lr ∝ d_model^{-0.5} · min(step^{-0.5}, step · warmup^{-1.5})。

现代 LLM 多数用线性 warmup + cosine decay,效果接近且更易理解。

11.2 Gradient clipping

attention 在某些 batch 上会产生异常大的梯度(特别是 softmax 进入饱和区时)。

Gradient clipping(按全局范数)把梯度幅度上限固定(典型值 1.0)。

不加 clipping,遇到 outlier batch 会让模型「炸掉」一次,整段训练被毁。

11.3 Mixed precision 训练

attention 的 QKᵀ 和 softmax 用 fp16/bf16 训练时,要小心数值溢出。

实际工程上:

FlashAttention 内部就是这么做的。

11.4 长上下文的「衰减」问题

当序列长度 L 很大(比如 32k+)时,softmax 的「赢者通吃」会让 attention 几乎只关注少数 token。

许多 token 的有效信息进不去 attention(α ≈ 0)。

解决方案多种多样:

每种方案都有自己的 tradeoff,长上下文专题(第三十二篇)会专门讲。

11.5 KV cache:推理时的关键优化

decoder-only 模型推理时,每生成一个新 token,前面所有 token 的 K 和 V 都已经算过了。

KV cache 把这些 K、V 缓存下来,新 token 只需要算自己的 K/V,跟前面的 cache 拼起来做 attention。

时间复杂度从每步 O(N²) 降到 O(N)(N 是当前已生成长度)。

代价:显存占用 O(N · d_model · num_layers · 2)(K 和 V 各一份)。

对长上下文推理,KV cache 占用极大——这就是 GQA、MQA 出现的原因。

KV cache 专题(第二十二篇)会讲实现、压缩、共享等技术。


十二、self-attention 的「失败模式」

12.1 Attention sink

某些 token(通常是 BOS / 特殊 token)会变成「吸引子」——所有 query 都在它身上分配较多权重。

这件事到 2023 年的 StreamingLLM 工作(Xiao 2023)里被系统化研究:模型推理时如果丢掉前几个 token 的 KV cache,性能会断崖式下降——因为那些 token 是 attention sink。

实践含义:前几个 token 的 KV 不能丢——即使你想截断长上下文。

12.2 Repetition under sharpening

模型在生成时,如果 self-attention 学到的 pattern 让某个 token 反复 attend 到「上一个相同 token」,会陷入重复生成。

GPT-2/3 时代的「重复 bug」很多源自这种自激振荡。

解决方案:repetition penalty、top-k/top-p 采样、frequency penalty 等。

12.3 Context dilution

当上下文非常长(比如 100k token),当前 query 想 attend 到中间某个具体 token,但权重被均匀稀释——所有 token 都拿到一点点 α,没有一个尖锐的 peak。

这是 long-context 模型的真实问题,到 2024 年底 NIAH(Needle In A Haystack)测试成为标准评估,很多模型在 50k+ context 上出现 attention 稀释。

12.4 Attention head pruning 的代价

研究发现很多 head 是「冗余」的——直接 prune(剪掉)几个 head,性能几乎不变。

哪几个 head 重要是任务相关的——同一个模型在 NLI 任务里关键的 head,可能在翻译任务里不重要。

不能简单 prune 一组 head 就在所有任务上用。

12.5 Lost-in-the-middle

Liu 等 2023 年的工作发现:长上下文里,模型对「中间位置」的信息检索能力显著低于「开头」和「结尾」。

直觉解释:开头有 attention sink 效应、结尾有 recency bias,中间被夹击。

这件事到 2026 年仍是 long-context 评估的重要维度——不能只看 NIAH 的端点测试,要在中间多个位置都测。

12.6 Attention 在分布外(OOD)输入上的退化

模型训练时见过的 attention pattern 是有限的。

遇到完全分布外的输入(比如训练时只见过英文的模型,输入一段藏文)时,attention 可能完全塌缩到无意义的位置。

这是 zero-shot 推理失败的常见原因之一——不仅是 token embedding 不熟悉,attention pattern 也找不到对的「靶子」。

12.7 Modality leakage

多模态模型里,文本 token 和图像 token 在同一 self-attention 内交互。

如果训练数据里某个模态主导(比如文本远多于图像),attention 会偏向多数模态,少数模态的信号被淹没。

到 2025 年的多模态 SFT 调优,「模态平衡」已经成为重要训练策略——单纯堆数据不够。


十三、关键概念回顾

走到这里,self-attention 该建立的几个核心直觉是:

self-attention 是 cross-attention 的特例,Q/K/V 都来自同一序列。

任意两 token 之间路径长度 = 1——这是 self-attention 取代 RNN 的核心理由。

self-attention 是 permutation-equivariant——对位置完全无知,必须靠位置编码注入位置信号。

multi-head 让不同 head 学不同的关系类型——共指、依存、局部上下文等。

attention ≠ 解释——α 是中间计算量,不是因果解释,看图要谨慎。

O(1) 不是免费的——代价是 O(N²) 计算和显存,引出了 FlashAttention、Linear Attention 等后续工作。


十四、常见误解

10.1 self-attention「让每个 token 看自己最像」

不对。

只有当 W_Q ≈ W_K(恒等附近)时才会出现「自己最像自己」。

学过的 W_Q ≠ W_K 几乎总是把对角线打散,让 attention 跨位置 sharpen。

10.2 self-attention 让模型「记住整个上下文」

不准确。

self-attention 在当前层让每个 token 看到所有 token 一次。

但「记住」需要多层堆叠 + 残差 + FFN 才能稳定保留——单层 self-attention 只是一次「混合」,不是「记忆」。

10.3 加了位置编码就「彻底解决」位置问题

不完全。

位置编码只是「注入」位置信号,但模型怎么用它仍要学。

不同位置编码方案(绝对/相对/RoPE/ALiBi)在外推性、长度泛化等方面差异巨大——这件事远没有「加个 PE 就完事」那么简单。

10.4 self-attention 和 cross-attention 是两种不同机制

它们是同一个机制,只是 Q 和 K/V 的来源不同。

把 self-attention 叫做「Q/K/V 同源的 attention」更准确。

10.5 self-attention 比 RNN 永远更好

不对。

在小数据 + 短序列 + 流式推理(streaming inference)场景,RNN/LSTM 仍然有优势——并非每个任务都要 self-attention。

但在「大模型 + 大数据 + 长依赖」场景,self-attention 几乎不可替代。

10.6 N² 是 self-attention 的根本上限

不是。

「在精确计算 attention」的前提下是 O(N²),但 FlashAttention 已经把常数大幅压低;Linear/Sparse Attention 把 N² 降成 N · log N 或 N。

到 2026 年长上下文模型(百万 token)已经实用化——N² 不再是不可逾越的墙。

10.7 self-attention 输出维度等于序列长度

不对。

输出形状是 (B, L, d_model),与输入相同。

L 是序列长度,d_model 是嵌入维度——两者是不同维度。

10.8 self-attention 没有归纳偏置(inductive bias)

不严格。

self-attention 确实没有 RNN 的「顺序性」、CNN 的「局部性」这些显式偏置。

但它有自己的偏置:permutation-equivariance(位置无关性)。

这个偏置看似中立,实际上在「数据量充足时」是优势,在「数据量稀缺时」是劣势——CNN 因为有局部性偏置,小数据上就能学到合理表示,self-attention 需要大数据才能补足这个偏置缺失。

这就是为什么 ViT 在小数据集(CIFAR、ImageNet-1k)上不一定打得过 CNN,但在大数据集(JFT-300M、ImageNet-21k)上反超。

10.9 self-attention 的「全连接」是坏事

不一定。

「每个 token attend 所有 token」看似是一种过度连接(overconnection)。

但配合 softmax 归一化 + 学习到的稀疏 pattern,模型实际上只在需要的位置 sharpen attention。

「全连接 + 学习选择」比「显式稀疏 + 没学习」更通用——这是 self-attention 在数据充足时的优势。

代价是 O(N²) 的最坏情况复杂度。

10.10 self-attention 等价于一种 graph neural network

接近正确但不完全。

self-attention 可以被看作是「全连接图上的 GAT」——每个 token 是一个节点,节点间全连接,attention weight 是边权。

但 self-attention 多了「位置编码」这个非图特性——纯 GAT 在无序图上工作,self-attention 必须有顺序信号。

这两者在「集合数据」(无序)上是等价的,在「序列数据」(有序)上有本质区别。


十五、下一步

下一篇 15|Scaled Dot-Product 会回答一个反复被提的问题:「为什么是除以 √d_k 不是 d_k?」

我们会用方差推导讲清楚:当 q、k 是 d_k 维独立同分布零均值单位方差向量时,q · k 的方差等于 d_k;除以 √d_k 把方差归一回 1,让 softmax 不饱和。

再往后:第 16 篇讲 multi-head,把今天提到的「不同 head 学不同 pattern」彻底拆开;第 17 篇讲 causal mask 和 KV cache;第 26 到 30 篇集中讲位置编码(绝对/Sinusoidal/Learned/Relative/RoPE/ALiBi)。

第 52 篇会回到这一篇埋的伏笔:attention 解释性问题,从 Jain & Wallace 到现代归因方法的全貌。

如果你想立刻动手验证 self-attention 的 permutation-equivariant 性质:拿第 13 篇的 SelfAttention 代码,构造一个长度 5 的输入 X,把 X 沿序列维度做一次 permute,比较两次输出(permute 前 vs permute 后再 permute 回去)——你会发现它们完全一致。

然后给 X 加一个固定 PE:再做一次同样的 permute 实验——你会发现这次输出不一致了。

这两个简单实验能让你对「位置编码奏效的机制」有直接的体感。


十六之一、补充:self-attention 的训练成本与 scaling laws

self-attention 的复杂度是 O(N² · d),而 FFN 是 O(N · d²)。

当 N 远小于 d(比如 N=512, d=4096)时,FFN 主导计算成本——attention 只占总 FLOPs 的 10-20%。

当 N 接近或超过 d 时,attention 主导——这是长上下文模型显著变贵的原因。

Chinchilla scaling laws(Hoffmann 2022)告诉我们:在固定 FLOPs 预算下,参数量 N_p 和数据量 D 应该按 N_p ∝ D 同步增长。

但这条定律默认 N(序列长度)固定——一旦序列长度也增长,attention 那部分会改变最优配比。

到 2026 年,长上下文 + scaling laws 的交互仍在被研究——这是一个开放问题。

16.1.1 一个工程上的对照表

短上下文(N ≤ 1k,d ≥ 4k):FFN 占主导,attention 加速性价比低。

中上下文(N ≈ d):两者相当,FlashAttention、KV cache 都开始有意义。

长上下文(N ≥ 4d):attention 主导,必须考虑 sparse / linear / paged attention。

这个表只是粗略经验,但能帮你判断「优化哪一块」最值得。

16.1.2 一个常被忽略的细节:activation memory

attention 的 softmax(QKᵀ) 矩阵本身就是 O(N²)——这是「显存瓶颈」而不是「FLOPs 瓶颈」。

FlashAttention 解决的正是这个:把这个 O(N²) 中间矩阵从 HBM 里去掉,改在 SRAM 里 tile-by-tile 计算

理解这一点你才能理解为什么 FlashAttention 对训练显存影响如此显著——它优化的是「中间存储」而不是「计算量」。


十六、参考文献

下面按相关度排序列出本篇直接引用与延伸阅读,每条附一句话提示其在本篇中的角色。

阅读建议:先看摘要、再回到本文找对应段落、再决定要不要精读全文。这种「问题驱动」的阅读比「从头读到尾」效率高得多。

  1. Vaswani, A. et al. “Attention Is All You Need.” NeurIPS 2017. self-attention 的源头论文,Fig 3 共指 head 可视化。
  2. Jain, S., Wallace, B. C. “Attention is not Explanation.” NAACL 2019. attention ≠ 解释 的经典论证。
  3. Wiegreffe, S., Pinter, Y. “Attention is not not Explanation.” EMNLP 2019. 与 2 配对的反驳。
  4. Clark, K., Khandelwal, U., Levy, O., Manning, C. D. “What Does BERT Look At? An Analysis of BERT’s Attention.” BlackBoxNLP 2019. BERT attention head 的实证分析。
  5. Voita, E. et al. “Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting.” ACL 2019. 不同 head 学不同 pattern 的实证支持。
  6. Vig, J. “A Multiscale Visualization of Attention in the Transformer Model.” ACL 2019 demo. BertViz 工具。
  7. Dosovitskiy, A. et al. “An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale.” ICLR 2021. ViT,self-attention 在视觉的应用。
  8. Jumper, J. et al. “Highly Accurate Protein Structure Prediction with AlphaFold.” Nature 2021. AlphaFold 2,self-attention 的非语言里程碑。
  9. Veličković, P. et al. “Graph Attention Networks.” ICLR 2018. self-attention 的图变体。
  10. Su, J. et al. “RoFormer: Enhanced Transformer with Rotary Position Embedding.” Neurocomputing 2024. RoPE,位置编码的现代主流。
  11. Press, O., Smith, N. A., Lewis, M. “Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation.” ICLR 2022. ALiBi。
  12. Shaw, P., Uszkoreit, J., Vaswani, A. “Self-Attention with Relative Position Representations.” NAACL 2018. Relative PE。
  13. Sutskever, I., Vinyals, O., Le, Q. V. “Sequence to Sequence Learning with Neural Networks.” NeurIPS 2014. RNN-based seq2seq 的代表,对照 self-attention 的「O(N) 跳数」问题。
  14. Hochreiter, S. “Untersuchungen zu dynamischen neuronalen Netzen.” Diploma thesis, TU München, 1991. 梯度消失问题的源头。
  15. Dao, T. et al. “FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness.” NeurIPS 2022. 对 O(N²) 显存问题的工程突破。
  16. Liu, N. F. et al. “Lost in the Middle: How Language Models Use Long Contexts.” TACL 2024. self-attention 在长上下文中位置偏置的实证研究。
  17. Xiao, G. et al. “Efficient Streaming Language Models with Attention Sinks.” ICLR 2024. attention sink 现象与 self-attention 的边界 token 异常。
  18. Bronstein, M. M. et al. “Geometric Deep Learning: Grids, Groups, Graphs, Geodesics, and Gauges.” arXiv 2104.13478. 用对称性理解架构归纳偏置的统一视角。
  19. Ainslie, J. et al. “GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints.” EMNLP 2023. self-attention 的 KV 共享变体,工程优化方向。
  20. Tay, Y. et al. “Efficient Transformers: A Survey.” ACM Computing Surveys 2022. self-attention 各种稀疏/线性变体的综述。

每篇都建议至少读摘要,时间充裕则精读 1, 2, 4, 8, 15——这五篇构成本系列对 self-attention 的核心知识地图。


← 上一篇:13|Q/K/V 三件套 | 下一篇:15|Scaled Dot-Product

同主题继续阅读

把当前热点继续串成多页阅读,而不是停在单篇消费。

2026-04-15 · transformer

【Transformer 与注意力机制】01|为什么要从这里开始

这是【Transformer 与注意力机制】系列的第一篇,承担两件事:一是把这套五十多篇文章为谁写、解决什么问题、彼此之间是什么关系交代清楚;二是为完全没基础的读者画出一条从向量、点积、矩阵乘法走到自注意力、再走到大语言模型的爬升路径,让你在投入时间之前先知道终点在哪、路上要经过哪些坎、读完之后你会、还不会做什么事。

2026-04-15 · transformer

【Transformer 与注意力机制】系列总览

从《Attention Is All You Need》出发,把注意力机制、Transformer 架构、训练范式、模型变体、推理工程、可解释性与未来架构串成一条 58 篇的深度博客线。


By .