Attention — Fanyi Pu

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

Multi-Head Attention

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

O = [\begin{matrix} O_{1} & \dots & O_{h} \end{matrix}] W_{o} O_{t} = Attention (X_{q} W_{q}^{(t)}, X_{k} W_{k}^{(t)}, X_{v} W_{v}^{(t)})

为了压缩 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 R^{s_{q} \times d_{c}^{(q)}}$ 和 $C_{k v} \in R^{s_{k v} \times d_{c}^{(k v)}}$ 两个低秩矩阵。这样子的话

W_{q} \in R^{d_{c}^{(q)} \times d_{q k}}, W_{k} \in R^{d_{c}^{(k v)} \times d_{q k}}, W_{v} \in R^{d_{c}^{(k v)} \times d_{v}}

然后考虑

Q K^{⊤} = C_{q} W_{q} {(C_{k v} W_{k})}^{⊤} = C_{q} W_{q} W_{k}^{⊤} C_{k v}^{⊤}

因为 $C_{q} W_{q} W_{k}^{⊤}$ 计算比较简单，我们把这部分合并，cache $C_{k v}$ 即可。

但问题是，我们之前因为需要得到 $K = C_{k v} W_{k}$ 然后 apply RoPE (Su et al., 2023)，现在 $C_{q} W_{q} W_{k}^{⊤}$ 合并了，无法在中间插入 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