Transformer 注意力机制详解

为什么需要注意力

在 Transformer 之前，处理序列数据的主流方案是 RNN 和 LSTM。它们按顺序逐个处理 token——第 3 个 token 必须等第 2 个 token 算完。

这个”必须等待”带来了两个致命问题：不能并行训练，长序列的早期信息会被遗忘。

2017 年，Vaswani 等人的《Attention Is All You Need》提出一个激进的想法：完全抛弃 RNN，只用注意力。这让整个序列可以同时处理，并行训练成为可能。关于论文本身的精读见《Attention Is All You Need》论文精读。

graph LR
    A[输入序列] -->|线性投影| B[Q K V 矩阵]
    B -->|Q×K^T| C[注意力分数矩阵]
    C -->|Softmax| D[注意力权重]
    D -->|×V| E[输出向量]

Self-Attention 的四步计算

以翻译句子 “The cat sat on the mat” 为例，当模型处理 “sat” 这个词时：

第一步：创建 Q、K、V

每个输入 token 的 embedding 通过三个不同的权重矩阵投影：

Q = input @ W_Q  # Query: 我在找什么？
K = input @ W_K  # Key:   我有什么信息？
V = input @ W_V  # Value: 我的实际内容是什么？

这三个矩阵的维度通常是 (seq_len, d_model)，其中 d_model = 512（原始 Transformer）。

第二步：计算注意力分数

将 Q 和 K 做点积，得到每个 token 对每个其他 token 的”关注度”：

scores = Q @ K.T  # shape: (seq_len, seq_len)
# scores[i][j] = token i 对 token j 的关注程度

第三步：缩放 + Softmax

点积结果可能很大（尤其是 d_model 较大时），这会导致 softmax 的梯度消失。除以 sqrt(d_k) 缩放后解决这个问题：

d_k = Q.shape[-1]
attention_weights = softmax(scores / sqrt(d_k), dim=-1)
# 每一行的和为 1，表示一个 token 对所有 token 的注意力分布

第四步：加权求和

output = attention_weights @ V  # shape: (seq_len, d_v)

每个 token 的输出是它关注的 token 的 Value 加权求和。

完整实现

import torch
import torch.nn.functional as F

def self_attention(X, W_Q, W_K, W_V):
    Q = X @ W_Q
    K = X @ W_K
    V = X @ W_V
    d_k = Q.shape[-1]
    scores = Q @ K.T / torch.sqrt(torch.tensor(d_k))
    weights = F.softmax(scores, dim=-1)
    return weights @ V, weights

# 示例：6 个 token，每个 512 维
X = torch.randn(6, 512)
W_Q = torch.randn(512, 64)  # d_k = 64
W_K = torch.randn(512, 64)
W_V = torch.randn(512, 64)

output, attention_map = self_attention(X, W_Q, W_K, W_V)
print(f"Output shape: {output.shape}")     # (6, 64)
print(f"Attention map shape: {attention_map.shape}")  # (6, 6)

attention_map 就是常说的”注意力热力图”——(6, 6) 的矩阵，第 i 行第 j 列表示 token i 对 token j 的关注权重。

Multi-Head Attention

单头注意力只能关注一种特征关系。多头注意力并行做多组 Q-K-V 投影，每组关注不同的特征子空间：

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model=512, n_heads=8):
        super().__init__()
        self.d_k = d_model // n_heads  # 每头的维度
        self.n_heads = n_heads
        # 合并所有头的投影矩阵
        self.W_Q = nn.Linear(d_model, d_model)
        self.W_K = nn.Linear(d_model, d_model)
        self.W_V = nn.Linear(d_model, d_model)
        self.W_O = nn.Linear(d_model, d_model)

    def forward(self, x):
        B, N, D = x.shape  # batch, seq_len, d_model
        # 投影并拆分多头
        Q = self.W_Q(x).view(B, N, self.n_heads, self.d_k).transpose(1, 2)
        K = self.W_K(x).view(B, N, self.n_heads, self.d_k).transpose(1, 2)
        V = self.W_V(x).view(B, N, self.n_heads, self.d_k).transpose(1, 2)
        # 缩放点积注意力
        scores = Q @ K.transpose(-2, -1) / math.sqrt(self.d_k)
        weights = F.softmax(scores, dim=-1)
        out = weights @ V
        # 合并多头并输出
        out = out.transpose(1, 2).contiguous().view(B, N, D)
        return self.W_O(out)

多个头各自学到不同的语言特征——有的关注语法结构，有的关注语义相似度，有的关注位置邻近性。

为什么注意力机制如此重要

Attention 不仅解决了并行训练问题，它还让模型具有了可解释性——你可以可视化注意力权重，看到模型在做预测时”关注”输入的哪个部分。这对于调试和信任 AI 系统至关重要。

GPT 系列的每一代都建立在 Transformer 架构之上，详见 GPT 系列模型演进史。推理阶段的效率优化也是注意力机制的延伸，参考 LLM 推理优化。