4. Transformer 基础：自注意力机制与多头注意力

你问我 Transformer 里最核心的是什么？我会毫不犹豫地说：自注意力机制。没有它，Transformer 跟 RNN 就没多大区别。我个人习惯把自注意力理解为“模型自己跟自己开会”——每个 token 都去问其他 token：“咱俩有关系吗？关系多强？” 然后根据答案重新调整自己的表示。

但光开会还不够。会场只有一个？信息太单一了。所以就有了多头注意力——开多个分会场，每个会场关注不同的关系，最后汇总讨论结果。说白了，就是用多个不同的视角去理解同一个句子。

嗯，现在我们从最底层的原理讲起。我在做分子生成项目时，发现很多同学把注意力机制当黑盒用，一遇到长序列就崩。你理解了底下这堆矩阵运算，才能在分子生成里跑得稳。

图4-1：自注意力 + 多头注意力全景流程 输入 X 线性变换 Wq / Wk / Wv Q K V 缩放点积注意力 Q * K^T / sqrt(dk) → Softmax → * V 分 h 头：每个头独立做 Attention 拼接所有头的结果 线性投影 → 输出 可选掩码 （对应位置填 -∞）

4.1 自注意力机制：让每个 token 都能看到全局

想象你有 4 个 token："The cat sat down"。没有自注意力时，每个词只知道自己的 embedding。有了自注意力，“sat”就能知道主语是“cat”，动作是“down”。这是怎么做到的？

先定义三个矩阵：Q（查询）、K（键）、V（值）。每个 token 都会生成自己的 Q、K、V。怎么生成？很简单——拿 token 的原始 embedding 乘三个不同的权重矩阵。

拿分子生成打个比方。假设我们有分子 SMILES 序列 CCO（乙醇）。每个原子 symbol 经过 embedding 后，变成一个向量 x。然后：

q = x * Wq
k = x * Wk
v = x * Wv

这三个 W 矩阵是模型要学的参数。我刚开始做分子项目时，习惯把 dk 设成 64，因为分子序列不长，64 足够捕捉化学键信息和官能团之间的依赖。

4.2 注意力分数的计算：点积、缩放与 Softmax

有了 Q、K、V，我们就要计算注意力分数了。核心公式就一句话：

注意力输出 = Softmax( (Q × KT) / sqrt(dk) ) × V

我给你拆开来看：

  1. Q × KT：每个 token 的 Q 和其他所有 token 的 K 做点积。点积越大，说明这两个 token 越“相关”。
  2. 除以 sqrt(dk)：防止点积数值太大，导致 Softmax 后的梯度消失。为什么是 sqrt？嗯，假设 Q 和 K 是独立同分布的均值为 0 方差为 1 的随机变量，点积的方差就是 dk。除以 sqrt(dk) 让方差变回 1，稳定训练。
  3. Softmax：把分数归一化成概率，和为 1。
  4. 乘以 V：用概率加权求和，得到新的表示。

来看一段简化的 PyTorch 计算代码：

import torch
import torch.nn.functional as F

def scaled_dot_product_attention(Q, K, V, mask=None):
    """
    Q, K, V: [batch_size, seq_len, d_k]
    mask: [batch_size, seq_len, seq_len], 可选
    """
    d_k = Q.size(-1)
    # [batch, seq_len, seq_len]
    scores = torch.matmul(Q, K.transpose(-2, -1)) / (d_k ** 0.5)
    
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
    
    attn_weights = F.softmax(scores, dim=-1)
    output = torch.matmul(attn_weights, V)
    return output, attn_weights
注意看 mask 那行：用 -1e9 填掉你想屏蔽的位置。在分子生成中，我们常用 padding mask 忽略填充位，也用 causal mask 防止 decoder 看到未来 token。我曾经因为忘记加 causal mask，训练时 loss 降到 0.01，推理时生成的全是垃圾——那是典型的“数据泄露”。

4.3 多头注意力：并行开会，多视角理解

