注意力机制：从入门到精通

🟢 入门版：面向高中生

想象你在读一本小说，读到这一句：

“小明把苹果放在了桌上。它很红。”

你的大脑会自动把”它”和”苹果”联系起来，而不是和”桌”联系起来。你怎么做到的？因为在你的理解里，“红”通常形容的是苹果，而不是桌子。你的注意力自然地聚焦在了”苹果”这个词上。

注意力机制做的事情，本质上就是一模一样的事。

当我们让 AI 阅读一段文字时，我们希望它也能像人一样——看到每个词的时候，知道该”关注”哪些其他词。比如在翻译”the bank of the river”时，AI 需要知道这里的”bank”是河岸，不是银行。怎么知道？靠的就是看看上下文里”river”这个词——注意力机制让 AI 能够”回头看”，把相关的信息提取出来。

具体怎么做呢？可以这么理解：

假设每个词都是一个同学，教室里有 N 个同学。现在老师问了一个问题，每个同学都需要综合全班的信息来回答自己的那部分。但不是所有同学的信息都同样重要——回答数学题时你更想听数学课代表的，回答历史题时你更想听历史课代表的。

注意力机制就是在做这个”分配注意力”的过程。对于每个词，它会去打量所有其他词，然后给每个词一个”重要性分数”。分数高的词，就对当前词的理解贡献更大。最后，把所有词的信息按照分数加权混合，就得到了这个词在上下文中的”完整理解”。

这个机制为什么厉害？因为在它之前，AI 处理语言主要靠 RNN，就像一个人只能从左到右一个词一个词地读，读完就忘了。而注意力机制让 AI 可以同时看到所有词，自由地建立任意两个词之间的关联。不管两个词隔了多远，注意力机制都能把它们连起来——这就是”全局视野”。

ChatGPT、翻译软件、写代码的 AI，底层都在用这个机制。它就是现代 AI 能够理解语言的基石。

🟡 中级版：面向理工科大学生

好，我们来看数学。

Self-Attention 的核心思路是：对于序列中的每个 token，计算它和所有其他 token 的关联程度，然后根据关联程度聚合信息。

三个关键概念：Query（查询）、Key（键）、Value（值）。

这个命名来自信息检索的类比。想象你去图书馆找书：

你心里想的”我要找什么”就是 Query
每本书封面上的标签和摘要就是 Key
书的实际内容就是 Value

图书馆会根据你的 Query 和每本书的 Key 的匹配程度，决定给你看哪些书的 Value。

在 Self-Attention 中，每个 token 同时扮演三种角色：

作为当前 token，它产生一个 Query（“我在找什么信息”）
每个候选 token 都有一个 Key（“我能提供什么信息”）
每个候选 token 还有一个 Value（“我的实际信息内容”）

Scaled Dot-Product Attention

设输入序列经过 embedding 后为矩阵 X（大小为 n×d），其中 n 是序列长度，d 是嵌入维度。

首先通过三个线性变换得到 Q、K、V：

Q = X · Wq,    K = X · Wk,    V = X · Wv

其中 Wq、Wk、Wv 是可学习的参数矩阵（大小 d×dk）。

然后计算注意力分数：

Attention(Q, K, V) = softmax(Q · K^T / √dk) · V

展开来看：Q·K^T 的结果是一个 n×n 的矩阵，每个元素 (i, j) 就是第 i 个 token 对第 j 个 token 的”关注度”。除以 √dk 是为了防止点积过大导致 softmax 梯度消失（这是 Transformer 论文中的一个关键细节）。softmax 把分数归一化为概率分布，最后乘以 V 得到加权后的输出。

直觉总结： 每个 token 的输出 = 所有 token 的 Value 按注意力权重加权求和。注意力权重越大，说明两个 token 的关联越强。

这就是 Self-Attention 的全部核心。简单、优雅、强大。整个 Transformer 架构在此基础上堆叠了多层、加入了残差连接和 Layer Normalization，但注意力计算本身就是这么干净。

🔴 高级版：面向 ML 研究者

Multi-Head Attention 的设计动机

单头注意力有一个问题：softmax 会导致输出倾向于关注少数几个 token（attention 的”聚焦”效应），这限制了模型同时捕捉多种关系的能力。

Multi-Head Attention 将 Q、K、V 投影到 h 个不同的子空间，每个”头”独立计算注意力，然后拼接：

MultiHead(Q, K, V) = Concat(head_1, ..., head_h) · Wo
head_i = Attention(Q·Wi^Q, K·Wi^K, V·Wi^V)

这等价于在不同的表示子空间中同时捕捉不同类型的依赖关系。比如一个头可能关注语法结构（主谓关系），另一个关注语义相似性，还有一个关注位置邻近性。实验表明，不同的头确实会学到不同的注意力模式——有些是局部的（关注相邻 token），有些是全局的（关注距离很远的 token）。

位置编码的必要性

Self-Attention 本身是 置换不变的（permutation invariant）——输入顺序打乱，输出（的集合）不变。这对 NLP 是致命的，因为语言有严格的语序。

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

PE(pos, 2i)   = sin(pos / 10000^(2i/d))
PE(pos, 2i+1) = cos(pos / 10000^(2i/d))

正弦编码的好处是可以外推到训练时没见过的长度，但实际效果有限。现代模型（GPT、LLaMA 等）普遍使用 可学习的绝对位置编码 或 RoPE（Rotary Position Embedding） 等相对位置编码。RoPE 通过旋转矩阵将相对位置信息注入 Q 和 K 的点积中，数学上非常优雅，且已被 LLaMA、PaLM 等模型广泛采用。

复杂度分析：O(n²d) 问题

标准 Self-Attention 的时间和空间复杂度都是 O(n²d)，其中 n² 来自 Q·K^T 的注意力矩阵。当序列很长时（比如长文档、视频），这个二次复杂度成为瓶颈。

空间瓶颈尤其严重： 注意力矩阵需要 O(n²) 的内存来存储，这对于大模型 + 长序列是致命的。一个 8B 参数的模型处理 128K 长度的序列，仅注意力矩阵就需要数十 GB 显存。

Flash Attention

Flash Attention（Dao et al., 2022）的核心思想是：不要显式地实例化完整的 n×n 注意力矩阵。

标准实现会先算出完整的 Q·K^T（写入 HBM），再算 softmax（写入 HBM），再乘 V（写入 HBM）。HBM 带宽是瓶颈。

Flash Attention 用了 tiling 技巧：将 Q、K、V 分块加载到 SRAM（高速缓存）中，在 SRAM 里完成注意力计算，只把最终结果写回 HBM。关键是利用了 softmax 的数学性质——可以在线（online）计算分块 softmax，不需要看到完整的矩阵。

这使得 IO 复杂度从 O(n²d) 降低到 O(n²d²/M)（M 是 SRAM 大小），实际效果是速度提升 2-4 倍，且内存使用从 O(n²) 降低到 O(n)（不需要存储完整注意力矩阵）。

后续的 Flash Attention 2 进一步优化了并行化，Flash Attention 3 针对 Hopper 架构（H100）做了异步分块等硬件级优化。这些工作使得训练和推理超长序列成为可能。

前沿方向

对于超长序列（百万 token 级别），社区在探索线性注意力（如 Linformer 用低秩近似、Performer 用随机特征）、稀疏注意力（如 Longformer、BigBird 用稀疏模式）以及状态空间模型（Mamba、RWKV）等替代方案。但 Flash Attention 的优化使得标准注意力在中等长度（~100K）下仍然是首选——工程上的精巧优化有时比架构上的根本改变更实用。