3、自注意力机制：Transformer的核心，QKV矩阵的诞生与计算流程

自注意力机制，说白了就是让模型学会“自己看自己”。

你想想看，当你在读一句话时，眼睛会不自觉地来回扫视，把当前词和前后文联系起来。比如“那只猫追到了老鼠，它很开心”——这个“它”指的是猫还是老鼠？人一眼就能判断，但机器不行。

自注意力要解决的就是这个问题。它让每个词都能“看到”序列里的所有其他词，然后决定该关注谁、忽略谁。Transformer 能火，靠的就是这个机制。

3.1 从“全连接”到“自注意力”

在 RNN 时代，信息是串行传递的。第 t 步只能看到前一步的隐藏状态，长距离依赖根本抓不住。我记得 2017 年做机器翻译项目时，用 LSTM 处理超过 30 个词的句子，效果就开始崩了——前面的主语到后面就“失忆”了。

自注意力的思路完全不同：

  • 并行计算：所有位置同时处理，不依赖时间步
  • 全局视野：每个位置直接和所有位置交互
  • 动态权重：关注哪些位置，由数据自己决定

核心洞察：自注意力本质上是“内容寻址”的加权求和。每个词根据自身内容，去查询其他词的内容，然后按相似度聚合信息。

3.2 QKV 矩阵：三剑客的诞生

QKV 这三个字母，我刚开始学的时候也觉得玄乎。其实没那么复杂——它们就是三个不同的线性投影。

假设输入序列有 n 个 ，每个 的向量维度是 d。我们做三件事：

  1. Q（Query，查询）：当前词“想找谁”
  2. K（Key，键）：其他词“我是谁”
  3. V（Value，值）：其他词“我有什么信息”

Q 和 K 做点积，得到注意力分数。分数越高，说明当前词越该关注那个位置。然后用 softmax 归一化，再和 V 加权求和。

嗯，这里要注意：Q、K、V 都是从同一个输入 X 通过不同的权重矩阵 WQ、WK、WV 算出来的。这就是“自”的含义——自己查自己。

# 伪代码：QKV 计算
import numpy as np

def compute_qkv(X, W_q, W_k, W_v):
    """
    X: (n, d) 输入序列
    W_q, W_k, W_v: (d, d_k) 权重矩阵
    """
    Q = np.dot(X, W_q)  # (n, d_k)
    K = np.dot(X, W_k)  # (n, d_k)
    V = np.dot(X, W_v)  # (n, d_v)
    return Q, K, V

个人习惯：我在实际项目中，通常把 d_k 设为 d / h，其中 h 是注意力头数。这样多头拼接后维度刚好对齐，省去额外的投影层。

3.3 注意力分数的计算流程

有了 Q、K、V，接下来就是标准的四步走：

步骤 操作 公式 输出形状
1 Q 与 K 点积 S = Q · KT (n, n)
2 缩放 S' = S / √dk (n, n)
3 Softmax 归一化 A = softmax(S') (n, n)
4 加权求和 Z = A · V (n, dv)

为什么要缩放？我踩过这个坑。当 dk 比较大时，点积的结果会很大，softmax 的梯度会变得非常小——说白了就是“梯度消失”。除以 √dk 能把方差拉回到 1 左右，训练就稳了。

我曾经在一个视觉语言模型项目中，忘了加缩放因子，结果训练 loss 死活降不下去。排查了两天才发现是注意力分数太大，softmax 几乎变成了 one-hot。加了缩放后，一个 epoch 就收敛了。

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

单头注意力有个问题：它只能学习一种注意力模式。但现实中，一个词可能因为不同原因关注不同位置。比如“苹果”这个词，在“苹果很好吃”和“苹果发布了新手机”里，关注的上下文完全不同。

多头注意力就是把 Q、K、V 分别投影到 h 个不同的子空间，每个头独立计算注意力，最后拼接起来：

# 多头注意力伪代码
def multi_head_attention(X, h, d_model):
    d_k = d_model // h
    
    # 每个头独立投影
    heads = []
    for i in range(h):
        Q_i = linear(X, W_q_i)  # (n, d_k)
        K_i = linear(X, W_k_i)  # (n, d_k)
        V_i = linear(X, W_v_i)  # (n, d_k)
        
        # 单头注意力
        Z_i = scaled_dot_product_attention(Q_i, K_i, V_i)
        heads.append(Z_i)
    
    # 拼接所有头
    Z = concat(heads, axis=-1)  # (n, d_model)
    
    # 输出投影
    Z = linear(Z, W_o)  # (n, d_model)
    return Z

我个人习惯用 h=8 或 h=12。头数太少，表达能力不够；头数太多，每个头分到的维度太少，反而学不到东西。8 到 16 是个比较稳妥的范围。

3.5 自注意力的完整流程图

下面这张图把整个流程串起来了。从输入 X 开始，到 QKV 的生成，再到注意力计算，最后输出 Z：

自注意力机制完整流程图 自注意力机制完整计算流程 输入序列 X Q = X·WQ K = X·WK V = X·WV S = Q · KT S' = S / √dk A = softmax(S') Z = A · V 多头：重复 h 次，拼接后投影 dk = dmodel / h

3.6 实际项目中的避坑指南

做自注意力的时候，有几个地方特别容易出问题：

  • 显存爆炸：注意力矩阵是 (n, n) 的，n=1024 时还好，n=4096 时就是 16M 个元素。我处理长序列时，经常要用到 FlashAttention 或者分块计算。
  • 数值稳定性：softmax 之前一定要做缩放。另外，如果序列长度变化很大，建议用 mask 把 padding 位置置为 -inf，防止无效位置参与计算。
  • 初始化：WQ、WK、WV 的初始化要小心。我习惯用 Xavier 初始化，方差设为 1/d，这样前向传播时注意力分数的方差能保持在合理范围。

小技巧：调试注意力机制时，可以把注意力矩阵 A 可视化出来。如果发现某一行几乎全是均匀分布，说明这个位置没学到有效的注意力——可能是梯度消失了，也可能是维度设置不合理。

3.7 自注意力的计算复杂度

最后聊一下复杂度。自注意力的计算量主要在 Q·KT 这一步，复杂度是 O(n²·d)。n 是序列长度，d 是维度。

操作 时间复杂度 空间复杂度
QKV 投影 O(n·d²) O(n·d)
注意力分数 O(n²·d) O(n²)
Softmax O(n²) O(n²)
加权求和 O(n²·d) O(n·d)

看到没？n² 是瓶颈。这也是为什么后来出现了各种高效注意力——Linformer、Performer、FlashAttention 等，都是在想办法把 n² 降下来。

不过对于视觉语言模型来说，图像 patch 的数量通常不会太大（比如 196 或 256），所以标准自注意力完全够用。我在做图文检索项目时，用 12 头、d=768 的配置，效果已经很好了。

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