单头注意力把整个序列的注意力分数压缩成一个加权和。但缺陷很明显——一个头只能关注一种类型的依赖。比如在分子中，C-C 单键和 C=C 双键同时存在，一个头很难同时捕捉键长和键角信息。

多头注意力就是解决方案。你把 Q、K、V 切成 h 份，每一份独立做注意力，最后拼起来再投影回去。

具体步骤是这样：

  1. 输入 X 分别过 WQ、WK、WV，得到 Q、K、V。
  2. 把 Q、K、V 按最后一维切分成 h 个头。比如 dmodel=512，h=8，每个头的维度 dk=64。
  3. 每个头独立做缩放点积注意力。
  4. 把所有头的结果拼接起来，再过一个线性投影 WO。

用代码实现更直观：

class MultiHeadAttention(torch.nn.Module):
    def __init__(self, d_model, n_heads, dropout=0.1):
        super().__init__()
        assert d_model % n_heads == 0
        
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads
        
        self.W_q = torch.nn.Linear(d_model, d_model)
        self.W_k = torch.nn.Linear(d_model, d_model)
        self.W_v = torch.nn.Linear(d_model, d_model)
        self.W_o = torch.nn.Linear(d_model, d_model)
        
        self.dropout = torch.nn.Dropout(dropout)
        
    def forward(self, Q, K, V, mask=None):
        batch_size = Q.size(0)
        
        # 1. 线性投影 + 分头
        Q = self.W_q(Q).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
        K = self.W_k(K).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
        V = self.W_v(V).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
        
        # 2. 每个头独立做 attention
        attn_output, attn_weights = scaled_dot_product_attention(Q, K, V, mask)
        
        # 3. 拼接所有头
        attn_output = attn_output.transpose(1, 2).contiguous()
        attn_output = attn_output.view(batch_size, -1, self.d_model)
        
        # 4. 线性投影输出
        output = self.W_o(attn_output)
        return output, attn_weights
我个人习惯把 n_heads 设成 8，d_model 设成 256 或 512。经验是：d_model 能被 n_heads 整除，而且单头维度 d_k 最好不小于 16。太小了每个头学不到什么信息。做分子生成时序列长度一般不超过 200，n_heads=8 够用。

4.4 为什么多头比单头强？

我们用一个表格来对比：

维度 单头注意力 多头注意力
Q / K / V 数量 1 组 h 组（并行）
计算量 低 略高（但可并行）
表达能力 只能捕捉一种模式 捕捉 h 种不同模式
对长距离依赖 弱，容易注意力涣散 鲁棒，多个头互补聚焦
实际应用 少量简单任务 几乎所有 SOTA Transformer

你知道最有趣的是什么吗？我曾经把训练好的 8 头注意力可视化，发现有的头专门关注“当前的 token 和下一个 token”的关系（类似 N-gram），有的头则关注远程的官能团——这正好对应了分子结构中的局部键合和远程空间效应。

注意：多头注意力不是多副本，是切分维度。很多人误以为 h=8 就是跑 8 次完整的注意力，然后取平均。不对！是先把 d_model 切成 8 段，每段 64 维，并行跑小维度的注意力，最后拼起来。这样参数量和计算量几乎是单头注意力的 h 倍——但因为你把维度拆小了，总量不变。

4.5 自注意力的计算复杂度

最后提一嘴复杂度。自注意力的计算量跟序列长度 L 的平方成正比，O(L² · dk)。为什么？因为每个 token 都要跟其他 L 个 token 算点积，所以是 L × L = L²。

这有什么实际影响？我做过一个分子生成实验，序列长度从 100 到 800，单头的显存占用从 2GB 暴涨到 32GB。如果序列太长，要考虑用 sparse attention 或者 local attention，但这些是后话了。

好在分子 SMILES 序列通常不超过 200 个 token，L² 完全扛得住。如果你做蛋白质多序列，那才真要头疼。

← 上一章 📖 返回目录 下一章 →