多头注意力（MHA）：从公式到实现

先只看一次加权求和：两个位置提供的 value 是 [2, 0] 和 [0, 4]，注意力权重是 0.75 和 0.25，读出的结果就是 [1.5, 1]。Attention 的计算从这里展开：权重怎么由 Q、K 得到，哪些位置要被遮住，多组这样的计算又怎样拼成多头注意力（MHA）。

核心计算：匹配并聚合信息

每个位置拿着自己的问题（query）去问所有位置的索引（key）。对得越上，就从那个位置的内容（value）里取越多。取回来的是一个加权平均。

多头会使用不同的投影，各算一份加权结果，再拼起来。有些头可能学到位置邻近或指代关系，但这些分工不是预先指定的，也不保证每个头都对应一个人能命名的功能。

Attention Matrix 表示什么

这里先讲 self-attention：$Q$、$K$、$V$ 来自同一个输入 $X$ 的三组可学习线性投影。Cross-attention 则可以从另一条序列取 K/V，详见原版 Transformer。

$$ \begin{gathered} Q=XW_Q\\ K=XW_K\\ V=XW_V \end{gathered} $$

它们不是三个单独的维度，也没有人为规定好的语义。本例令 $d_v=d_k$：序列有 $T$ 个 token 时，$Q,K,V$ 都是 $(T,d_k)$；一般情况下 V 的宽度可以不同。公式赋予了它们不同的计算角色：Q/K 用来算权重，V 是之后被加权汇总的向量。

完整计算只有下面五步。

1. 每两个 token 算一个标量分数

$$ S=\frac{QK^\top}{\sqrt{d_k}},\qquad S_{ij}=\frac{\mathbf q_i^\top\mathbf k_j}{\sqrt{d_k}} $$

$S$ 的形状是 $(T,T)$，这就是 attention score matrix。第 $i$ 行表示“当前位置 $i$ 想从哪里取信息”，第 $j$ 列表示“候选（candidate）来源位置 $j$”。其中每个格子 $S_{ij}$ 只是一个标量；这里还没有使用 $V$。

2. Decoder-only 模型先把未来位置遮住

GPT 在位置 $i$ 预测下一个 token 时，只允许使用位置 $i$ 及其左边的信息。因果 mask 写成：

$$ M_{ij}= \begin{cases} 0, & j\le i\\ -\infty, & j>i \end{cases} $$

三个 token 的分数矩阵会变成：

$$ S+M= \begin{bmatrix} s_{11} & -\infty & -\infty\\ s_{21} & s_{22} & -\infty\\ s_{31} & s_{32} & s_{33} \end{bmatrix} $$

mask 不是让模型只看前一个 token，而是让它看到自己和之前的所有 token，同时看不到未来。训练时所有位置会并行计算；如果不遮住右上角，前面的位置就能直接读取后面的正确答案，训练目标会发生信息泄漏。双向 encoder 通常没有 causal mask，但仍可能使用 padding mask。

3. 对每一行做 softmax

$$ \begin{gathered} A=\operatorname{softmax}_{j}(S+M) \end{gathered} $$

Softmax 沿列索引 $j$ 进行，因此每一行满足：

$$ \sum_j A_{ij}=1 $$

由于 $e^{-\infty}=0$，被 mask 的位置权重为 0。这里要求每行至少有一个允许的位置，且允许位置的分数有限；整行都屏蔽时，普通 softmax 没有有效的归一化分母，不能指望它自动返回全零。$A$ 才是 attention weight matrix：第 $i$ 行给出了当前位置从允许位置各取多少信息。

4. 用这一行权重汇总所有 value

$$ O=AV,\qquad \mathbf o_i=\sum_j A_{ij}\mathbf v_j $$

所以 $QK^\top$ 只负责决定“权重是多少”，真正被取出并混合的是 $V$。attention matrix 的一个格子是标量，而输出 $\mathbf o_i$ 是一个 $d_k$ 维向量。

5. 多个头各算一套，再合并

每个头都有自己的投影、score matrix 和 attention weights，因此可以学到不同的匹配方式。所有头的输出先拼接，再经过 $W_O$ 投影回模型维度。

整条数据流可以压缩成一句：

$$ \begin{gathered} S=QK^\top/\sqrt{d_k},\\ A=\operatorname{softmax}_{\rm row}(S+M),\\ O=AV. \end{gathered} $$

