第9章：Transformer：动态拓扑的注意力革命

每一个 Attention Head 都在问:此刻,哪些部分对哪些部分重要?

---

一、2017年6月的一篇论文

2017年6月12日,Google Brain的Ashish Vaswani等八位作者在arXiv上传了一篇论文,标题是”Attention Is All You Need”。

这个标题很狂妄。当时机器翻译的主流架构是RNN(循环神经网络)和LSTM(长短期记忆网络)——它们统治NLP领域已经超过二十年。而这篇论文说:你们都不需要了,注意力机制就够了。

更狂妄的是,他们是对的。

五个月后,这篇论文被NIPS 2017接收。两年后,BERT、GPT-2横空出世,全部基于Transformer架构。再过三年,GPT-3证明了Transformer可以扩展到1750亿参数。

今天,几乎所有的大型语言模型——从ChatGPT到Claude——都是Transformer的变体。

但在2017年6月,这只是一个激进的想法:能否完全抛弃循环结构?

1.1 顺序依赖的代价

RNN的工作方式很直观:从左到右逐个处理单词,每一步都把当前单词和之前的”记忆”结合起来,更新隐状态。

$$h_t=f(h_{t-1},x_t)$$

这个设计看起来很自然——毕竟人类也是从左到右阅读句子的。但它有两个致命缺陷:

缺陷1:无法并行化

要计算第10个词的隐状态,你必须先计算第9个词的隐状态;要计算第9个词的,必须先计算第8个词的……这个链条无法并行化。

在GPU时代,这是灾难性的。GPU有数千个核心,可以同时计算数千个操作。但RNN强制你串行计算——一次只能处理一个词。

对于长度为N的序列,RNN需要N步串行计算。即使每步只需要1毫秒,1000个词的句子就需要1秒。

缺陷2:长程依赖的信息衰减

假设你要翻译这个句子:

"The cat, which we found in the garden last summer, was very friendly."

主语”cat”和谓语”was”之间隔了12个词。RNN需要把”cat”的信息通过12步传递到”was”。

每一步传递,信息都会衰减。即使有LSTM的门控机制,长距离传递仍然困难——这是梯度消失问题的另一种表现。

Bengio等人在1994年就证明了:RNN在学习长程依赖时,梯度会指数级衰减。LSTM缓解了这个问题,但没有根本解决。

1.2 Vaswani的赌注

Vaswani等人的想法是:如果每个词都能直接”看到”其他所有词,就不需要通过隐状态传递信息了。

这就是Self-Attention(自注意力)机制。它让序列中的每个位置,都能直接计算与其他所有位置的相关性,然后用这个相关性加权求和。

没有循环,没有顺序依赖,完全可以并行计算。

代价是什么?计算复杂度从 $O(N)$ 变成 $O(N^2)$。

对于长度为N的序列,Self-Attention需要计算 $N\times N$ 的注意力矩阵——每个词对每个词的相关性。当N=512时,这是262,144个数。当N=2048时,这是4,194,304个数。

这是一个大胆的赌注:用 $O(N^2)$ 的空间和计算,换取并行化和全局感受野。

结果证明,这个赌注是值得的。

---

停顿一下

"结果证明这个赌注是值得的"——但值得，凭什么？

Transformer能工作，是因为Self-Attention这个设计在结构上比RNN更接近语言的"真实结构"？还是仅仅因为 $O(N^2)$ 的算力更容易被GPU的并行性利用，让它能用更多的数据和计算量来弥补架构缺陷？

换句话说：Transformer的成功，是因为它的归纳偏置正确，还是因为它的归纳偏置最弱——弱到可以用蛮力数据来填充？

"没有归纳偏置"本身是一种归纳偏置。Transformer假设了"任意位置的词都可能对任意其他位置重要"，这个假设对语言来说是对的吗？对所有任务来说是对的吗？

先把这个问题放着。

---

二、Self-Attention的数学:一个优雅的对称性

Transformer的核心是Self-Attention机制。它要解决的问题很简单:

给定一个序列,如何让每个位置都能”看到”其他所有位置,并决定哪些位置重要?

