深入剖析 Transformer 自注意力机制：从数学原理到工程实现

引言

2017 年，Google 在论文《Attention Is All You Need》中提出了 Transformer 架构，彻底抛弃了 RNN 和 CNN 的序列建模方式，仅凭注意力机制就完成了序列到序列的映射。如今，从 BERT、GPT 到 LLaMA、DeepSeek，几乎所有主流大语言模型都以自注意力（Self-Attention）作为核心计算单元。

然而，"注意力"这个词在科普文章中被反复提及，真正理解其内部数学结构、计算流程和工程优化细节的人却不多。本文不谈宏观概念，只讲技术细节——从公式的每一项含义，到矩阵运算的内存布局，再到 FlashAttention 的分块技巧，力求把这一机制拆解到"能自己手写实现"的程度。

一、自注意力的数学定义

1.1 基本公式

给定输入序列 X ∈ R^(n×d)（n 为序列长度，d 为模型维度），自注意力的计算过程为：

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

其中：

• Q = X·W_Q，K = X·W_K，V = X·W_V，三个权重矩阵 W_Q, W_K ∈ R^(d×d_k)，W_V ∈ R^(d×d_v)
• QK^T ∈ R^(n×n) 是查询与键的点积矩阵，衡量每对 token 之间的相似度
• √d_k 是缩放因子，防止点积值过大导致 softmax 梯度消失
• softmax 沿最后一个维度（即 K 的方向）归一化
• 最终输出 ∈ R^(n×d_v)，与输入形状一致

1.2 为什么要除以 √d_k

这不是一个随意的选择。假设 Q 和 K 的每个元素是均值为 0、方差为 1 的独立随机变量，那么点积 q·k = Σ(i=1→d_k) q_i·k_i 的均值为 0，方差为 d_k。当 d_k 较大时（例如 64 或 128），点积的值会很大，使得 softmax 输出接近 one-hot 分布——梯度几乎为零，训练无法进行。除以 √d_k 将方差拉回 1，保持梯度在合理范围内。

1.3 softmax 的数值稳定性

直接计算 e^(x_i) 会溢出。实际实现中采用 max-subtraction 技巧：

softmax(x)_i = e^(x_i - max(x)) / Σ_j e^(x_j - max(x))

减去最大值后，所有指数的输入 ≤ 0，e 的值在 (0, 1] 之间，不会溢出。这一步看似简单，但在后续 FlashAttention 的分块计算中会带来非平凡的处理难度——后面会详细展开。

二、多头注意力（Multi-Head Attention）

2.1 结构

单头注意力只有一个投影空间，表达能力有限。多头注意力将 d 维空间拆成 h 个子空间，每个头独立做注意力，再拼接：

MultiHead(Q, K, V) = Concat(head_1, ..., head_h) · W^O

head_i = Attention(Q·W_i^Q, K·W_i^K, V·W_i^V)

其中 W_i^Q, W_i^K ∈ R^(d×d_k)，W_i^V ∈ R^(d×d_v)，通常 d_k = d_v = d/h。

2.2 为什么有效

每个头可以关注不同类型的依赖关系：有的头关注局部语法（相邻 token），有的头关注长距离语义（代词指代），有的头关注位置模式。在可视化研究中，这种"分工"现象被反复观察到。需要注意的是，这种分工并非人为设计，而是训练过程中自发涌现的——这也是注意力机制可解释性讨论的焦点之一。

2.3 工程实现的关键细节

实际代码中不会真的创建 h 个独立的矩阵乘法。标准做法是：

1. 做一次大矩阵投影：Q = X·W_Q ∈ R^(n×d)
2. reshape 成 R^(n×h×d_k)，再 transpose 成 R^(h×n×d_k)
3. 所有头并行做注意力计算（batch matmul）
4. 输出 reshape 回 R^(n×d)，乘以 W^O

这样整个前向传播只需要 4 次矩阵乘法（W_Q, W_K, W_V, W^O），而不是 4h 次。在 GPU 上，这比循环逐头计算快一个数量级。