也就是：三个投影产生 Q/K/V；Q 和 K 生成 token 两两之间的权重；mask 删除不允许的信息路径；每行 softmax 归一化；最后用这些权重对 V 求加权和。

MHA 里的 softmax 确实提供了非线性，但它主要在 token 维度上决定“和谁通信”。FFN 的激活函数（activation function）则在每个 token 的特征维度上做非线性变换；两者作用不同。

单个头在算什么

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

拆到单个 query $\mathbf{q}_i$ 看：

$$ \begin{gathered} s_{ij}=\mathbf q_i^\top\mathbf k_j/\sqrt{d_k},\\ \alpha_{ij}=\frac{\exp(s_{ij})}{\sum_{j\prime}\exp(s_{ij\prime})},\\ \mathbf o_i=\sum_j\alpha_{ij}\mathbf v_j. \end{gathered} $$

在没有 attention dropout 时，每行权重非负且和为 1，所以单个头在输出投影 $W_O$ 之前的结果是 value 的凸组合。开头的 [1.5, 1] 就是一个例子。别把这个结论套到整个 Transformer block：dropout、输出投影和残差连接（residual connection）都会改变它；权重本身也随输入变化，并不是固定的平均。

为什么除以 $\sqrt{d_k}$

先看它改变了什么。假设一个头的维度为 64，某个 query 对两个 key 的点积（inner product）分数是 [8, −8]：

不缩放时，几乎全部权重都落在第一个位置上。缩放以后，第二个位置还能参与输出。这不是说权重越平均越好，而是不希望仅仅因为维度变大，softmax 就过早变得很尖。

这里要分清两件事：除以一个正数会改变分数差，也会改变权重；减去同一个数不改变分数差，只是让指数运算更稳。[1000, 1000] 仍然应该得到 [0.5, 0.5]，不能笼统地说“输入大就会饱和”。

手写时用 head dimension $d_k$，不是整个模型的 $d_\text{model}$。如果模型还有 QK normalization 或可学习的温度，要以它的实际公式为准；缩放能缓解一个问题，不是训练稳定性的保证。

手写反向：另一个位置的梯度也要算进去

两个位置各分到一半权重：$\alpha=(0.5,0.5)$。如果 loss 只取第一个权重，增加第一个分数会让 loss 上升；增加第二个分数会抢走第一个位置的权重，让 loss 下降。因此，对两个分数的梯度是 $(0.25,-0.25)$，而不是 $(0.25,0)$。

import numpy as np

def softmax_backward(probabilities, upstream):
    weighted_mean = np.sum(probabilities * upstream, axis=-1, keepdims=True)
    return probabilities * (upstream - weighted_mean)

probabilities = np.array([0.5, 0.5])
assert np.allclose(softmax_backward(probabilities, np.array([1.0, 0.0])),
                   [0.25, -0.25])
assert np.allclose(softmax_backward(probabilities, np.ones(2)), [0.0, 0.0])

为什么使用多头：目的不是增加维度

$$ \begin{gathered} \text{head}_i=\\ \operatorname{Attention}(XW_i^Q,XW_i^K,XW_i^V),\\ \operatorname{MultiHead}=\\ \operatorname{Concat}(\text{head}_1,\ldots,\text{head}_h)W^O. \end{gathered} $$

单头对所有 value 通道只用一套注意力权重。把固定宽度拆成 $h$ 个头，可以在相近的投影和矩阵乘法成本下得到多套这样的关系；不过分别保存各头的注意力矩阵、启动各自的 kernel，也可能带来额外开销。下面三点说明其中的取舍。

从零实现时最容易错的六处

实现时先检查这六处，比盯着最终 loss 猜原因更直接：

第 1 条为什么必须这样。 投影输出是 (B, T, d_model)，其中 d_model 这一维是 $h$ 个头首尾相接排列的。所以要先把最后一维拆成 (H, dh)，再把 H 挪到前面。直接 view(B,H,T,dh) 会横跨时间维乱切，得到的每个「头」是一堆不相干位置的碎片——形状对，数值全错，而且不报错。

用两个分数检查 mask。 [0, 0] 的 softmax 是 [0.5, 0.5]。只允许第一个位置时，先屏蔽得到 [1, 0]；softmax 后只置零得到 [0.5, 0]，输出就缩小了。如果再除以剩余权重之和，数学上也能得到 [1, 0]，但有限精度下可能遇到剩余权重下溢，所以通常在 softmax 前屏蔽。下面的教学代码会拒绝整行被屏蔽的输入，而不是悄悄产生 NaN。