数学形式出奇地简洁。对于输入序列 $X\in {\mathbb{R}}^{N\times d}$ (N个词,每个词d维),计算三个矩阵:

$$Q=X{W}_Q,K=X{W}_K,V=X{W}_V$$

这三个矩阵有不同的语义角色:

• Query(查询):”我在找什么?”——当前位置想要关注的模式

• Key(键):”我能提供什么?”——每个位置提供的索引信息

• Value(值):”我的内容是什么?”——每个位置的实际信息

然后计算注意力:

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

让我拆解这个公式的每一步:

步骤1:计算相似度 $Q{K}^T\in {\mathbb{R}}^{N\times N}$

每个query和每个key做点积。第 $i$ 行第 $j$ 列的元素是 $q_i\cdot k_j$ ——表示位置 $i$ 对位置 $j$ 的”兴趣程度”。

点积为什么能衡量相似度?因为 $q\cdot k=|q||k|\cos \theta$ ——当两个向量方向相同时,点积最大。

步骤2:缩放 除以 $\sqrt{d_k}$

这是一个关键的技术细节。当 $d_k$ 很大时,点积的方差会很大——可能出现非常大的正值或负值。这会导致softmax饱和:最大值对应的概率接近1,其他位置接近0。

在初始化假设下， $Q$ 与 $K$ 的各维独立且方差为常数，因此点积 $Q\cdot K$ 的方差与维度 $d$ 成正比，采用 $\frac{1}{\sqrt{d_k}}$ 进行缩放可以将其方差归一化到 $O(1)$ 的稳定范围，从而避免 $softmax$ 进入饱和或过于平坦的区域。在训练过程中，尽管表示分布会发生偏移，但通过层归一化对输入分布的重整，以及残差连接对梯度流的稳定作用，模型能够自适应地维持attention score的尺度在合理范围内，从而保证训练稳定性。

步骤3:归一化 softmax按行归一化

将相似度转化为概率分布。第 $i$ 行的softmax输出是一个概率向量,和为1,表示位置 $i$ 对所有其他位置的注意力分配。

步骤4:加权求和 乘以 $V$

用注意力权重对value加权求和。位置 $i$ 的输出是:

$${\text{out}}_i=\sum _{j=1}^N{\text{attention}}_{ij}\cdot v_j$$

这是一个软寻址(soft addressing)机制:不是硬性选择某个位置,而是对所有位置加权平均,权重由相似度决定。

这个机制的美妙之处:

1. 全局感受野:每个位置都能直接”看到”所有其他位置,不需要通过隐状态传递

2. 完全并行:所有位置的输出可以同时计算,没有顺序依赖

3. 动态拓扑:注意力矩阵是根据输入内容动态生成的,不是固定的权重

---

三、自己动手:看注意力在”看”什么

拿出一张纸,手工计算一个简单例子。

输入句子:“cat sat mat”(3个词,为了简化)

假设每个词已经嵌入为2维向量:

cat: [1, 0]
sat: [0, 1]
mat: [1, 1]

步骤1:计算Q, K, V

简化起见,假设 $W_Q=W_K=W_V=I$ (单位矩阵),所以 $Q=K=V=X$ 。

步骤2:计算注意力分数

$$Q{K}^T=\left[\begin{array}{cc}1& 0\\ 0& 1\\ 1& 1\end{array}\right]\left[\begin{array}{ccc}1& 0& 1\\ 0& 1& 1\end{array}\right]=\left[\begin{array}{ccc}1& 0& 1\\ 0& 1& 1\\ 1& 1& 2\end{array}\right]$$

除以 $\sqrt{d_k}=\sqrt{2}\approx 1.41$:

$$\text{scores}=\left[\begin{array}{ccc}0.71& 0& 0.71\\ 0& 0.71& 0.71\\ 0.71& 0.71& 1.41\end{array}\right]$$

步骤3:应用softmax(按行)

对第一行$[0.71,0,0.71]$:

$$\text{softmax}=\frac{[e^{0.71},e^0,e^{0.71}]}{\sum }\approx [0.42,0.16,0.42]$$

