3. Transformer基础：自注意力机制、多头注意力

聊到Transformer，最核心的东西就两个：自注意力和多头注意力。很多人一上来就被各种公式吓住了，其实说白了，它要解决的就是一句话：让模型在生成分子序列时，知道该「看」输入的哪些部分。

我做AI制药这些年，最深的体会是：小分子和自然语言本质上是同构的——一个是一串原子，一个是一串token。所以理解自注意力，就是理解Transformer为什么能折腾出那么多新分子。

3.1 自注意力是什么？

自注意力（Self-Attention）的核心操作，就是对输入序列中每个位置，根据其他所有位置的信息，重新计算这个位置的表示。

举个例子，假设我们有一个分子序列：C C = O （丙酮的简化表示）。

  • 看到第一个C，需要知道它和第二个C是单键还是双键。
  • 看到第二个C，需要留意后面有个O双键。
  • 看到O，要回头确认它前面到底连了几个C。

自注意力就是干这个的——每个位置去「查询」所有位置，然后加权融合信息。

3.1.1 核心计算：Q、K、V

每个输入token会生成三个向量：Query（查询）、Key（键）、Value（值）。

  • Query：问我当前这个位置「想找谁」。
  • Key：每个位置回答「我有什么信息」。
  • Value：每个位置真正要贡献的「内容」。

注意力分数 = Query 与所有 Key 的点积，然后 softmax 得到权重，最后用权重加权求和所有 Value。

公式（简化记忆级）

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

其中 √d_k 是为了防止点积过大导致softmax梯度消失。d_k 是Key的维度。

我的一点经验：刚开始调参时，我总喜欢把d_k设得很大，以为信息多。结果发现注意力分数直接炸了，softmax输出几乎变成one-hot。后来才明白，scale是必须的。如果你像我一样踩过这个坑，记得检查一下这里的缩放系数。

3.2 多头注意力（Multi-Head Attention）

一个注意力头可能只学到了「这个C和那个O关系很强」，但分子里还有键角、环结构、手性等等信息。怎么办？多头注意力就是并行跑多个自注意力头，每个头关注不同的子空间。

3.2.1 多头怎么组合？

流程很简单：

  1. 将输入的Q、K、V分别线性投影到h个不同的低维空间（每个头维度 = d_model / h）
  2. 每个头独立计算自注意力
  3. 将所有头的输出拼接起来，再做一个线性变换
