第 4 章扩展：注意力机制的变体

在线 Notebook

对应的交互式版本可在 Google Colab 打开，统一使用方式见 第0章说明。

标准注意力并不是所有场景的唯一选择。所谓“注意力变体”，通常是在调整以下某一项：哪些位置可以连接、如何计算相似度、多少查询头共享键和值、怎样表示位置，或者怎样在硬件上执行同一个计算。

图 E4.1：先确定瓶颈，再选择注意力方案。不同分支解决的问题不同，不能只按“更快”或“更强”排列。

E4.1 先分清五类改动

改动维度	代表方法	主要目标	是否改变标准注意力的连接或结果
连接范围	局部、滑动窗口、块稀疏、全局 token	降低长序列计算量	改变连接范围
数值近似	低秩投影、核函数近似	避免显式构造完整注意力矩阵	通常是近似计算
KV 共享	MQA、GQA	减少解码时的 KV cache 和内存带宽	改变参数化方式
位置信息	相对位置偏置、RoPE	让注意力分数包含顺序或相对距离	改变位置注入方式
执行方式	分块、融合算子、FlashAttention	减少显存读写和中间结果	数学结果保持等价或近似等价

这五类方法可以组合。例如，一个生成模型可以同时使用 RoPE、GQA 和 FlashAttention：三者分别处理位置、KV cache 和硬件执行问题。

E4.2 相似度与掩码

注意力的共同结构是“根据查询与键计算权重，再对值加权求和”：

$$\mathrm{Attention}(Q,K,V)=\mathrm{softmax}(S+M)V,$$

其中 $S$ 是相似度分数，$M$ 是可选掩码。

缩放点积注意力使用：

$$S=\frac{Q{K}^{\mathrm{\top }}}{\sqrt{d_k}}.$$

加法注意力和双线性注意力只是换了一种分数函数：

$$s(q,k)=v^{\mathrm{\top }}\tanh (W_{q}q+W_{k}k),s(q,k)=q^{\mathrm{\top }}Wk.$$

不能脱离模型规模、硬件和任务断言某个分数函数一定“更强”。Transformer 常用缩放点积，是因为它可以高效地写成矩阵乘法。

自回归生成还需要因果掩码：

$$M_{ij}=\{\begin{array}{ll}0,& j\le i,\\ -\mathrm{\infty },& j>i.\end{array}$$

它禁止位置 $i$ 读取未来 token。因果掩码解决的是信息泄漏，不是计算加速。

E4.3 改变连接范围：局部与稀疏注意力

完整自注意力让每个位置连接全部 $n$ 个位置，分数矩阵有 $n^2$ 个元素。局部或稀疏注意力只保留部分连接，例如：

• 滑动窗口：每个位置只读取附近 $w$ 个位置；
• 分块连接：先在块内计算，再通过少量跨块连接交换信息；
• 全局 token：为少数位置保留全局读写能力。

若每个位置只连接 $w$ 个位置，注意力分数计算量约为 $O(nw)$。但实际速度还取决于稀疏模式能否被硬件和算子有效利用。

代价也很直接：没有连接的两个位置不能在这一层直接交换信息。选择稀疏模式时，应检查任务所需的远距离依赖能否通过多层传播或全局连接到达。

E4.4 近似注意力：低秩与核方法

另一条路线是不删除连接，而是近似注意力计算。例如：

• 低秩方法把序列维投影到更小的维度；
• 核方法改写或近似 softmax 相似度，使计算顺序可以重排。

这些方法的复杂度取决于投影维度、特征数和具体实现，不能统一写成无条件的 $O(n)$。近似误差也可能随序列长度、数据分布和任务变化，因此至少要同时比较：任务指标、峰值显存、吞吐量和目标长度下的稳定性。

E4.5 改变 KV 共享：MHA、GQA 与 MQA

设查询头数为 $h_q$，键值头数为 $h_{kv}$：

• 多头注意力（MHA）：通常 $h_{kv}=h_q$；
• 分组查询注意力（GQA）：$1<h_{kv}<h_q$；
• 多查询注意力（MQA）：$h_{kv}=1$。

多个查询头共享同一组键和值，可以减少自回归解码时需要保存和读取的 KV cache。若其他维度不变，KV cache 中键和值的规模会大致按 $h_{kv}/h_q$ 缩减；这不等于整个模型参数量按同样比例缩减。

共享越多，缓存越小，但模型表达和训练结果是否保持，需要通过目标任务验证。GQA 提供的是共享粒度的中间选项，不保证在所有模型上都优于 MHA 或 MQA。

E4.6 注入位置信息：相对偏置与 RoPE

注意力本身不包含 token 顺序。相对位置偏置直接把距离相关项加入分数：

$$S_{ij}=\frac{q_i^{\mathrm{\top }}{k}_j}{\sqrt{d_k}}+b(i-j).$$

RoPE 则在计算点积前，对查询和键的二维分量施加与位置相关的旋转。以一个二维分量对为例：

$$R_m=\left[\begin{array}{cc}\cos (m\theta )& -\sin (m\theta )\\ \sin (m\theta )& \cos (m\theta )\end{array}\right],q_m^{\prime }=R_m{q}_m,k_n^{\prime }=R_n{k}_n.$$

由于 $R_m^{\mathrm{\top }}{R}_n=R_{n-m}$，旋转后的点积包含相对位移 $n-m$。这解释了 RoPE 如何把相对位置信息带入分数，但不意味着模型无需验证就能稳定处理任意超出训练范围的长度。

E4.7 优化执行：FlashAttention

普通实现可能把完整分数矩阵和 softmax 中间结果反复写入显存。FlashAttention 通过分块计算、在线 softmax 和算子融合，减少高带宽显存与片上存储之间的数据搬运。

它优化的是同一个注意力计算的执行过程，不是把完整注意力改成稀疏注意力。加速比例取决于序列长度、数据类型、硬件、掩码和实现版本，因此应在目标环境中实测，而不是使用固定倍数。

E4.8 如何选择

先记录瓶颈，再选择对应改动：

1. 训练时分数矩阵过大：考虑局部、稀疏或近似注意力。
2. 自回归解码的 KV cache 或带宽成为瓶颈：考虑 GQA 或 MQA。
3. 数学结构不想改变，但显存读写开销高：考虑 FlashAttention 一类的精确执行优化。
4. 任务依赖顺序、距离或长度变化：比较位置编码方案，并单独测试训练长度外的行为。

比较方案时固定模型规模、输入长度、批量大小和精度，至少记录任务指标、峰值显存、训练吞吐和解码延迟。只有这些条件一致，“更快”或“更省”才有可比性。

最低实践任务

在第 4 章已有自注意力实现上完成以下实验：

1. 加入因果掩码，检查注意力矩阵上三角是否为零；
2. 将完整注意力改为窗口大小为 $w$ 的局部注意力；
3. 对两个版本分别记录允许连接的元素数，并解释它为何从 $n^2$ 变为约 $nw$；
4. 不要求给出速度结论；若测量运行时间，必须同时写出设备、序列长度、批量大小和预热次数。

完成标准是：能够指出自己改动的是连接范围，而不是 KV 共享、位置编码或执行内核。

延伸阅读

• Vaswani et al., Attention Is All You Need：缩放点积注意力与多头注意力。
• Shaw et al., Self-Attention with Relative Position Representations：相对位置表示。
• Su et al., RoFormer：旋转位置编码。
• Shazeer, Fast Transformer Decoding: One Write-Head is All You Need：MQA。
• Dao et al., FlashAttention：面向 IO 的精确注意力实现。

---

返回： 第 4 章：Transformer 详解