完整的注意力矩阵(近似):

$$A=\left[\begin{array}{ccc}0.42& 0.16& 0.42\\ 0.16& 0.42& 0.42\\ 0.27& 0.27& 0.46\end{array}\right]$$

解读: - 第1行:“cat”对自己(0.42)和”mat”(0.42)的注意力相近,对”sat”较低(0.16) - 第3行:“mat”对所有词都有注意力,但对自己最高(0.46)

步骤4:加权求和

第1个词的输出:

$${\text{out}}_1=0.42\times [1,0]+0.16\times [0,1]+0.42\times [1,1]=[0.84,0.58]$$

这个向量融合了”cat”和”mat”的信息,因为注意力权重高。

关键观察:注意力矩阵 $A$ 是动态生成的——它不是固定的权重,而是根据输入内容计算出来的。这就是”动态拓扑”的含义。

---

四、多头注意力:分工与协作

单个注意力头只能捕获一种关系模式。但语言中的关系是多样的:

• 句法关系:主语-谓语,修饰语-中心词

• 语义关系:代词-先行词,同义词

• 位置关系:相邻词,远距离依赖

Transformer用多头注意力(Multi-Head Attention)解决这个问题:并行运行多个注意力头,每个头学习不同的关系模式。

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

其中每个头:

$${\text{head}}_i=\text{Attention}(Q{W}_i^Q,K{W}_i^K,V{W}_i^V)$$

直觉:把 $d$ 维空间分成 $h$ 个子空间,每个子空间 $d/h$ 维。每个头在自己的子空间里做注意力,最后拼接起来。

实验观察(来自BERT的可视化研究):

• Head 1:关注相邻词(局部语法)
• Head 2:关注代词和先行词(指代消解)
• Head 3:关注句子开头和结尾(全局结构)
• Head 4:关注标点符号(句法边界)

不同的头自发地学会了不同的语言学功能,尽管没有显式监督。

---

五、O(N²)的代价:全局感受野的热力学账单

Self-Attention的时间复杂度是 $O(N^{2}d)$ ,空间复杂度是 $O(N^2)$ 。

这不是一个可以被优化掉的常数因子,这是架构的本质代价。

5.1 为什么是N²?

计算 $Q{K}^T$ 需要 $N\times N$ 次点积,每次点积需要 $d$ 次乘法。总计 $O(N^{2}d)$ 次操作。

存储注意力矩阵需要 $N\times N$ 个浮点数。当 $N=512$ 时,这是262,144个数。当 $N=2048$ 时,这是4,194,304个数。

对比其他架构:

• RNN: $O(N{d}^2)$ ——线性于序列长度,但顺序依赖

• CNN: $O(Nk{d}^2)$ ——线性于序列长度,但感受野受限于卷积核大小$k$

• Transformer: $O(N^{2}d)$ ——全局感受野,但平方复杂度

5.2 信息论视角:全连接的熵成本

从信息论角度看, $O(N^2)$ 复杂度是全连接拓扑的必然结果。

RNN是链式拓扑:每个节点只连接前一个节点。信息流是单向的,从左到右。这是一个低熵结构——连接模式是确定的,不需要存储。

Transformer是全连接拓扑:每个节点连接所有其他节点。这是一个高熵结构——连接权重(注意力矩阵)是动态的,必须存储。

Shannon的信息论告诉我们:描述一个 $N$ 节点全连接图的连接权重,需要 $\mathrm{\Omega }(N^2)$ 比特信息——因为有 $N^2$ 条边。

这不是实现问题,这是信息论下界。

5.3 物理学视角:Landauer原理的约束

Landauer原理说:擦除1比特信息至少需要 $k_{B}T\ln 2$ 的能量。

Transformer在每个前向传播中,都要计算并存储 $N^2$ 个注意力权重。在下一个token时,这些权重被”擦除”(重新计算)。

这意味着:处理长度为 $N$ 的序列,Transformer至少需要 $O(N^2{k}_{B}T\ln 2)$ 的能量——这是热力学下界。

