4、多头自注意力(MSA)详解：为什么需要多头？多头如何工作？计算复杂度分析。
好，咱们进入正题。前面聊了单头自注意力，说白了就是一个注意力函数在运作。但你会发现，单头注意力有个问题——它太“专一”了。我打个比方，就像你只用一只眼睛看世界，虽然也能看清，但深度感和细节总会差那么点意思。
多头注意力（Multi-Head Self-Attention, MSA）就是来解决这个问题的。它让模型同时从多个角度去“看”输入序列。嗯，这里要注意，不是简单的重复，而是每个头关注不同的子空间。
为什么需要多头？
我个人习惯把这个问题拆成两点来讲：
- 表达能力受限：单头注意力只能学习一种注意力分布。比如在句子“我喜欢吃苹果，也喜欢吃香蕉”中，一个头可能只关注“苹果”和“吃”的关系，忽略了“香蕉”。多头就能同时捕捉多种关系。
- 信息瓶颈：单头把所有的信息压缩到一个注意力矩阵里，信息会互相干扰。我在项目中遇到过，用单头做图像分类，模型老是关注背景噪声，后来换成多头，每个头各司其职，效果立马就上来了。
核心观点：多头注意力本质上是让模型拥有多个“视角”，每个视角关注不同的特征子空间。这就像团队协作，一个人再强也比不过一群各有所长的人。
多头如何工作？
说白了，就是先把输入拆成多份，分别做注意力，最后再拼回去。具体流程如下：
- 线性投影：输入序列 X 分别经过 h 组不同的 W_Q、W_K、W_V 矩阵，得到 h 组 Q、K、V。
- 缩放点积注意力：每组 Q、K、V 独立计算注意力输出 head_i = Attention(Q_i, K_i, V_i)。
- 拼接与投影：将所有 head 拼接起来，再经过一个输出投影矩阵 W_O。
公式长这样：
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)你想想看，每个头其实是在不同的子空间里做注意力。比如在 MobileViT 里，我通常设 h=4 或 h=8。头数太少，表达能力不够；头数太多，计算量又上去了。
避坑指南：我曾经在项目中把头数设成 16，结果模型训练特别慢，而且效果没提升。后来发现，对于小模型（比如 MobileViT），4-8 个头就足够了。头数不是越多越好，得看模型容量。
计算复杂度分析
咱们来算笔账。假设输入序列长度为 N，每个 的维度是 d，头数为 h，每个头的维度是 d_k = d / h。
单头注意力的复杂度是 O(N² · d)。多头呢？
- 每个头的复杂度：O(N² · d_k) = O(N² · d / h)
- h 个头加起来：O(h · N² · d / h) = O(N² · d)
看到了吗？多头注意力的计算复杂度和单头是一样的。为什么？因为每个头的维度变小了，总计算量被均摊了。
但这里有个细节：实际计算量会略高。因为多了线性投影和拼接操作。不过这些开销相比注意力计算本身，可以忽略不计。
| 操作 | 单头复杂度 | 多头复杂度 |
|---|---|---|
| Q/K/V 投影 | O(N · d²) | O(N · d²)（不变） |
| 注意力计算 | O(N² · d) | O(N² · d)（不变） |
| 输出投影 | O(N · d²) | O(N · d²)（不变） |
所以，多头注意力在几乎不增加计算量的前提下，大幅提升了模型的表达能力。这也是为什么它成了 Transformer 家族的标配。
注意：虽然复杂度相同，但多头注意力的内存占用会略高。因为需要存储 h 个注意力矩阵。在 MobileViT 这种轻量级模型上，我建议 h 不超过 8，否则显存会吃紧。
SVG 流程图：多头自注意力机制
这张图把流程画得很清楚了。输入先做线性投影，然后分头计算注意力，最后拼接再投影。每一步都清晰可见。
好了，多头自注意力的核心就这些。记住：多头不是简单的重复，而是并行地在不同子空间里学习。这个思想在 MobileViT 里被用到了极致，后面我们会看到它是怎么和卷积结合的。
总结：多头自注意力用几乎相同的计算量，换来了更强的表达能力。它让模型能同时关注序列中不同位置、不同特征维度的关系。这就是为什么它这么香。
可可精品资料库