三、计算复杂度与瓶颈分析

3.1 复杂度

核心瓶颈在于 QK^T 和 attn·V 两步：计算量 O(n²·d)，中间矩阵 O(n²)。当 n = 8192（GPT-3 的上下文窗口），n² = 67M，注意力矩阵占显存约 134MB（FP32），看似不大；但当 n = 128K（长上下文模型），n² = 16.4B，仅注意力矩阵就需要约 65GB 显存——这远超单卡容量。

3.2 内存墙问题

在标准实现中，QK^T 的结果要写回 HBM（高带宽内存），softmax 要从 HBM 读回，再写回，最后乘 V 又要从 HBM 读回。对于 n = 4096, d = 4096，注意力矩阵有 16M 个元素，每次读写 64MB，来回三次就是 192MB 的 HBM 带宽消耗。而实际计算只需要几毫秒，大部分时间花在等内存搬运上——这就是"内存墙"。

四、位置编码：让注意力感知顺序

自注意力本身是置换不变的——打乱输入顺序，输出只是对应打乱，不会产生不同的计算结果。这对语言建模是不可接受的，"猫追狗"和"狗追猫"必须不同。位置编码就是解决这个问题的。

4.1 正弦/余弦编码（原始 Transformer）

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

优点：对任意相对位置 k，PE(pos+k) 可以表示为 PE(pos) 的线性函数，理论上能泛化到训练中未见过的序列长度。缺点：实践中长序列外推效果并不理想。

4.2 旋转位置编码（RoPE）

RoPE 是目前大模型的主流选择（LLaMA、Qwen、DeepSeek 均采用）。核心思想：不把位置信息加到输入上，而是通过旋转矩阵编码到 Q 和 K 上，使得 q_m^T · k_n 只依赖于相对位置 m - n。

具体地，将 d 维向量两两分组，对第 m 个位置的第 i 组施加旋转：

[q_2i  ]     [cos(m·θ_i)  -sin(m·θ_i)]   [q_2i  ]
[q_2i+1]  →  [sin(m·θ_i)   cos(m·θ_i)] · [q_2i+1]

其中 θ_i = 10000^(-2i/d)。K 同理旋转。这样点积自然地成为相对位置 m - n 的函数。RoPE 的优势在于：不引入额外参数、保持注意力矩阵的稀疏性友好、支持通过 NTK-aware 插值扩展上下文窗口。

4.3 ALiBi（Attention with Linear Biases）

另一种思路：不改 Q 和 K，直接在注意力分数上减去一个与距离成正比的偏置：

score_ij = q_i · k_j - m_h · |i - j|

其中 m_h 是每个头的固定斜率。越远的 token 惩罚越大，且不同头的惩罚速率不同。ALiBi 在长序列外推上表现优于正弦编码，但不如 RoPE 灵活，目前使用较少。

五、FlashAttention：突破内存墙的分块算法

5.1 核心思想

FlashAttention（Dao et al., 2022）的关键洞察：注意力计算不需要把完整的 n×n 矩阵物化到 HBM。通过分块（tiling）和在线 softmax（online softmax），将计算限制在 SRAM（片上高速缓存）中完成，只读写最终输出。

5.2 在线 softmax 的数学技巧

标准 softmax 需要看到整行才能计算。分块后，每块只有部分数据，如何增量更新？关键公式：

设已处理块的最大值为 m、指数和为 ℓ、加权值为 v，新来一块数据的最大值为 m’、指数和为 ℓ’、加权值为 v’，则合并后：