当 $N=2048$ ,室温 $T=300K$ 时,这大约是 ${10}^{-18}$ 焦耳。看起来很小,但当你处理数十亿个token时,这会累积成显著的能量消耗。

GPT-3训练消耗了约1287 MWh的电力。其中相当一部分,是用于计算和存储注意力矩阵。

5.4 这个代价值得吗?

对于长度 $<2048$ 的序列,答案是肯定的。Transformer在机器翻译、文本生成、问答等任务上全面超越RNN。

但对于长序列(文档级翻译、长文本摘要、基因序列分析), $O(N^2)$ 成为瓶颈。

这催生了一系列”高效Transformer”的研究:

• Sparse Attention(稀疏注意力):只计算部分注意力,如局部窗口+全局token

• Linear Attention(线性注意力):用核技巧将复杂度降到 $O(N{d}^2)$

• Reformer:用局部敏感哈希近似注意力

• Longformer:混合局部窗口和全局注意力

但所有这些方法都有代价:要么牺牲全局感受野,要么引入近似误差。

没有免费的午餐:全局感受野和线性复杂度,你只能选一个。这不是工程限制,这是信息论和热力学的基本约束。

---

六、线性注意力：用加法结合律换掉 softmax

$O(N^2)$ 的根源在于 $Q{K}^T$ ——这个矩阵必须完整算出来，才能对每一行做 softmax 归一化。

那么，如果去掉 softmax 呢？

---

8.1 核函数视角

回忆标准注意力：

$${\text{out}}_i=\frac{\sum _j\exp (q_i\cdot k_j)\cdot v_j}{\sum _j\exp (q_i\cdot k_j)}$$

softmax 里的 $\exp (q\cdot k)$ 可以看作一个核函数 $\kappa (q,k)$ ，它衡量 query 和 key 的相似度。

核函数有一个普遍的近似方案：特征映射。如果存在函数 $\varphi (\cdot )$，使得：

$$\kappa (q,k)\approx \varphi (q{)}^T\varphi (k)$$

那么注意力就可以改写为：

$${\text{out}}_i=\frac{\sum _j\varphi (q_i{)}^T\varphi (k_j)\cdot v_j}{\sum _j\varphi (q_i{)}^T\varphi (k_j)}$$

关键：利用加法结合律和矩阵乘法结合律，把括号换一个顺序：

$${\text{out}}_i=\frac{\varphi (q_i{)}^T(\sum _j\varphi (k_j)v_j^T)}{\varphi (q_i{)}^T(\sum _j\varphi (k_j))}$$

现在分子分母里的 $\sum _j\varphi (k_j)v_j^T$ 和 $\sum _j\varphi (k_j)$，可以提前计算一次，存为两个固定矩阵。这两个矩阵的计算代价是 $O(N{d}^2)$ ——线性于序列长度。

所有 $N$ 个 query 共享这两个矩阵，每个 query 的计算是一次矩阵-向量乘法，代价 $O(d^2)$ 。

总复杂度： $O(N{d}^2)$ ，线性注意力诞生了。

---

8.2 线性注意力的致命缺陷：圆滑性

这听起来很完美。但线性注意力在实践中一直不如标准注意力——尤其是在需要精确、尖锐的注意力时（比如指代消解、关键词匹配）。

原因是：softmax天然地产生尖锐的分布，线性核不能。

softmax 的指数函数会大幅放大大的分数、压制小的分数。结果是注意力权重高度集中——某个位置可能拿走 0.95 的权重，其他位置只剩下 0.05。

线性核 $\varphi (q{)}^T\varphi (k)$ 没有这种”尖锐化”效果。所有位置的权重更加均匀地分布，大家都能说上话。

这看起来更”公平”，但恰恰是问题所在。

用一个比喻：你在一个嘈杂的派对上想听清一个人说话（信号），softmax 注意力会让你把 95% 的注意力集中在那个人身上，5% 分给其他人（噪声）。线性注意力会把注意力平均分给所有人——你能听到那个人，但噪声也被同等放大，信号被淹没了。

线性注意力太圆滑了——它把噪声和信号一视同仁地累加进来，稀释了尖锐度。

