第2章 Multi-Head Attention（MHA）深度解析

2、Multi-Head Attention（MHA）深度解析：计算原理、参数量与显存占用分析

好，咱们今天来啃一块硬骨头——Multi-Head Attention。说实话，MHA 是整个 Transformer 里最核心、也最“吃”资源的部分。很多同学跑不动大模型，或者训练时 OOM，十有八九是这里没搞明白。

我个人习惯，在调优之前，先把它的“底裤”扒干净。说白了，就是搞清楚三件事：它怎么算的？有多少参数？吃多少显存？ 这三件事搞定了，你才能知道瓶颈在哪，怎么去优化。

MHA 的基础是单头注意力，也就是 Scaled Dot-Product Attention。公式很简单：

Attention(Q, K, V) = softmax(Q * K^T / sqrt(d_k)) * V

嗯，这里要注意，d_k 是每个头的维度。为什么要除以 sqrt(d_k)？因为点积的结果会随着维度增大而变大，softmax 的梯度会变得非常小，模型就学不动了。我在项目中遇到过，有同事没加这个缩放因子，训练 loss 死活降不下去，查了半天才发现是这里的问题。

整个计算流程，我拆成三步：

你想想看，这个过程的计算复杂度是 O(n²·d)，n 是序列长度。序列一长，这个平方项就是灾难。

单头注意力只能关注一种“关系”。但一句话里，词与词之间的关系是复杂的。比如“他喜欢打篮球，因为很帅”，这里的“他”和“篮球”、“帅”都有关系，但侧重点不同。

MHA 的做法很简单：把 Q、K、V 分别投影到 h 个不同的子空间，然后并行地做 h 次注意力计算。最后把结果拼起来，再投影回去。

公式长这样：

MultiHead(Q, K, V) = Concat(head_1, ..., head_h) * W_O
其中 head_i = Attention(Q * W_Q_i, K * W_K_i, V * W_V_i)

这里 W_Q_i、W_K_i、W_V_i 是每个头的投影矩阵，W_O 是输出投影矩阵。

我刚开始学的时候，总觉得这 h 个头是“分开”的。其实不是，它们在计算上是并行的。你可以把 Q 直接拆成 h 块，一次矩阵乘法搞定所有头的投影。这才是工业级的做法。

核心要点：MHA 的本质，是在不显著增加计算量的前提下，让模型从多个角度去理解序列。每个头可以关注不同的位置组合。

咱们来算算账。假设模型维度是 d_model，头数是 h，每个头的维度 d_k = d_model / h。

MHA 的参数来自四个投影矩阵：

所以，MHA 的总参数量是 4 * d_model²。

举个例子，BERT-base 的 d_model=768，那么 MHA 的参数量就是 4 * 768² ≈ 2.36M。这占了整个模型参数的很大一部分。

我的经验：很多同学以为头数 h 会影响参数量，其实不会。因为 d_k = d_model / h，所有头的投影矩阵加起来，参数量还是 d_model²。h 只影响计算方式和显存占用，不影响参数量。

显存占用，才是真正让人头疼的地方。MHA 的显存占用主要来自两部分：参数和中间激活值。

参数显存：就是上面算的 4 * d_model²，每个参数如果是 float32，占 4 字节。这部分是固定的，不大。

中间激活值显存：这才是大头！

在计算过程中，我们需要保存 Q、K、V、注意力分数矩阵、softmax 后的概率矩阵、以及加权后的输出。其中，注意力分数矩阵的形状是 [batch_size, h, seq_len, seq_len]。

你想想看，当 seq_len=1024，h=16，batch_size=32 时，这个矩阵的大小是：

32 * 16 * 1024 * 1024 * 4 bytes = 2GB

这还只是一个注意力分数矩阵！再加上 Q、K、V 和输出，轻松突破 4GB。这就是为什么长序列训练时，显存会爆炸。

避坑指南：我曾经在训练一个长文档模型时，seq_len 设到了 2048，结果 32GB 的 V100 直接 OOM。后来才发现，注意力分数矩阵占了 8GB 多。解决办法是使用梯度检查点（Gradient Checkpointing），或者改用 Flash Attention。

既然知道了问题在哪，咱们就聊聊怎么优化。我总结了几个常用的方法：

减少头数 h：在 d_model 不变的情况下，减少 h 可以降低注意力分数矩阵的大小。但要注意，头数太少会影响模型表达能力。
使用 Flash Attention：这是目前最流行的优化方法。它通过分块计算和重计算，避免了显式地保存整个注意力分数矩阵，显存占用从 O(n²) 降到 O(n)。
稀疏注意力：只计算部分位置的注意力分数，比如局部窗口注意力或全局+局部注意力。这在处理超长序列时非常有效。
混合精度训练：把参数和激活值从 float32 降到 float16 或 bfloat16，显存直接减半。

我个人习惯，在项目初期先用标准 MHA 跑通，然后根据显存瓶颈，逐步引入这些优化技巧。不要一上来就搞花里胡哨的，先把基线跑出来再说。

MHA 的核心，就是通过多头并行，让模型从不同角度理解序列。它的参数量是 4 * d_model²，不算大。但它的显存占用，尤其是注意力分数矩阵，是 O(n²) 的，这才是真正的性能瓶颈。

嗯，这一章的内容就到这里。下一章咱们聊聊 FFN 的优化，那也是个大头。