实验：让手写版本与库实现对上

import numpy as np

def softmax(logits, axis=-1):
    values = np.asarray(logits)
    if values.dtype.kind not in "biuf":
        raise ValueError("Expected real logits")
    converted = values.astype(np.result_type(values.dtype, np.float32), copy=False)
    if values.dtype.kind in "iu" and any(
            int(original) != int(floating) for original, floating in zip(values.flat, converted.flat)):
        raise ValueError("Integer logits must be exactly representable in the computation dtype")
    values = converted
    if np.any(np.isnan(values) | np.isposinf(values)):
        raise ValueError("Logits must be finite or -inf for masked positions")
    if np.any(np.all(np.isneginf(values), axis=axis)):
        raise ValueError("Each query must have at least one allowed key")
    shifted = values - np.max(values, axis=axis, keepdims=True)
    weights = np.exp(shifted)
    return weights / np.sum(weights, axis=axis, keepdims=True)


def attention(q, k, v, mask=None):
    """q,k,v: (B, H, T, d_head)   mask: True = 屏蔽"""
    d = q.shape[-1]
    scores = q @ k.swapaxes(-2, -1) / np.sqrt(d)      # (B, H, Tq, Tk)
    if mask is not None:
        scores = np.where(mask, -np.inf, scores)      # softmax 之前
    w = softmax(scores, axis=-1)
    return w @ v, w

class MultiHeadAttention:
    def __init__(self, d_model, n_head, seed=0):
        assert d_model % n_head == 0
        rng = np.random.default_rng(seed)
        s = 1.0 / np.sqrt(d_model)
        self.n_head, self.d_head = n_head, d_model // n_head
        self.Wq, self.Wk, self.Wv, self.Wo = (
            rng.uniform(-s, s, (d_model, d_model)) for _ in range(4))

    def split(self, x):                               # (B,T,C) -> (B,H,T,dh)
        b, t, _ = x.shape
        return x.reshape(b, t, self.n_head, self.d_head).transpose(0, 2, 1, 3)

    def __call__(self, x, mask=None):
        q, k, v = self.split(x @ self.Wq), self.split(x @ self.Wk), self.split(x @ self.Wv)
        out, w = attention(q, k, v, mask)
        b, h, t, dh = out.shape
        out = out.transpose(0, 2, 1, 3).reshape(b, t, h * dh)   # 合头
        return out @ self.Wo, w

def causal_mask(t):
    return np.triu(np.ones((t, t), dtype=bool), k=1)  # 严格上三角 = 未来

import math, torch, torch.nn as nn

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, n_head, bias=False):
        super().__init__()
        assert d_model % n_head == 0
        self.n_head, self.d_head = n_head, d_model // n_head
        self.wq = nn.Linear(d_model, d_model, bias=bias)
        self.wk = nn.Linear(d_model, d_model, bias=bias)
        self.wv = nn.Linear(d_model, d_model, bias=bias)
        self.wo = nn.Linear(d_model, d_model, bias=bias)

    def forward(self, x, mask=None):
        B, T, C = x.shape
        split = lambda p: p(x).view(B, T, self.n_head, self.d_head).transpose(1, 2)
        q, k, v = split(self.wq), split(self.wk), split(self.wv)

        att = (q @ k.transpose(-2, -1)) / math.sqrt(self.d_head)
        if mask is not None:
            att = att.masked_fill(mask, float("-inf"))
        if torch.isneginf(att).all(dim=-1).any():
            raise ValueError("Each query must have at least one allowed key")
        att = att.softmax(dim=-1)

        y = (att @ v).transpose(1, 2).contiguous().view(B, T, C)   # 合头
        return self.wo(y), att

def causal_mask(t, device=None):
    return torch.triu(torch.ones(t, t, dtype=torch.bool, device=device), diagonal=1)

import torch.nn.functional as F

output = F.scaled_dot_product_attention(
    q, k, v,
    attn_mask=None if mask is None else ~mask,
    dropout_p=0.0,
)