更形式化地说：线性注意力计算的是 Key-Value 外积矩阵 $\sum _j\varphi (k_j)v_j^T$ 的一次线性读取。这个矩阵把所有历史信息压进了一个固定大小的矩阵（ $d\times d$ 维），像一个有限容量的黑板。当序列变长时，新信息不断叠加，旧信息被覆写。

这正是线性注意力的本质：它实际上是一个有限容量的循环状态，而不是真正的全局感受野。

---

七、状态空间模型：从线性注意力到选择性记忆

线性注意力的分析揭示了一个更深的问题：全局感受野和线性复杂度，能否同时拥有？

答案是：不能——除非引入一个新的维度：选择性。

这就是状态空间模型（SSM）以及 Mamba 的核心思想。

---

9.1 从线性注意力看出 SSM 的影子

线性注意力的递推形式：

$$S_t=S_{t-1}+\varphi (k_t)v_t^T,{\text{out}}_t=\varphi (q_t{)}^T{S}_t$$

其中 $S_t\in {\mathbb{R}}^{d\times d}$ 是一个累积的”状态矩阵”。

这个递推式和经典的线性状态空间模型（Linear SSM）几乎完全一样：

$$h_t=A{h}_{t-1}+B{x}_t,y_t=C{h}_t$$

$h_t$ 是隐状态，$A$ 是状态转移矩阵，$B$ 是输入投影，$C$ 是输出投影。

线性注意力的 $S_t$ 就是 $A=I$（单位矩阵）时的 SSM——它只知道”累加”，不知道”遗忘”，所以旧信息永远留在状态里，造成覆写和稀释。

SSM 的改进：让 $A$ 不等于单位矩阵，而是一个可以选择性”衰减”旧信息的矩阵。

---

9.2 HiPPO：用正交多项式压缩历史

2020年，Albert Gu 等人提出了 HiPPO（High-order Polynomial Projection Operator）框架，给出了 $A$ 矩阵的理论构造：

