ML Revisit: Attention

参考了 (Purshow, 2026; Weng, 2023; 苏剑林, 2024)。

Multi-Head Attention

Multi-Head Attention (MHA) (Vaswani et al., 2023) 大概就是把 QKV 分成多组：

O = \begin{bmatrix} O_1&\cdots&O_h \end{bmatrix}W_o\\ O_t = \mathrm{Attention}\left(X_qW_q^{(t)}, X_kW_k^{(t)}, X_vW_v^{(t)}\right)

为了压缩 KV cache，Multi-Query Attention (MQA) (Shazeer, 2019) 让 $$K$$ 和 $$V$$ 合并成一个 head。但是这样可能导致压缩过猛，所以 Grouped-Query Attention (GQA) (Ainslie et al., 2023) 将 head 分成 $$g$$ 组，每组 share 同一个 $$W_k$$ 和 $$W_v$$ 。

相同的 8 个 Query heads自回归推理时缓存的 K/V

MHAMulti-Head Attention

8 组 K/V100%

GQAGrouped-Query Attention

2 组 K/V25%

MQAMulti-Query Attention

1 组 K/V12.5%

连线表示各 Query head 使用哪组 K/V；GQA 以 8 个 Query heads 分成 2 组为例。下方比例为相对 MHA 的 KV cache 大小，假设序列长度、每个 head 的维度、层数、batch size 和缓存精度相同。

Multi-Head Latent Attention

Multi-Head Latent Attention (MLA) (DeepSeek-AI et al., 2024) 尝试将 QKV 低秩压缩。

首先将 $$X$$ 压缩成 $C_q\in\mathbb{R}^{s_q\times d_{c}^{(q)}}$ 和 $C_{kv}\in\mathbb{R}^{s_{kv}\times d_{c}^{(kv)}}$ 两个低秩矩阵。这样子的话

W_q\in\mathbb{R}^{d_c^{(q)}\times d_{qk}},\quad W_{k}\in \mathbb{R}^{d_c^{(kv)}\times d_{qk}}, \quad W_{v}\in\mathbb{R}^{d_c^{(kv)}\times d_{v}}

然后考虑

QK^{\top} = C_qW_q\left(C_{kv}W_k\right)^{\top} = C_qW_{q}W_{k}^{\top}C_{kv}^{\top}

因为 $C_qW_qW_k^{\top}$ 计算比较简单，我们把这部分合并，cache $C_{kv}$ 即可。

但问题是，我们之前因为需要得到 $K=C_{kv}W_k$ 然后 apply RoPE (Su et al., 2023)，现在 $C_qW_qW_k^{\top}$ 合并了，无法在中间插入 rope。解决方案是给 $$Q$$ 和 $$K$$ 单独加一个 RoPE 分支： $$Q$$ 的 RoPE 分支从 $$C_q$$ 投影得到， $$K$$ 的 RoPE 分支直接从 $$X$$ 投影得到，并在所有 heads 间共享一个 $$K$$ 。

下面沿用上文的矩阵记号展开 MLA 的结构：

Query 分支K/V 内容分支RoPE 分支推理时缓存

从下往上读 · 横向滑动查看完整结构 →

每层的 KV cache$C_{kv}$ 和 $K_R$

从下往上读，沿用上文的矩阵记号：内容分支为 $Q=C_qW_q$、$K=C_{kv}W_k$、$V=C_{kv}W_v$。下标 $R$ 表示额外的 RoPE 分支，其中 $K_R$ 在所有 heads 间共享；方括号表示沿特征维度拼接。图中展开 K/V 是为了展示结构；推理时可将$W_k$ 吸收到 Query 侧，将 $W_v$ 吸收到输出投影 $W_o$ 中，无需缓存展开后的多头 K/V。

References

Ainslie, J., Lee-Thorp, J., de Jong, M., Zemlyanskiy, Y., Lebrón, F., & Sanghai, S. (2023). GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. arxiv.org

DeepSeek-AI, Liu, A., Feng, B., Wang, B., Wang, B., Liu, B., Zhao, C., Dengr, C., Ruan, C., Dai, D., Guo, D., Yang, D., Chen, D., Ji, D., Li, E., Lin, F., Luo, F., Hao, G., Chen, G., … Xie, Z. (2024). DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model. arxiv.org

Purshow. (2026). Attention Architecture Notes. github.com

Shazeer, N. (2019). Fast Transformer Decoding: One Write-Head is All You Need. arxiv.org

Su, J., Lu, Y., Pan, S., Murtadha, A., Wen, B., & Liu, Y. (2023). RoFormer: Enhanced Transformer with Rotary Position Embedding. arxiv.org

Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., & Polosukhin, I. (2023). Attention Is All You Need. arxiv.org

Weng, L. (2023). The Transformer Family Version 2.0. Lilianweng.Github.Io. lilianweng.github.io

苏剑林. (2024, May). 缓存与效果的极限拉扯：从MHA、MQA、GQA到MLA. spaces.ac.cn

Multi-Head Attention

Multi-Head Latent Attention

References

Cite this post