MultiHead(Q,K,V) = Concat(head_1, head_2, ..., 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、W_i^V 是每个头的投影矩阵，W_O 是输出投影矩阵。

我曾经踩过的坑：头数太多（比如64个头），每个头的维度只剩个位数，注意力几乎退化成了随机权重。我个人习惯是：d_model=512时用8个头，每个头64维；如果序列很长（比如蛋白质），12~16个头也够用，再大就不太行了。

3.3 自注意力在分子生成中的关键作用

在分子序列生成中，自注意力帮模型做三件事：

  1. 长程依赖捕捉：例如预测一个遥远侧链对活性位点的影响。
  2. 局部结构构建：例如一个六元环内的C-C键顺序。
  3. 消歧义：比如同一个原子符号在不同环境下（芳香环 vs 脂肪链）意义不同。

你看，这跟自然语言里的“词义消歧”是不是很像？

3.3.1 自注意力的矩阵维度

假设输入序列长度 = n，每个token的表示维度 = d_model。

组件维度变化
输入 X(n, d_model)
Q、K、V（每个头）(n, d_k) 或 (n, d_v)
注意力矩阵 (Q×K^T)(n, n) — 空间复杂度O(n²)
单个头输出(n, d_v)
多头拼接后(n, h×d_v) = (n, d_model)
最终输出(n, d_model)

这里有个明显的瓶颈：注意力矩阵是n×n。分子序列如果很长（比如几千个原子），显存直接炸了。我遇到过几次OOM，后来要么截断序列，要么换用稀疏注意力。

3.4 自注意力的变体：我们为什么不用原始版本？

原始Transformer里的自注意力是「全局全连接」的——每个位置看所有位置。但在分子生成里，很多长程信息其实没必要（比如两个相距50个原子的碳原子，一般不会直接成键）。

常见的改进：

  • 局部注意力：每个位置只看前后窗口大小内的token。
  • 稀疏注意力：预先定义一些连接模式（比如对角线、膨胀模式）。
  • 线性注意力：用核技巧把O(n²)降到O(n)。

实际建议：如果你刚开始做分子生成，先用全局注意力把基础跑通。后面再根据你的序列长度和GPU显存去换稀疏版本。我在项目里常用的是「局部+全局稀疏」的混合模式，效果和速度都不错。

自注意力与多头注意力结构图 自注意力与多头注意力核心流程 输入序列 X (n, d_model) 线性投影 Q 线性投影 K 线性投影 V 注意力分数矩阵 Q×K^T / √d_k → softmax 加权求和 → 自注意力输出 (n, d_v) 多头注意力 (h=4 示例) 头1: Q,K,V 头2: Q,K,V 头3: Q,K,V 头4: Q,K,V 拼接所有头 → 线性投影 → 多头输出 (n, d_model) 图注： • 左侧：标准自注意力 —— 输入经三个线性投影得到Q/K/V，计算注意力分数后加权求和。 • 右侧：多头注意力 —— 将Q/K/V投影到h个子空间并行计算，最后拼接并线性变换。 • 注意：每个头的输出维度 = d_model / h，保证总参数不变。

3.5 动手实现一个小例子

用PyTorch写一个自注意力模块，也就几十行。我建议你自己敲一遍，比看十遍文章都管用。

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

class SelfAttention(nn.Module):
    def __init__(self, d_model, d_k=None, d_v=None):
        super().__init__()
        self.d_model = d_model
        self.d_k = d_k if d_k is not None else d_model
        self.d_v = d_v if d_v is not None else d_model
        self.W_q = nn.Linear(d_model, self.d_k, bias=False)
        self.W_k = nn.Linear(d_model, self.d_k, bias=False)
        self.W_v = nn.Linear(d_model, self.d_v, bias=False)

    def forward(self, x, mask=None):
        # x: (batch, seq_len, d_model)
        Q = self.W_q(x)
        K = self.W_k(x)
        V = self.W_v(x)
        scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.d_k ** 0.5)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        attn_weights = F.softmax(scores, dim=-1)
        output = torch.matmul(attn_weights, V)
        return output, attn_weights

# 用一下
seq = torch.randn(2, 10, 512)  # batch=2, seq_len=10, d_model=512
attn = SelfAttention(d_model=512)
out, weights = attn(seq)
print(out.shape)  # (2, 10, 512)

关键点：

  • 线性层不加bias，因为后续有LayerNorm。
  • mask用来屏蔽padding位置（在分子生成中，不同长度的序列需要padding）。
  • 输出形状和输入一致，这是Transformer残差连接的前提。

然后多头注意力就是把这个SelfAttention跑h次，再把结果拼起来。我习惯用一个包装类：

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, h=8):
        super().__init__()
        assert d_model % h == 0, "d_model must be divisible by h"
        self.h = h
        self.d_k = d_model // h
        self.heads = nn.ModuleList([
            SelfAttention(d_model, d_k=self.d_k, d_v=self.d_k)
            for _ in range(h)
        ])
        self.W_o = nn.Linear(d_model, d_model, bias=False)

    def forward(self, x, mask=None):
        outputs = []
        for head in self.heads:
            out, _ = head(x, mask)
            outputs.append(out)       # 每个 out: (batch, seq, d_k)
        concat = torch.cat(outputs, dim=-1)  # (batch, seq, d_model)
        return self.W_o(concat)

mha = MultiHeadAttention(d_model=512, h=8)
out_mha = mha(seq)
print(out_mha.shape)  # (2, 10, 512)

小技巧：实际项目中不要像我这样一个一个for循环头，效率太低。应该用合并的线性投影 + reshape + transpose并行算。但为了理解原理，上面的写法更清晰。

3.6 一些避坑和个人习惯

  • 关于scale：√d_k这个分母千万别忘了。我见过有人实现时写成了√d_model，结果注意力全糊了。
  • 关于mask：分子生成通常是自回归的，需要因果mask（每个位置只看之前的位置）。实现时注意一下上三角矩阵的掩码。
  • 关于头数选择：如果d_model=256，头数选4或8；如果d_model=1024，选16。不是越多越好。
  • 关于初始化：建议用xavier均匀初始化或kaiming正态初始化。我用默认的线性层初始化有时会导致训练初期注意力全均分。

嗯，到这里自注意力和多头注意力的核心你已经掌握了。后面我们会用它搭出完整的Transformer编码器和解码器，然后就能生成分子了。

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