Multi-head Latent Attention (MLA)

The use of kv-cache with MHA at inference time becomes the bottleneck for efficiency due to its large size.

Transclude of SDPA#a383c1

Transclude of MHA#2094e9

MLA, equipped with low-rank key-value joint compression, achieves better performance than MHA while requiring a significantly smaller amount of KV cache. The core of this can be expressed as:

$$\begin{array}{rl}{\text{c}}_t^{KV}& =W^{DKV}{\text{h}}_t\\ {\text{k}}_t^C& =W^{UK}{\text{c}}_t^{KV}\\ {\text{v}}_t^C& =W^{UV}{\text{c}}_t^{KV}\end{array}$$

where ${\text{c}}_t^{KV}\in {\mathbb{R}}^{d_c}$ is the compressed latent vector for keys and values for a particular token; $d_c(\ll d_h{n}_h)$ denotes the KV compression dimension; $W^{DKV}\in {\mathbb{R}}^{d_c\times d}$ is the down-projection matrix; and $W^{UK},W^{UV}\in {\mathbb{R}}^{d_h{n}_h\times d_c}$ are the up-projection matrices for keys and values, respectively.

During inference, MLA only needs to cache ${\text{c}}_t^{KV}$, so its KV cache has only $d_{c}l$ elements, where $l$ denotes the number of layers. This method of compression is conceptually similar to LoRA, which is used for fine-tuning.

In addition, during inference, since $W^{UK}$ can be absorbed into $W^Q$, and $W^{UV}$ can be absorbed into $W^O$, we even do not need to computes keys and values out for attention.

During training, low-rank compression can also be applied to the queries (although it cannot reduce the size of the KV cache):

$$\begin{array}{rl}{\text{c}}_t^Q& =W^{DQ}{\text{h}}_t\\ {\text{q}}_t^C& =W^{UQ}{\text{c}}_t^Q\end{array}$$

where ${\text{c}}_t^Q\in {\mathbb{R}}^{d_c^{\prime }}$ is the compressed latent vector for queries for a particular token; $d_c^{\prime }(\ll d_h{n}_h)$ denotes the query compression dimension; $W^{DQ}\in {\mathbb{R}}^{d_c^{\prime }\times d}$, $W^{UQ}\in {\mathbb{R}}^{d_h{n}_h\times d_c^{\prime }}$ are the down-projection and up-projection matrices for queries, respectively.

Decoupled Rotary Position Embedding

As established by LLaMA (and used by DeepSeek), RoPE is an effective method of position embedding. However, RoPE is incompatible with MLA.

RoPE is position-sensitive for both keys and queries. If we apply RoPE for the keys (${\text{k}}_t^C$), the up-projection (*$W^{UK}$) will be coupled with a position sensitive RoPE matrix.

References

• DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model