../code/attention_torch.py 把四种实现跑在同一组权重上并互相对答案：

  from-scratch vs F.scaled_dot_product_attention : 1.19e-07
  from-scratch vs nn.MultiheadAttention          : 1.19e-07
  from-scratch torch vs pure NumPy               : 1.19e-07
  every attention row sums to 1                  : 1.19e-07
  weight on any future position                  : 0.00e+00

能让从零实现的版本和 nn.MultiheadAttention 对上，比仅仅「能跑」更重要。对不上时定位差异，本身就是有效的调试练习。nn.MultiheadAttention 把 $W_Q, W_K, W_V$ 存成一个拼接的 in_proj_weight，而 nn.Linear 存的是 $(out, in)$，因此搬到 NumPy 时需要转置。

但一次输出相同，还不能说明迁移完了。可以按下面的顺序查：

完整脚本中的 library_version 会新建一个模块并复制参数，只用于对照；它的参数梯度不会回到原来的 mha。仓库测试在三条 PyTorch 路径上覆盖了 bias、float32 / float64、逐头 mask 和梯度检查；NumPy 对照仍是无 bias 的小例子。这里都关闭 dropout，不是速度或显存测试。接口约定见 PyTorch 2.14 MultiheadAttention。

接上 KV cache，mask 要从哪里开始？

前面画的是正方形：5 个 query 对 5 个 key。接上缓存后，假设已经有位置 0、1、2，现在一次追加位置 3、4。只计算两个新 query，但它们要读 5 个 key。

因果关系没有变，仍然是“不能看未来”。变的是query 的第 0 行，现在代表位置 3，不是整段文本的位置 0。

用一组小数字检查：令所有 Q/K 分数为 0，五个 value 分别为 [10, 20, 30, 40, 50]。每个 query 对允许的位置取平均，就能直接看出 mask 是否错位。

截至 PyTorch 2.14（2026-10 核对文档），F.scaled_dot_product_attention(..., is_causal=True) 对非方阵采用左上角对齐，不会根据 KV 长度自动推断 query 在末尾。这里用实际位置构造 attn_mask，不同时传 is_causal=True。这是这个 API 的约定，不是说所有 attention kernel 都采用同一种对齐方式。SDPA 文档

import torch
import torch.nn.functional as F

def cached_attention_example():
    query = torch.zeros(1, 1, 2, 2, dtype=torch.float64)
    key = torch.zeros(1, 1, 5, 2, dtype=torch.float64)
    value = torch.tensor([10., 20., 30., 40., 50.], dtype=torch.float64)
    value = value.view(1, 1, 5, 1)
    query_positions = torch.tensor([3, 4])
    key_positions = torch.arange(5)
    allowed = key_positions[None, :] <= query_positions[:, None]
    aligned = F.scaled_dot_product_attention(
        query, key, value, attn_mask=allowed, dropout_p=0.0,
    )
    upper_left = F.scaled_dot_product_attention(
        query, key, value, is_causal=True, dropout_p=0.0,
    )
    return aligned, upper_left, allowed

aligned, upper_left, allowed = cached_attention_example()
torch.testing.assert_close(aligned.flatten(), torch.tensor([25., 30.], dtype=torch.float64))
torch.testing.assert_close(upper_left.flatten(), torch.tensor([10., 15.], dtype=torch.float64))
print(aligned.flatten().tolist(), upper_left.flatten().tolist())

如果只追加位置 4 一个 query，且传入的五个 key 全都有效，就没有未来位置需要挡住，可以不加 causal mask。有 padding、滑动窗口或批内长度不同，就不能直接照搬这个简化。 还要记得，functional SDPA 的 dropout 由 dropout_p 控制，外层调用 eval() 不会替它把这个参数改成 0。

上面的数字只检验 mask，不是完整 decoder。完整模型还要检查 RoPE 位置、缓存身份，以及缓存前后的 logits 是否一致，见 KV cache。

面试常见问题

自检

继续阅读

没有位置编码（positional encoding）、也没有依赖顺序的 mask 时，self-attention 对 token 置换是等变的：输入换序，输出跟着换序，并不是输出完全不变。Causal mask 本身已经区分了前后；位置编码还能提供更直接的位置信号。接着看残差连接和归一化如何帮助训练，再回到原版 Transformer 把整块拼起来。

本章公式的起点是 Attention Is All You Need §3.2；实现中的分块显存优化见 FlashAttention。

快速学习：Multi-Head Attention 的一句话与边界条件