Multi-Head Latent Attention: Boosting Inference Efficiency

Introduction

• 作者提出 Multi-head Latent Attention (MLA)，通过将 KV 压缩为 Compressed Latent KV，在减小 KV cache 的同时保持模型精度
  

Method

Low-Rank Key-Value Joint Compression

• MLA 将 KV vectors ${\mathbf{k}}_t,{\mathbf{v}}_t\in {\mathbb{R}}^{d_h{n}_h}$ 压缩为 latent vector ${\mathbf{c}}_t^{KV}\in {\mathbb{R}}^{d_c}$，从而在推理时仅需保存 latent vector ${\mathbf{c}}_t^{KV}$ 而无需保存 KV cache ($d_c\ll d_h{n}_h$，$d_h$ 为 head dim，$n_h$ 为 #heads)
  $$\begin{array}{c}{\mathbf{c}}_t^{KV}=W^{DKV}{\mathbf{h}}_t\\ {\mathbf{k}}_t^C=W^{UK}{\mathbf{c}}_t^{KV}\\ {\mathbf{v}}_t^C=W^{UV}{\mathbf{c}}_t^{KV}\end{array}$$其中，$W^{DKV}\in {\mathbb{R}}^{d_c\times d}$，$W^{UK},W^{UV}\in {\mathbb{R}}^{d_h{n}_h\times d_c}$. 这样每个 token 对应的 KV cache 数据量由原来的 $2n_h{d}_{h}l$ 降低到了 $d_{c}l$，$l$ 为 Transformer 层数，这样一来，在设计 LLM 架构参数时甚至可以把 $d_h$ 设置得比 $d/h_n$ 更大，这样不仅不会增加 KV cache，还可以进一步提升模型能力
• MLA 在推理时无需用 $W^{UK},W^{UV}$ 重新计算出 ${\mathbf{k}}_t^C,{\mathbf{v}}_t^C$，而是将 $W^{UK},W^{UV}$ 分别融到模型权重 $W^Q,W^O$ 里，不会带来额外的推理开销
  $${\mathbf{q}}_t^T{\mathbf{k}}_t^C=(W_{(h)}^Q{\mathbf{h}}_t{)}^T(W_{(h)}^{UK}{\mathbf{c}}_t^{KV})={\left({\left(W_{(h)}^{UK}\right)}^T{W}_{(h)}^Q{h}_t\right)}^T{\mathbf{c}}_t^{KV}$$$${\left(\sum _{j=1}^t{\mathbf{p}}_j{\mathbf{v}}_j^C\right)}^T{W}_{(h)}^O={\left(\sum _{j=1}^t{\mathbf{p}}_j{W}_{(h)}^{UV}{\mathbf{c}}_j^{KV}\right)}^T{W}_{(h)}^O={\left(\sum _{j=1}^t{\mathbf{p}}_j{\mathbf{c}}_j^{KV}\right)}^T{\left(W_{(h)}^{UV}\right)}^T{W}_{(h)}^O$$其中，$W_{(h)}^Q,W_{(h)}^O\in {\mathbb{R}}^{d_h\times d_h{n}_h}$，$W_{(h)}^{UK},W_{(h)}^{UV}\in {\mathbb{R}}^{d_h\times d_c}$ 为 head $h$ 对应的权重参数

Decoupled Rotary Position Embedding

• 上述对 KV cache 的低秩压缩无法直接与 RoPE 兼容，因为 RoPE 要给 $\mathbf{q},\mathbf{k}$ 做内积之前进行旋转，这导致 $W^{UK}$ 无法融到 $W^Q$ 里，每次推理时都需要重新从 ${\mathbf{c}}^{KV}$ 计算 $\mathbf{k}$，从而增加大量推理开销。为此，MLA 采用 decoupled RoPE，给每个 attn 层额外增加 multi-head queries ${\mathbf{q}}_{t,i}^R\in {\mathbb{R}}^{d_h^R}$ 和共享的 key ${\mathbf{k}}_t^R\in {\mathbb{R}}^{d_h^R}$ 用于存储 RoPE 位置信息，这样只需要同时存储 ${\mathbf{c}}^{KV}$ 和 ${\mathbf{k}}^R$ 即可，MLA 所需的 KV cache 数据量增加为 $(d_c+d_h^R)l$
  $$\begin{array}{rl}[{\mathbf{q}}_{t,1}^R;{\mathbf{q}}_{t,2}^R;...;{\mathbf{q}}_{t,n_h}^R]={\mathbf{q}}_t^R& =\mathrm{RoPE}(W^{QR}{\mathbf{h}}_t),\\ {\mathbf{k}}_t^R& =\mathrm{RoPE}(W^{KR}{\mathbf{h}}_t),\\ {\mathbf{q}}_{t,i}& =[{\mathbf{q}}_{t,i}^C;{\mathbf{q}}_{t,i}^R],\\ {\mathbf{k}}_{t,i}& =[{\mathbf{k}}_{t,i}^C;{\mathbf{k}}_t^R],\\ {\mathbf{o}}_{t,i}& =\sum _{j=1}^t{\mathrm{Softmax}}_j(\frac{{\mathbf{q}}_{t,i}^T{\mathbf{k}}_{j,i}}{\sqrt{d_h+d_h^R}}){\mathbf{v}}_{j,i}^C,\\ {\mathbf{u}}_t& =W^O[{\mathbf{o}}_{t,1};{\mathbf{o}}_{t,2};...;{\mathbf{o}}_{t,n_h}],\end{array}$$其中，$W^{QR}\in {\mathbb{R}}^{d_h^R{n}_h\times d},W^{KR}\in {\mathbb{R}}^{d_h^R\times d}$

References

• DeepSeek-AI, et al. “DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model.” arXiv preprint arXiv:2405.04434 (2024).
• 苏剑林. (May. 13, 2024). 《缓存与效果的极限拉扯：从MHA、MQA、GQA到MLA 》[Blog post]. Retrieved from https://kexue.fm/archives/10091