multi-head attention 怎么算？参数量、FLOPs 多少？
为什么要 multi-head？

英文问法Walk me through multi-head attention with shapes. · How many parameters and FLOPs? · Why multiple heads? · Implement multi-head self-attention.

考过面经：NVIDIA LLM 日常实习 2023（"MultiheadAttention的计算流程，参数量，FLOPs，以及具体实现需要注意什么"，牛客）、NVIDIA DL Algorithm Engineer（"Implement multi headed self attention"，Glassdoor）、MAI Senior L63（手写 self-attention 和 masked self-attention）；他司：Apple、Meta MLE 手写 MHA（PracHub）、字节 AI infra 实习"mha手撕"。完整练习见 ml-coding。

先弄清楚：每个概念是什么

Q、K、V projection: 每个 token 乘三个 \(d\times d\) 矩阵。Q 是我要找什么，K 是我有什么，V 是被找到后交出的内容。
Head: 把 \(d\) 维切成 \(h\) 段，每段 \(d_k = d/h\) 维，各自独立做一遍 attention。拆头只是 reshape + transpose，不加参数、不多算。例：\(d = 512\)，\(h = 8\) → 每个 head 64 维。
Mask: softmax 前把不能看的位置设成 \(-\infty\)，权重就是 0。causal mask 遮未来，padding mask 遮补齐用的 token。

30 秒回答

Multi-head attention projects the input X of shape (b, s, d) into queries, keys and values with three d×d matrices, splits each into h heads of size d_k = d/h, and runs scaled dot-product attention in every head: softmax(QKᵀ/√d_k + mask)·V, with scores of shape (b, h, s, s). The heads are concatenated and multiplied by W_O. Parameters are 4d², plus 4d with biases, independent of h. Forward FLOPs are about 8sd² for the projections plus 4s²d for the two attention matmuls. Multiple heads attend to different subspaces and positions at once, at the cost of one full-width head.

计算流程：每一步的 shape

例子：\(b = 2\)，\(s = 3\)，\(d = 8\)，\(h = 2\)，所以 \(d_k = 4\)。

步骤	操作	shape	例子
投影	\(Q = XW_Q\)，K、V 同理	(b, s, d)	(2, 3, 8)
拆头	reshape 成 (b, s, h, d_k)，再交换 s、h	(b, h, s, d_k)	(2, 2, 3, 4)
score	\(QK^\top/\sqrt{d_k}\) + mask，softmax	(b, h, s, s)	(2, 2, 3, 3)
加权	\(AV\)	(b, h, s, d_k)	(2, 2, 3, 4)
合头	交换回来，reshape，再乘 \(W_O\) 混合各 head	(b, s, d)	(2, 3, 8)

参数量和 FLOPs

\[ W_Q, W_K, W_V, W_O \in \mathbb R^{d\times d} \;\Rightarrow\; 4d^2\ (+\,4d\ \text{bias}) \]

例：\(d = 512\) → \(4\cdot 512^2 = 1{,}048{,}576\)，加 bias 是 1,050,624。和 \(h\) 无关，\(h\) 变大只让 score 矩阵 (b, h, s, s) 的 memory 线性增长。

FLOPs（每条序列 forward）：4 个投影各 \(2sd^2\)，\(QK^\top\) 和 \(AV\) 各 \(2s^2d\)，合计 \(8sd^2 + 4s^2d\)，\(s > 2d\) 时 \(s^2\) 项占主导。整层怎么数见 Transformer 复杂度。

为什么要 multi-head

一个 head 只有一组 attention 权重，不同关系会被平均到一起。\(h\) 个 head 在不同 subspace 里各算一组，能同时看不同位置。原论文 Table 3 里 \(h = 1\) 比最好的设置低 0.9 BLEU，\(h = 32\) 也变差。

手撕代码（NumPy）

def mha(X, Wq, Wk, Wv, Wo, h, mask=None):
    # X: (b, s, d)；W*: (d, d)；mask: (s, s)，True = 不能看
    b, s, d = X.shape
    dk = d // h
    def split(M):                    # (b, s, d) → (b, h, s, dk)
        return M.reshape(b, s, h, dk).transpose(0, 2, 1, 3)
    Q, K, V = split(X @ Wq), split(X @ Wk), split(X @ Wv)
    S = Q @ K.transpose(0, 1, 3, 2) / np.sqrt(dk)    # (b, h, s, s)
    if mask is not None:
        S = np.where(mask, -np.inf, S)
    S = S - S.max(-1, keepdims=True)                 # 防 overflow
    A = np.exp(S)
    A = A / A.sum(-1, keepdims=True)
    O = (A @ V).transpose(0, 2, 1, 3).reshape(b, s, d)  # concat
    return O @ Wo

必答点

流程：投影 → 拆成 \(h\) 个 head → \(\mathrm{softmax}(QK^\top/\sqrt{d_k} + \text{mask})V\) → concat → \(W_O\)
score 的 shape 是 (b, h, s, s)；\(d_k = d/h\)，缩放用 \(\sqrt{d_k}\) 不是 \(\sqrt{d}\)
参数 \(4d^2\)（+\(4d\) bias），与 \(h\) 无关
FLOPs ≈ \(8sd^2 + 4s^2d\)
实现：先 reshape (b, s, h, d_k) 再 transpose；mask 填 \(-\infty\)；softmax 先减 max

追问

拆头为什么不能直接 reshape 成 (b, h, s, d_k)？

直接 reshape 会把不同 token 的维度塞进同一个 head。例：\(s = 3\)，\(d = 8\)，head 0 应该是每个 token 的前 4 维，却拿到 token 0 的 8 维和 token 1 的前 4 维。

cross-attention 的 shape 是什么？

Q 长度 \(s_q\)（decoder），K、V 长度 \(s_k\)（encoder），score 是 (b, h, \(s_q\), \(s_k\))，输出长度跟 Q 一样。

一整行都被 mask 掉会怎样？

整行是 \(-\infty\)，减 max 时 \(-\infty - (-\infty)\) = nan。用 −1e9 代替 \(-\infty\)，或者把这一行输出置 0。

MQA / GQA 改了什么？

K、V 只投影出 \(g < h\) 组，每组被 \(h/g\) 个 query head 共享，\(W_K\)、\(W_V\) 变成 \(d\times g d_k\)。见 MQA / GQA。

来源Vaswani et al. 2017 §3.2.2（pp.4–5）、Table 3 rows (A)（p.9）· Dive into Deep Learning §11.5.1–11.5.2 · Ainslie et al. 2023（GQA）

multi-head attention 怎么算？参数量、FLOPs 多少？为什么要 multi-head？