自注意力机制——Transformer 的“心脏”

说实话，我第一次接触 Transformer 的时候，最让我眼前一亮的就是自注意力机制。它不像 RNN 那样按时间步一个一个处理，也不像 CNN 那样只看局部窗口。自注意力一次就把序列里所有位置的关系都摸清了。我记得在做一个基因表达数据矫正的项目时，序列长度有几千个位点，RNN 跑起来慢得让人抓狂，换成自注意力之后，速度和效果都有了质的飞跃。

自注意力的核心思想很简单：让每个位置都去“关注”序列中其他所有位置。它通过三个可学习的权重矩阵——Query、Key、Value——来完成这个任务。你想想看，这就像你在图书馆里找书：Query 是你想查的关键词，Key 是每本书的标签，Value 是书的内容。系统根据 Query 和 Key 的匹配程度，把最相关的 Value 加权提取出来。

从公式到直觉

标准自注意力的计算分四步走，我建议你把它们记在脑子里：

  1. 线性投影：输入序列 \( X \) 分别乘以三个权重矩阵 \( W_Q, W_K, W_V \)，得到 Query、Key、Value 矩阵。
  2. 计算注意力分数：用 \( Q \) 和 \( K^T \) 做点积，得到相似度矩阵。
  3. 缩放与归一化：点积结果除以 \( \sqrt{d_k} \)（防止数值爆炸），再经过 softmax 得到注意力权重。
  4. 加权求和：用注意力权重去加权 Value，得到最终输出。

整个过程的数学表达式就是：

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

这里 \( d_k \) 是 Query/Key 的维度。嗯，这里要注意：如果不做缩放，当 \( d_k \) 很大的时候，点积结果会很大，softmax 的梯度会变得非常小——这就叫“梯度饱和”。我在一次训练蛋白质序列模型时踩过这个坑，输出全变成 NaN 了，后来一查就是忘加缩放因子。

自注意力的两个“亲戚”

很多人把自注意力（Self-Attention）和传统的注意力搞混。我在这里给你列个对比表，方便你记：

类型 Query 来源 Key/Value 来源 典型应用
自注意力 (Self-Attention) 序列自身 序列自身 Transformer 编码器、BERT
交叉注意力 (Cross-Attention) 一个序列 (如解码器) 另一个序列 (如编码器输出) 机器翻译解码器、T5
因果注意力 (Causal Attention) 序列自身 (带 mask) 序列自身 (只看到过去) GPT 等自回归模型

我们做批次效应矫正时，用的主要是自注意力——因为我们想让每个基因位点的表达量，参考其他位点的表达模式来调整。

多头注意力：让模型“多角度看问题”

单头注意力只能学到一种“关注模式”。但现实数据里，一个位点可能需要同时关注它的局部邻居、远距离调控区域、甚至同源基因。多头注意力就是为此而生的。

它把 Query、Key、Value 分成 \( 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 )

我习惯把 \( h \) 设成 8 或 12，具体看序列长度。序列越长，头数可以适当多些——但别太多，不然每个头关注的信息会重叠。我踩过的一个坑是头数设到 64，结果训练时显存暴涨，效果反而还不如 8 头。

位置编码：给序列带上“顺序感”

自注意力本身是“置换等变的”——打乱输入顺序，输出也会跟着乱。但批次效应矫正的数据中，位点是有顺序的（比如基因组坐标）。所以我们得给模型注入位置信息。

最经典的方法是正弦波位置编码：

PE(pos, 2i)   = sin( pos / 10000^(2i/d_model) )
PE(pos, 2i+1) = cos( pos / 10000^(2i/d_model) )

这个公式的好处是：不同位置之间的距离关系能被模型隐式学到。当然，现在也有很多用可学习位置编码的，但在处理超长序列时，正弦波依然是不错的选择。

实战避坑指南

我曾经在训练一个 10000 位点的基因表达矫正模型时，softmax 之前的值直接溢出了 Python 的 float16 精度。后来加了两层保护：

  • 在 softmax 之前减去最大值（注意是沿着最后一个维度），即 scores -= scores.max(dim=-1, keepdim=True).values
  • 缩放因子不要忘记乘以 1/sqrt(d_k)

如果你使用 PyTorch，可以直接调用 torch.nn.MultiheadAttention，它内置了缩放、dropout 和 mask 支持。我建议你自己动手写一个自注意力层，理解透彻后再用封装好的——我在面试实习生时发现很多人只会调包，一问细节就卡壳。

一个最小可用的自注意力代码

下面是我常用的一个简单自注意力模块，去掉注释后不到 20 行：

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

class SelfAttention(nn.Module):
    def __init__(self, d_model, d_k, d_v):
        super().__init__()
        self.W_Q = nn.Linear(d_model, d_k)
        self.W_K = nn.Linear(d_model, d_k)
        self.W_V = nn.Linear(d_model, d_v)
        self.scale = d_k ** 0.5

    def forward(self, x, mask=None):
        # x shape: (batch, seq_len, d_model)
        Q = self.W_Q(x)  # (b, n, d_k)
        K = self.W_K(x)
        V = self.W_V(x)

        scores = torch.matmul(Q, K.transpose(-2, -1)) / self.scale  # (b, n, n)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)
        attn = F.softmax(scores, dim=-1)
        out = torch.matmul(attn, V)
        return out, attn

这个模块可以直接用在你的批次效应矫正模型里。注意 mask 参数：在解码器里用它来防止看到未来的 token；在编码器里一般不用 mask（或者用 padding mask 只忽略填充位）。

一张图总结自注意力计算流程

下面这张 SVG 图展示了自注意力从输入到输出的完整流水线。我把它画得尽量清楚，方便你对照代码理解。

自注意力机制计算流程 输入 X Q = X·W_Q K = X·W_K V = X·W_V Q × Kᐪ (相似度矩阵) 除以 √d_k Softmax (注意力权重) 注意力 × V 输出 Z 输入: (batch, seq, d_model) 可选 Mask 屏蔽

从这张图你可以看到：输入 X 经过三个线性层变成 Q、K、V；Q 和 K 做点积、缩放、softmax 得到权重；然后用权重去加权 V，得到最终的输出。整个过程就是一次“全局推理”。

我个人的习惯是画这种流程图的时候，一定把“缩放”和“softmax”分开画。因为很多新手容易把这两个步骤合并理解，结果调试时找不到数值溢出的根源。你在自己实现的代码里也建议把每一步的中间结果打印出来看看——相信我，这能帮你省下不少 debug 时间。

好了，本章的自注意力机制就讲到这里。它的核心概念——QKV、缩放点积、多头、位置编码——在后续的 Transformer 架构搭建和批次效应矫正实战中都会反复用到。你不需要一下子全记住，但至少要理解“每个位置怎么去关注其他位置”这件事。下次我们就要用这些组件组装出一个能矫正批次效应的 Transformer 模型了。

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