$$A_{nk}=-\{\begin{array}{ll}(2n+1{)}^{1/2}(2k+1{)}^{1/2}& \text{if }n>k\\ n+1& \text{if }n=k\\ 0& \text{if }n<k\end{array}$$

直觉：HiPPO-$A$ 让隐状态 $h_t$ 成为历史输入 $x_{1:t}$ 在勒让德多项式基下的最优投影。就像用傅里叶级数的前几项近似一个函数——你用有限的系数，捕获了历史的”主要成分”。

这解决了线性注意力的覆写问题：旧信息不是被均等地累加，而是被按重要性压缩成一个低维的”历史摘要”。

---

9.3 Mamba：选择性状态空间模型

2023年，Albert Gu 和 Tri Dao 提出了 Mamba，核心创新是选择性机制（Selective SSM, S6）。

标准 SSM 的 $A,B,C$ 矩阵是固定的，与输入无关。这意味着模型以同样的方式处理每一个 token——无论这个 token 多么重要或无关紧要。

Mamba 让 $B,C$ 和步长 $\mathrm{\Delta }$ 依赖于输入：

$$B_t=B(x_t),C_t=C(x_t),{\mathrm{\Delta }}_t=\mathrm{\Delta }(x_t)$$

离散化后的状态转移：

$${\overline{A}}_t=\exp ({\mathrm{\Delta }}_{t}A),{\overline{B}}_t={\mathrm{\Delta }}_t{B}_t$$$$h_t={\overline{A}}_t{h}_{t-1}+{\overline{B}}_t{x}_t,y_t=C_t{h}_t$$

选择性的意义：

• ${\mathrm{\Delta }}_t$ 大：${\overline{A}}_t\approx 0$ ，状态大幅衰减——模型主动遗忘旧信息，关注当前输入
• ${\mathrm{\Delta }}_t$ 小：${\overline{A}}_t\approx I$ ，状态几乎不变——模型主动记忆旧信息，忽略当前输入

这不是固定的遗忘曲线，而是由内容驱动的选择性记忆。对于重要的 token，模型降低衰减率，让它的信息在状态中保留更久；对于噪声 token，模型加速遗忘。

---

9.4 为什么选择性解决了线性注意力的问题？

回到”圆滑性”的问题：线性注意力把所有信息均等地累加进状态，导致信号被噪声稀释。

Mamba 通过选择性 ${\mathrm{\Delta }}_t$ 解决了这个问题：

• 看到重要信息 → ${\mathrm{\Delta }}_t$ 大 → 状态剧烈更新，旧噪声被清除
• 看到噪声 → ${\mathrm{\Delta }}_t$ 小 → 状态几乎不变，噪声被忽视

这就是 softmax 的”尖锐性”在循环架构中的对应物：不是在空间维度上集中注意力，而是在时间维度上选择性记忆。

---

9.5 代价：局部性回来了

Mamba 在序列长度上是 $O(N)$，而且可以在训练时并行（通过并行扫描算法），在推理时循环（每步 $O(d)$ 常数时间）。

但它有一个代价：选择性是局部的。

${\mathrm{\Delta }}_t$ 只依赖于当前 token $x_t$，不能直接看到序列中其他位置的 token。这意味着 Mamba 不能像 Transformer 那样做”跨越式”的长程匹配——它的长程记忆是通过逐步压缩历史得到的，而不是直接检索。

对于需要精确匹配远距离 token 的任务（如：找到 1000 个 token 之前出现的某个特定词），Mamba 的性能会弱于 Transformer。

这是不可避免的：你没办法既有全局直接访问，又有线性复杂度。Mamba 用”压缩历史的间接访问”换取了线性代价。

---

9.6 Transformer、线性注意力、Mamba：三角权衡

架构	时间复杂度	空间复杂度	长程能力	精确匹配能力
Transformer	$O(N^{2}d)$	$O(N^2)$	全局直接访问	强（softmax 尖锐）
线性注意力	$O(N{d}^2)$	$O(d^2)$	有限容量状态	弱（圆滑，噪声稀释）
Mamba (SSM)	$O(Nd)$	$O(d)$	压缩式历史记忆	中（选择性记忆）

没有免费的午餐。全局感受野、线性复杂度、精确匹配能力——你最多只能同时拥有其中两个。

这是信息论的基本约束，不是工程问题。

---

八、位置编码：Transformer对”顺序”的理解

Self-Attention有一个奇怪的性质:它对词序不敏感。

如果你把”cat sat mat”打乱成”mat cat sat”,注意力矩阵会改变(因为 $Q{K}^T$ 的行列顺序变了),但计算过程完全相同。Transformer本身没有”第1个词”、“第2个词”的概念。

但语言是有顺序的。“cat sat mat”和”mat sat cat”意思不同。

解决方案:位置编码(Positional Encoding)。在输入嵌入中加入位置信息:

$$\text{input}={\text{word}}_{\text{embedding}}+{\text{position}}_{\text{encoding}}$$

原始Transformer使用正弦位置编码:

$$P{E}_{(pos,2i)}=\sin (pos/{10000}^{2i/d})$$$$P{E}_{(pos,2i+1)}=\cos (pos/{10000}^{2i/d})$$

其中 $pos$ 是位置, $i$ 是维度索引。

为什么用正弦函数? 因为它有一个优雅的性质:

$$P{E}_{pos+k}=f(P{E}_{pos})$$

也就是说,相对位置可以表示为绝对位置的线性函数。这让模型更容易学习”相对位置关系”。

正弦位置编码是一种绝对位置编码方法，其编码是确定性的且不包含可学习参数，因此不会引入额外的学习噪声，然而由于这个编码方法通过与token embedding相加的方式注入的，这样token表示中的内容信息与绝对位置信息是强耦合的，从而在注意力计算中难以被显式解耦。其局限在于——并未显式建模token之间的相对位置关系，模型需要通过绝对位置编码间接推导这些关系，从而增加了学习负担。

此外，这种隐式建模方式可能在一定程度上影响模型对长距离依赖的捕获能力以及对未见序列长度的外推能力。

后续研究提出了相对位置编码:

• RoPE(旋转位置编码):直接在注意力计算中编码相对位置
• ALiBi:在注意力分数上加入位置偏置,距离越远权重越低

这些方法在长序列外推上表现更好——训练时见过512长度,测试时可以处理2048长度。

---

九、伪代码：Transformer的完整前向传播

让我把Transformer的核心算法形式化。

算法1:Self-Attention层

import torch
import torch.nn.functional as F

def self_attention(X, W_Q, W_K, W_V):
    # X: (N, d_model), W_Q/W_K/W_V: (d_model, d_k)
    Q = X @ W_Q          # (N, d_k)
    K = X @ W_K          # (N, d_k)
    V = X @ W_V          # (N, d_v)
    d_k = Q.shape[-1]
    scores = Q @ K.T / d_k ** 0.5              # (N, N)
    attn   = F.softmax(scores, dim=-1)         # 按行归一化
    return attn @ V                            # (N, d_v)
# 时间复杂度 O(N²·d)，空间复杂度 O(N²)

算法2:多头注意力

import torch
import torch.nn as nn

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        assert d_model % num_heads == 0
        self.h  = num_heads
        self.dk = d_model // num_heads
        self.W_Q = nn.Linear(d_model, d_model, bias=False)
        self.W_K = nn.Linear(d_model, d_model, bias=False)
        self.W_V = nn.Linear(d_model, d_model, bias=False)
        self.W_O = nn.Linear(d_model, d_model, bias=False)

    def forward(self, X):
        B, N, D = X.shape
        # 投影并拆分为多头 (B, h, N, dk)
        Q = self.W_Q(X).view(B, N, self.h, self.dk).transpose(1, 2)
        K = self.W_K(X).view(B, N, self.h, self.dk).transpose(1, 2)
        V = self.W_V(X).view(B, N, self.h, self.dk).transpose(1, 2)
        attn = F.softmax(Q @ K.transpose(-2, -1) / self.dk**0.5, dim=-1)
        out  = (attn @ V).transpose(1, 2).reshape(B, N, D)  # 拼接各头
        return self.W_O(out)

算法3:完整Transformer块

class TransformerBlock(nn.Module):
    def __init__(self, d_model, num_heads, d_ff):
        super().__init__()
        self.attn  = MultiHeadAttention(d_model, num_heads)
        self.ff1   = nn.Linear(d_model, d_ff)
        self.ff2   = nn.Linear(d_ff, d_model)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)

    def forward(self, X):
        X = self.norm1(X + self.attn(X))                         # 注意力 + 残差 + LN
        X = self.norm2(X + self.ff2(F.relu(self.ff1(X))))        # FFN + 残差 + LN
        return X