m_new = max(m, m')

ℓ_new = e^(m - m_new) · ℓ + e^(m' - m_new) · ℓ'

v_new = [e^(m - m_new) · ℓ · v + e^(m' - m_new) · ℓ' · v'] / ℓ_new

这三个公式允许逐块累积 softmax 结果，不需要回看已写出的数据。这就是 FlashAttention 能避免 n×n 矩阵物化的根本原因。

5.3 分块策略

将 Q, K, V 分别按行分块：

• Q 分成 B_r 行一块（block row）
• K, V 分成 B_c 列一块（block column）
• 外层循环遍历 Q 的块，内层循环遍历 K, V 的块
• 每个内层迭代在 SRAM 中完成一块 B_r × B_c 的注意力计算，用在线 softmax 更新累积器
• 一行处理完后，将最终结果写回 HBM

块大小 B_r, B_c 的选择取决于 SRAM 大小，通常使得 B_r × B_c × sizeof(float) ≤ SRAM capacity。在 A100 上，SRAM 约 192KB，B_r = B_c = 128 是常见配置。

5.4 性能提升的本质

FlashAttention 不是减少了计算量——FLOPs 仍然是 O(n²·d)。它减少的是 HBM 读写次数：从 O(n² + nd) 降到 O(n²·d / M)（M 为 SRAM 大小）。在 A100 上，这带来了约 2-4 倍的端到端加速，同时将峰值显存从 O(n²) 降到 O(n)——后者使得训练更长上下文成为可能。

5.5 FlashAttention-2 的改进

v2 的主要优化：

1. 减少非矩阵乘法运算：重新排列循环顺序，将 Q 放在外循环、K/V 放在内循环，使得每次内循环的 softmax 更新只涉及更少的行，减少了 rescale 操作。
2. 更好的并行性：沿序列维度和 head 维度同时并行，充分利用 GPU 的 SM（Streaming Multiprocessor）数量。
3. 更精确的分块：根据 d 的大小自适应选择 B_r，使得 warp 级别的利用率最大化。

六、KV Cache：推理阶段的关键优化

6.1 问题

自回归生成时，每生成一个 token 就要重新计算整个序列的注意力。第 t 步的输入长度为 t，总计算量为 Σ(t=1→n) t = O(n²)，与序列长度呈平方关系。

6.2 解决方案

注意到每一步的 K 和 V 只是在前面追加一行，之前计算过的部分不变。因此可以缓存历史 K 和 V，每步只计算新 token 的 q, k, v，然后与缓存拼接做注意力：

# 生成第 t 个 token
q_new = x_t @ W_Q           # (1, d)
k_new = x_t @ W_K           # (1, d)
v_new = x_t @ W_V           # (1, d)
K_cache = concat(K_cache, k_new)  # (t, d)
V_cache = concat(V_cache, v_new)  # (t, d)
attn = softmax(q_new @ K_cache.T / sqrt(d)) @ V_cache  # (1, d)

这样每步的计算量降为 O(t·d)，总计算量降为 O(n²·d/2)——虽然渐进复杂度相同，但常数因子减半，且避免了重复的矩阵投影。

6.3 KV Cache 的显存开销

KV Cache 的大小为 2 × n × d × L × sizeof(float)（2 表示 K 和 V，L 为层数）。以 LLaMA-2-70B 为例：d = 8192，L = 80，n = 4096，FP16：

2 × 4096 × 8192 × 80 × 2 bytes = 10.7 GB

这比模型权重本身（约 140GB FP16）小，但在多用户并发推理时，每个请求都需要独立的 KV Cache，显存成为主要瓶颈。这也是 PagedAttention、KV Cache 量化等技术出现的动机。

6.4 Multi-Query Attention 与 Grouped-Query Attention

标准多头注意力中，每个头有独立的 K 和 V，KV Cache 随头数线性增长。MQA（Multi-Query Attention）让所有头共享同一组 K 和 V，将 KV Cache 缩小 h 倍，但质量略有下降。GQA（Grouped-Query Attention）是折中方案：将 h 个头分成 g 组，每组内共享 K 和 V。当 g = 1 时退化为 MQA，g = h 时退化为标准 MHA。LLaMA-2-70B 使用 g = 8，在质量和效率之间取得了较好的平衡。

七、注意力变体与稀疏化

7.1 滑窗注意力（Sliding Window Attention）

限制每个 token 只关注窗口大小 w 内的 token，注意力矩阵从稠密变为带状。计算复杂度降为 O(n·w·d)。Mistral 7B 使用 w = 4096 的滑窗注意力，通过堆叠多层让信息间接传播到全序列（第 L 层的有效感受野为 L × w）。

7.2 稀疏注意力（Sparse Attention）

在注意力矩阵中人为设置稀疏模式：局部窗口 + 少量全局 token + 随机连接。BigBird 证明了这种模式在理论上保持了 Transformer 的 TC0（复杂度类）表达能力，同时将计算复杂度降为 O(n)。

7.3 线性注意力（Linear Attention）

将 softmax 近似为核函数 φ(q)·φ(k)，利用矩阵乘法结合律重排计算顺序：

Attention = φ(Q) · (φ(K)^T · V)

先算 φ(K)^T · V ∈ R^(d×d)（与 n 无关），再乘 φ(Q)。这样计算复杂度降为 O(n·d²)，与序列长度呈线性关系。代价是 softmax 的非线性被移除，表达能力下降，在语言建模任务上效果不如标准注意力。Linear Transformer、RWKV 等模型采用了这类思路。

八、从注意力看大模型架构演进

回顾以上技术，可以看出大模型架构演进的一条清晰主线：

每一代改进都在回答同一个核心问题：如何在保持注意力建模能力的同时，降低其 O(n²) 的计算和内存代价。FlashAttention 从工程层面回答了这个问题（不改变数学，只改变实现），稀疏/线性注意力从数学层面回答了这个问题（改变注意力本身的定义），而 GQA/MQA 则从架构层面回答了推理效率的问题。

九、手写实现：最小化的自注意力

将以上理论浓缩为一段可运行的 PyTorch 代码：

import torch
import torch.nn.functional as F
import math

class MultiHeadSelfAttention(torch.nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        assert d_model % n_heads == 0
        self.d_k = d_model // n_heads
        self.n_heads = n_heads

        # Q, K, V 合并为一个矩阵投影，减少 kernel launch 开销
        self.qkv = torch.nn.Linear(d_model, 3 * d_model, bias=False)
        self.o_proj = torch.nn.Linear(d_model, d_model, bias=False)

    def forward(self, x):
        """
        x: (batch, seq_len, d_model)
        return: (batch, seq_len, d_model)
        """
        B, N, D = x.shape
        H, d_k = self.n_heads, self.d_k

        # 1. 投影 + 分头: (B, N, D) -> (B, N, 3D) -> 3 × (B, H, N, d_k)
        qkv = self.qkv(x).reshape(B, N, 3, H, d_k).permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]  # 各为 (B, H, N, d_k)

        # 2. 缩放点积注意力
        scale = 1.0 / math.sqrt(d_k)
        attn = (q @ k.transpose(-2, -1)) * scale  # (B, H, N, N)
        attn = F.softmax(attn, dim=-1)

        # 3. 加权求和 + 合并头: (B, H, N, d_k) -> (B, N, D)
        out = (attn @ v).transpose(1, 2).reshape(B, N, D)

        # 4. 输出投影
        return self.o_proj(out)

这段代码约 30 行，涵盖了本文第一、二节讨论的所有核心概念：QKV 投影、多头分拆、缩放点积、softmax 归一化、头合并、输出投影。实际生产代码中，第 2-3 步会被替换为 F.scaled_dot_product_attention（PyTorch 2.0+ 内置，底层调用 FlashAttention），但理解这段代码是理解一切优化的前提。

结语

自注意力机制的精妙之处在于：它的数学定义极其简洁——一行公式，三个矩阵乘法，一个 softmax——但其工程实现却衍生出了 FlashAttention 的分块技巧、KV Cache 的推理优化、GQA 的架构折中、RoPE 的位置编码等一整套技术体系。

理解这些细节的意义不仅在于"知道大模型怎么工作的"，更在于当你在实际工程中遇到显存溢出、推理延迟过高、长上下文效果差等问题时，能准确定位瓶颈所在，并选择正确的优化手段。注意力机制不是黑箱——它是一系列可拆解、可优化、可替换的数学运算，而本文试图把这种可拆解性展现出来。