自注意力与多头注意力

注意力机制回答一个朴素的问题：当前 token 应该「看」序列里的哪些位置？ 自注意力让序列中每个位置都与其他位置交互，权重由内容动态决定。

缩放点积注意力

对输入 $X \in \mathbb{R}^{n \times d}$，先经三个线性变换得到查询、键、值：

$Q = XW_Q,\quad K = XW_K,\quad V = XW_V$

注意力权重是查询与键的点积相似度，除以 $\sqrt{d_k}$ 防止内积过大 导致 softmax 落入饱和区：

$\mathrm{Attention}(Q, K, V) = \mathrm{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right) V$

import torch
import torch.nn.functional as F

def scaled_dot_product_attention(q, k, v, mask=None):
    d_k = q.size(-1)
    scores = q @ k.transpose(-2, -1) / (d_k ** 0.5)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float('-inf'))
    weights = F.softmax(scores, dim=-1)
    return weights @ v

为什么叫「缩放」

若 $q, k$ 各分量独立且均值为零，内积的方差随维度线性增长到 $d_k$。 不缩放时 softmax 的输入落在饱和区，梯度趋近于零，训练不稳定。

多头注意力

单头注意力只在一个子空间里聚合信息。多头把 $d_{model}$ 切成 $h$ 份 并行计算，再拼接投影回原维度：

$\mathrm{MultiHead}(Q, K, V) = \mathrm{Concat}(\mathrm{head}_1, \dots, \mathrm{head}_h)\, W_O$

其中 $\mathrm{head}_i = \mathrm{Attention}(QW_Q^i, KW_K^i, VW_V^i)$。 不同头可以学会不同的关注模式：句法关系、共指、局部依赖……

复杂度

• 计算 $QK^\top$：$\mathcal{O}(n^2 d)$
• softmax：$\mathcal{O}(n^2)$
• 加权求和：$\mathcal{O}(n^2 d)$

序列长度 $n$ 的平方项是大模型长上下文的主要瓶颈，也是 FlashAttention、稀疏注意力等优化的出发点（见第 2 章）。