# 完整 Transformer = 堆叠多个 TransformerBlock

---

十、注意力的哲学：理解还是检索？

Transformer的成功引发了一个深刻的问题:注意力机制是在”理解”关系,还是在”检索”相关信息?

考虑这个句子:

"The trophy doesn't fit in the suitcase because it is too big."

人类知道”it”指代”trophy”(而不是”suitcase”),因为我们理解”太大”的逻辑——trophy太大所以放不进去。

Transformer也能正确地让”it”的注意力集中在”trophy”上。但它是真的理解了因果逻辑,还是只是学到了统计模式:“it”通常指代前面的名词,而”too big”附近的”it”更可能指代第一个名词?

这个问题没有简单答案。但有一些线索:

证据1:注意力模式的可解释性

研究者发现,某些注意力头确实学到了语言学上有意义的模式：主语-谓语一致性、修饰语-中心词关系、代词-先行词绑定。

这些模式不是随机的,而是系统性的。

证据2:分布偏移下的脆弱性

但当测试数据的分布和训练数据不同时,这些模式会崩溃。如果训练时”it”总是指代第一个名词,测试时遇到”it”指代第二个名词的情况,模型会失败。

这表明Transformer学到的可能是统计捷径,而不是深层的因果理解。

证据3:注意力≠解释

更微妙的是,注意力权重本身可能不是模型决策的真正原因。2019年的研究表明,即使随机打乱注意力权重,模型的输出有时也不会改变太多——因为真正的信息可能编码在value向量中,而不是注意力分布中。

