深度学习中的注意力机制回顾：自注意力与多头注意力
各位好，我是你们的老朋友。今天咱们聊聊注意力机制——这玩意儿在Transformer里是核心，也是EvoFormer能搞蛋白质结构预测的基础。我记得刚入行那会儿，看到“Attention”这个词还以为是人脑注意力，后来才发现是向量之间的加权关系。好了，闲话少说，咱们先把自注意力和多头注意力捋一遍，后面讲EvoFormer才能跟得上。
自注意力：一个向量怎么“看”其他向量
自注意力的想法很简单：每个位置的输出，应该是所有位置输入的加权和。权重由当前位置与每个其他位置的“匹配度”决定。用数学说就是：给你一组向量 X，你先通过三个线性变换得到 Q（Query）、K（Key）和 V（Value）。
核心公式（Scaled Dot‑Product Attention）：
Attention(Q, K, V) = softmax( Q × Kᵀ / √dₖ ) × V其中 √dₖ 是缩放因子，防止点积太大让softmax梯度消失。dₖ 是Key向量的维度。
我做第一个翻译模型时，没加这个缩放，结果训练到一半梯度全炸了。后来加上了√dₖ，loss才开始下降。嗯，这里要特别注意。
遮罩（Mask）也很关键。在解码器里，你不能让当前位置看到未来的单词。做法就是把未来位置的注意力分数设成负无穷，softmax之后就会变成0。我习惯把遮罩矩阵跟分数矩阵做逐元素加法，像这样：
# 假设 seq_len=3，左下三角为0表示可看，右上角设为-1e9
mask = [[0, -1e9, -1e9],
[0, 0, -1e9],
[0, 0, 0]]
scores = Q @ K.T / sqrt(d_k) + mask
weights = softmax(scores, dim=-1)自注意力的复杂度是O(n²d)，n是序列长度。所以长序列时你会很痛苦——我优化过32k长度的蛋白质序列，那个矩阵占显存大到哭。但它的好处是每个位置都能直接交互，感受野是全局的。
多头注意力：为什么不只算一次
单头注意力只能学习一种匹配模式。但你可以想想，句子里的词可能同时需要关注语法、语义、位置信息等多个方面。所以论文提出把Q、K、V各分成h个头（比如h=8），每个头独立计算自注意力，再把结果拼起来投影一次。说白了就是让模型“从不同角度去看同一件事”。
我项目里踩过一个坑：头数设得太多，每个头的维度太小（比如2~4维），结果每个头都学不到什么东西。后来我按经验把每个头的维度设在32~64之间，效果才稳定下来。具体的通道数关系：
| 参数 | 含义 | 常用值 |
|---|---|---|
| dₘₒₕₑₗ | 模型总维度（输入/输出） | 512 / 768 / 1024 |
| h | 头的数量 | 8 / 12 / 16 |
| dₖ = dᵥ | 每个头的Key/Value维度 | dₘₒₕₑₗ / h（通常64左右） |
多头注意力的计算流程可以写成：
- 对输入做线性变换，得到Q、K、V（形状都是 [B, L, dₘ]）。
- 把最后一维拆成
[h, dₖ]，再转置成[B, h, L, dₖ]。 - 每个头分别计算
Attention(Qₕ, Kₕ, Vₕ)，得到 [B, h, L, dₖ]。 - 把h和dₖ合并回 [B, L, dₘ]，再做一次线性投影。
你可以在代码里用 einops.rearrange 来优雅地完成分头和合并，比reshape更不易出错。
from einops import rearrange
q = rearrange(q, 'b l (h d) -> b h l d', h=num_heads)
attn = scaled_dot_product_attention(q, k, v)
out = rearrange(attn, 'b h l d -> b l (h d)')单头 vs 多头：一张表说清楚
| 特性 | 单头注意力 | 多头注意力 |
|---|---|---|
| 表达能力 | 只能学到一种关系 | 多种子空间，并行学习不同关系 |
| 参数量 | 相对较小，但QKV矩阵维度大 | 总参数量接近（拆分后投影矩阵更小，但多了一个合并投影） |
| 计算复杂度 | O(n²dₘ) | 同样是O(n²dₘ)，拆分不增加额外计算，只是分块并行 |
| 训练稳定性 | 较易过拟合单一模式 | 更稳定，不同头可互补，提升泛化 |
| 典型应用 | 早期序列对齐、简单阅读理解 | Transformer、BERT、GPT、EvoFormer |
你想想看，如果不搞多头，你这个模型就只能用一种“眼光”去理解输入。实际任务里，同一个词可能需要同时捕捉“它是什么词性”、“它在句子里的位置”、“它跟前几个词的关系”。没有多头，这些信息就得挤在一个低维空间里，效果自然差。
注意力机制在EvoFormer中的特殊地位
EvoFormer之所以叫“进化版Transformer”，就是因为它用注意力机制处理了多序列比对（MSA）和结构信息。但核心的注意力模块还是自注意力和多头注意力的变体。只不过EvoFormer里会多了“Pairwise attention”和“Triangle attention”这些衍生品。
明白基础，你才能理解后面那些花哨的变体。别小看这一章，我曾经面试过候选人，连multi-head attention的分头逻辑都说不清楚。大家一定要亲手写一遍backward propagation（虽然现在都是用自动微分），但理解每个张量的形状变化，是深度学习的基本功。
避坑提醒：当我用fp16混合精度训练时，注意力分数 Q@K.T / sqrt(d_k) 里的除法一定要做，否则fp16范围不够，softmax会全变成NaN。另外，如果序列特别长（比如蛋白质的L=1024），建议用FlashAttention或者memory efficient attention来避免OOM。我早期的一个项目因为没有做这些优化，显存直接被撑炸了三次……后来老老实实换了FlashAttention。
本节知识导图：注意力机制的结构
下面这张图把自注意力和多头注意力的核心流程串了起来。你可以把它当作脑图，边看代码边对照。
注意图中“分头”这一步在单头模式下可以跳过。另外，EvoFormer的Pairwise attention其实是把MSA信息和结构信息作为额外的Key/Value注入，但核心的scaled dot‑product机制是一样的。
口头总结一下
自注意力让每个位置都能全局交互，但一次只能学一种关系。多头注意力把单头复制几份，让它们各学各的，然后合并。就是这么简单。你千万别被“多头”两个字吓到，拆开看就是一个分组卷积的思想。下回别人问你，你就说：“自注意力是看一遍，多头是看好几遍，每遍关注不同方面。”——印象会深刻很多。