注意力是一个有用的工具,但它不是理解的充分条件。

---

十一、可视化：注意力头在”看”什么

图9.1:多头注意力的分工

---

Transformer 定义了推理的基础设施。下一章，我们将看到它如何与搜索结合——在状态空间中寻找推理的最优路径。

悬而未决

• 注意力机制是否真的在”理解”关系?

  还是只是在做复杂的模式匹配?考虑这个测试:如果你随机打乱注意力权重,但保持value向量不变,模型的输出会改变多少?Jain & Wallace (2019)发现:在某些任务上,打乱注意力权重对输出影响很小。这说明什么?是注意力权重不重要,还是真正的信息编码在value向量里?

• O(N²)的代价能否被根本性突破?

  所有”高效Transformer”都在牺牲全局感受野或引入近似。但这是必然的吗?信息论下界说:描述 $N$ 个节点的全连接图需要 $\mathrm{\Omega }(N^2)$ 比特。但也许我们不需要存储完整的注意力矩阵——也许存在一种隐式表示,用 $O(N)$ 空间编码 $N^2$ 的连接模式?

• 位置编码的本质是什么?

  为什么正弦函数有效?RoPE(旋转位置编码)为什么在长序列外推上更好?位置信息的最优编码方式是什么?这个问题等价于:如何在保持置换不变性的同时,注入顺序信息?

• 多头注意力的头数应该是多少?

  BERT用12个头,GPT-3用96个头。这个数字是如何选择的?是否存在最优头数的理论?一个可能的角度:每个头学习一个子空间的关系模式,头数应该等于”语言关系的本征维度”。但这个维度是多少?

• 注意力权重能否作为可解释性工具?

  如果注意力不等于解释(Jain & Wallace的论证),我们应该如何理解模型的决策过程?一个更深的问题:可解释性是否要求因果性?也就是说,我们需要的不是”模型在看什么”,而是”如果改变输入,模型的输出会如何变化”?

• Transformer的归纳偏置是什么?

  CNN有局部性偏置(卷积核),RNN有顺序性偏置(循环结构)。Transformer的归纳偏置是什么?”没有归纳偏置”本身是一种归纳偏置吗?这个问题触及架构即先验的哲学:每个架构都在假设观测数据的某种结构,只是有些假设比其他假设更隐蔽。

---

延伸阅读

• Vaswani et al. (2017). Attention Is All You Need — Transformer原始论文，开启了NLP的新纪元 → [arXiv:1706.03762]

• Katharopoulos et al. (2020). Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention — 线性注意力的核函数推导，加法结合律换来 $O(N)$ 复杂度及其圆滑性代价 → [arXiv:2006.16236]

• Gu, A. et al. (2020). HiPPO: Recurrent Memory with Optimal Polynomial Projections — 勒让德多项式最优历史压缩，SSM 的理论基础 → [arXiv:2008.07669]

• Gu, A. & Dao, T. (2023). Mamba: Linear-Time Sequence Modeling with Selective State Spaces — 选择性 SSM，用内容驱动的 ${\mathrm{\Delta }}_t$ 在时间维度上实现 softmax 的尖锐性 → [arXiv:2312.00752]

• Devlin et al. (2018). BERT: Pre-training of Deep Bidirectional Transformers — 双向Transformer预训练，展示了注意力机制的强大 → [arXiv:1810.04805]

• Su et al. (2021). RoFormer: Enhanced Transformer with Rotary Position Embedding — RoPE旋转位置编码，改进长序列外推 → [arXiv:2104.09864]

• Press et al. (2021). Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation — ALiBi位置偏置方法

• Clark et al. (2019). What Does BERT Look At? An Analysis of BERT’s Attention — 注意力可视化与可解释性研究

• Tay et al. (2020). Efficient Transformers: A Survey — 高效Transformer架构综述

• Jain & Wallace (2019). Attention is not Explanation — 对注意力作为解释工具的批判性分析

• Wiegreffe & Pinter (2019). Attention is not not Explanation — 对上述批判的反驳,注意力的辩护