Deep Dive into Multi-head Latent Attention (MLA): What does "be absorbed" mean?

Purpose

The purpose of this article is to organize "with equations" the following expression from the Multi-head Latent Attention (MLA) section of the DeepSeek V2 paper. Primarily, it organizes "which weights" and "how they are absorbed" using "equations."
(I might turn this into a more detailed explanatory article later.)

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 compute keys and values out for attention.

This article assumes that you have an understanding of the Multi-head Latent Attention (MLA) described in the DeepSeek V2 paper.
https://arxiv.org/abs/2405.04434

Note that regarding the claim that it still holds even after the introduction of RoPE, I will not handle that "this time."
However, the basic claim remains the same. I might write an article about it eventually.

I am still in the process of learning.
If I have made any errors in the theoretical derivation, please let me know.

In doing so, I would appreciate it if you could explain clearly what is wrong and why, so that I can understand it.

This article is also intended to seek guidance from those who understand the DeepSeek V2 paper.
At the very bottom of the article, there are some questions I am struggling with because I don't understand them. I would be very happy if someone knowledgeable could teach me!

That said, I have described MLA in detail—which has very wide "gaps between the lines" in the paper—so I would be happy if this helps someone's understanding.
Also, I wanted to write an article that current AI cannot write.
(Even if you ask something like o1 pro for an explanation, it likely won't provide one this detailed.)

Additionally, I have also written a paper commentary article on the recently popular DeepSeek-R1.
Unlike this article, I tried to write it so that those who are not familiar with theory would find it interesting! So please take a look if you'd like.
https://zenn.dev/asap/articles/34237ad87f8511

Organizing MLA with Equations

Scaled Dot-Product Attention is expressed as follows:

$$\text{Attention}_{(i)}(Q_{(i)}, K_{(i)}, V_{(i)}) = \text{softmax} \left( \frac{Q_{(i)}K_{(i)}^\top}{\sqrt{d_h}} \right) V_{(i)}$$

Based on the above, let's examine the derivation for MLA.

Introduction of Symbols

MLA is primarily aimed at compressing the KV cache and queries.
Therefore, we will introduce symbols to handle compressed KV and query representations.

Compressed KV Representation and KV

Here, we describe the compressed KV representation and the original KV.

In the DeepSeek V2 paper, equations are written for specific tokens, but since I thought it would be boring to do the same, I will write them down here using matrices that summarize all tokens.

I am doing this as a tribute to the Attention Is All You Need paper, which also described Attention using the formula above.

In MLA, we introduce the compressed KV representation $\mathbf{C}^{KV} \in \mathbb{R}^{T \times d_c}$.
$T$ is the total number of tokens, and $d_c$ is the dimensionality of the compressed KV representation (512 dimensions in the paper).

This compressed representation is created as follows:
(The following transformation is applied to each token. In other words, parallel processing is performed for the number of tokens.)

$$\begin{equation} \mathbf{c}_t^{KV} = W^{DKV} \mathbf{h_t} \end{equation}$$

$$\begin{equation} \mathbf{C}^{KV} = \begin{bmatrix} (\mathbf{c}_1^{KV})^T \\ (\mathbf{c}_2^{KV})^T \\ \vdots \\ (\mathbf{c}_T^{KV})^T \end{bmatrix} \end{equation}$$

Here, $t$ denotes the features of the $t$-th token.
Also, $\mathbf{h_t} \in \mathbb{R}^{d}$ is the output of the previous layer at the $t$-th token, where $d$ is the output dimensionality of the previous layer (5120 dimensions in the paper).

Furthermore, $W^{DKV} \in \mathbb{R}^{d_c \times d}$ is the down-projection matrix.
$\mathbf{c}_t^{KV} \in \mathbb{R}^{d_c}$ is the compressed KV representation at the $t$-th token.
And $\mathbf{C}^{KV} \in \mathbb{R}^{T \times d_c}$ is a matrix that summarizes $\mathbf{c}_t^{KV}$ for all tokens.

Using the compressed KV representation $\mathbf{C}^{KV}$, we can restore the keys and values.

$$\begin{equation} \mathbf{K} = \left(W^{UK} (\mathbf{C}^{KV})^T\right)^T \end{equation}$$

$$\begin{equation} \mathbf{V} = \left(W^{UV} (\mathbf{C}^{KV})^T\right)^T \end{equation}$$

$W^{UK}, W^{UV} \in \mathbb{R}^{d_hn_h \times d_c}$ are the up-projection matrices.
They restore all key and value representations for the multi-head attention from the compressed KV.
$d_h$ is the dimensionality per head (128 dimensions in the paper), and $n_h$ is the number of heads (128 heads in the paper).
$\mathbf{K}, \mathbf{V} \in \mathbb{R}^{T \times d_hn_h}$ are the uncompressed key and value representations, respectively.

Compressed Query Representation and Q

Next, queries are also compressed to reduce parameters during training.

$$\begin{equation} \mathbf{c}_t^{Q} = W^{DQ} \mathbf{h_t} \end{equation}$$

$$\begin{equation} \mathbf{C}^{Q} = \begin{bmatrix} (\mathbf{c}_1^{Q})^T \\ (\mathbf{c}_2^{Q})^T \\ \vdots \\ (\mathbf{c}_T^{Q})^T \end{bmatrix} \end{equation}$$

Also, $W^{DQ} \in \mathbb{R}^{d_c' \times d}$ is the down-projection matrix.
$d_c'$ is the dimensionality of the compressed query representation (1536 dimensions in the paper).
$\mathbf{c}_t^{Q} \in \mathbb{R}^{d_c'}$ is the compressed query representation at the $t$-th token.
$\mathbf{C}^{Q} \in \mathbb{R}^{T \times d_c'}$ is the compressed query representation for the entire sequence of tokens.

Using this compressed query representation $\mathbf{C}^{Q}$, the original query can be restored.

$$\begin{equation} \mathbf{Q} = \left(W^{UQ} (\mathbf{C}^{Q})^T\right)^T \end{equation}$$

$W^{UQ} \in \mathbb{R}^{d_hn_h \times d_c'}$ is the up-projection matrix.
It restores all query representations for the multi-head attention from the compressed query representation.
$\mathbf{Q} \in \mathbb{R}^{T \times d_hn_h}$ is the uncompressed query representation.

Difference between regular MHA (Multi-Head Attention) and MLA

In regular MHA, Q, K, and V are each expressed as follows:

Key representation in MHA

$$\begin{equation} \mathbf{k}_t = W^{K} \mathbf{h_t} \end{equation}$$

$$\begin{equation} \mathbf{K} = \begin{bmatrix} (\mathbf{k}_1)^T \\ (\mathbf{k}_2)^T \\ \vdots \\ (\mathbf{k}_T)^T \end{bmatrix} \end{equation}$$

Therefore,

$$\begin{equation} \mathbf{K} = \begin{bmatrix} (W^{K} \mathbf{h_1})^T \\ (W^{K} \mathbf{h_2})^T \\ \vdots \\ (W^{K} \mathbf{h_T})^T \end{bmatrix} \end{equation}$$

Note that $\mathbf{k}_t \in \mathbb{R}^{d_hn_h}$, $W^{K} \in \mathbb{R}^{d_hn_h \times d}$, and $\mathbf{K} \in \mathbb{R}^{T \times d_hn_h}$.

The following are almost the same, so I have hidden them. Please take a look if you are interested.

Query representation in MHA

Query representation in MHA

$$\begin{equation} \mathbf{q}_t = W^{Q} \mathbf{h_t} \end{equation}$$

$$\begin{equation} \mathbf{Q} = \begin{bmatrix} (\mathbf{q}_1)^T \\ (\mathbf{q}_2)^T \\ \vdots \\ (\mathbf{q}_T)^T \end{bmatrix} \end{equation}$$

Therefore,

$$\begin{equation} \mathbf{Q} = \begin{bmatrix} (W^{Q} \mathbf{h_1})^T \\ (W^{Q} \mathbf{h_2})^T \\ \vdots \\ (W^{Q} \mathbf{h_T})^T \end{bmatrix} \end{equation}$$

Note that $\mathbf{q}_t \in \mathbb{R}^{d_hn_h}$, $W^{Q} \in \mathbb{R}^{d_hn_h \times d}$, and $\mathbf{Q} \in \mathbb{R}^{T \times d_hn_h}$

Value representation in MHA

Value representation in MHA

$$\begin{equation} \mathbf{v}_t = W^{V} \mathbf{h_t} \end{equation}$$

$$\begin{equation} \mathbf{V} = \begin{bmatrix} (\mathbf{v}_1)^T \\ (\mathbf{v}_2)^T \\ \vdots \\ (\mathbf{v}_T)^T \end{bmatrix} \end{equation}$$

Therefore,

$$\begin{equation} \mathbf{V} = \begin{bmatrix} (W^{V} \mathbf{h_1})^T \\ (W^{V} \mathbf{h_2})^T \\ \vdots \\ (W^{V} \mathbf{h_T})^T \end{bmatrix} \end{equation}$$

Note that $\mathbf{v}_t \in \mathbb{R}^{d_hn_h}$, $W^{V} \in \mathbb{R}^{d_hn_h \times d}$, and $\mathbf{V} \in \mathbb{R}^{T \times d_hn_h}$

As you can see, they are constructed in almost exactly the same way.
So, let's look at the difference between MHA and MLA, focusing on the key representation.

Key representation in MLA

What is important are the following three equations described in the "Compressed KV Representation and KV" section.

$$\mathbf{c}_t^{KV} = W^{DKV} \mathbf{h_t}$$

$$\mathbf{C}^{KV} = \begin{bmatrix} (\mathbf{c}_1^{KV})^T \\ (\mathbf{c}_2^{KV})^T \\ \vdots \\ (\mathbf{c}_T^{KV})^T \end{bmatrix}$$

$$\mathbf{K} = \left(W^{UK} (\mathbf{C}^{KV})^T\right)^T$$

These correspond to equation numbers (1), (2), and (3), respectively.
Let's combine them.

From equations (2) and (3):

$$\begin{equation} \mathbf{K} = \left(W^{UK} [\mathbf{c}_1^{KV}, \mathbf{c}_2^{KV}, \cdots ,\mathbf{c}_T^{KV}]\right)^T \end{equation}$$

From equations (1) and (11):

$$\mathbf{K} = \left(W^{UK} [W^{DKV} \mathbf{h_1} ,W^{DKV} \mathbf{h_2} ,\cdots ,W^{DKV} \mathbf{h_T}]\right)^T$$

$$= [W^{UK}W^{DKV} \mathbf{h_1} ,W^{UK}W^{DKV} \mathbf{h_2} ,\cdots ,W^{UK}W^{DKV} \mathbf{h_T}]^T$$

So, aligning the notation with how we have written it so far:

$$\begin{equation} \mathbf{K} = \begin{bmatrix} (W^{UK}W^{DKV} \mathbf{h_1})^T \\ (W^{UK}W^{DKV} \mathbf{h_2})^T \\ \vdots \\ (W^{UK}W^{DKV} \mathbf{h_T})^T \end{bmatrix} \end{equation}$$

Note that for clarity, $\mathbf{K} \in \mathbb{R}^{T \times d_hn_h}$ and $W^{UK}W^{DKV} \mathbf{h_t} \in \mathbb{R}^{d_hn_h}$.

By the way, the key representation of MHA was as follows:

$$\begin{equation} \mathbf{K} = \begin{bmatrix} (W^{K} \mathbf{h_1})^T \\ (W^{K} \mathbf{h_2})^T \\ \vdots \\ (W^{K} \mathbf{h_T})^T \end{bmatrix} \end{equation}$$

Therefore, if training is done appropriately, from equations (12) and (13), we have $W^{K} = W^{UK}W^{DKV}$.
In this case, the degree of degradation is determined by the rank of the matrix $W^{K}$ and the degree of freedom $d_c$ of the matrices $W^{UK}$ and $W^{DKV}$.
In the unlikely event that the rank of the matrix $W^{K}$ is smaller than the degree of freedom $d_c$ of the matrices $W^{UK}$ and $W^{DKV}$, no degradation will occur with proper training.

Additionally, while $W^{K} \in \mathbb{R}^{d_hn_h \times d}$ has $d_hn_h \times d$ parameters, when split into $W^{UK} \in \mathbb{R}^{d_hn_h \times d_c}$ and $W^{DKV} \in \mathbb{R}^{d_c \times d}$, the number of parameters becomes $d_hn_h \times d_c + d_c \times d$.
Since $d$ is very large compared to $d_c$, it is possible to approximate it with fewer parameters than the former.

In the world of deep learning, it has often been observed that approximating by reducing the number of trainable parameters leads to improved accuracy (as seen in the world of CNNs for image processing), so the fact that MLA performed better than MHA this time might be due to that perspective as well.

What does "be absorbed" mean??

Now, let's verify that the following content from the paper actually holds.

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 compute keys and values out for attention.

Absorption of $W^{UK}$

Assuming a single head, if we focus on a specific head $i$, Attention can be expressed as follows:

$$\begin{equation} \text{SDPA}_{(i)}(Q_{(i)}, K_{(i)}, V_{(i)}) = \text{softmax} \left( \frac{Q_{(i)}K_{(i)}^\top}{\sqrt{d_h}} \right) V_{(i)} \end{equation}$$

Note that $\text{SDPA}_{(i)} \in \mathbb{R}^{T \times d_h}$.
SDPA stands for Scaled Dot-Product Attention.

The part

$W^{UK}$ can be absorbed into $W^Q$

relates to $Q_{(i)}K_{(i)}^\top$, so we will focus on that part.

However, since we are examining this head by head, the symbols used so far are divided as follows:
$Q_{(i)} \in \mathbb{R}^{T \times d_h}$, $K_{(i)} \in \mathbb{R}^{T \times d_h}$, $W^{UQ}_{(i)} \in \mathbb{R}^{d_h \times d_c'}$, $W^{UK}_{(i)} \in \mathbb{R}^{d_h \times d_c}$
The following remain unchanged as they are used across heads:
$\mathbf{C}^{Q} \in \mathbb{R}^{T \times d_c'}$, $\mathbf{C}^{KV} \in \mathbb{R}^{T \times d_c}$

Now, in MLA, $Q_{(i)}K_{(i)}^\top$ can be written out as follows:

$$Q_{(i)}K_{(i)}^\top = \left(W^{UQ}_{(i)} (\mathbf{C}^{Q})^T\right)^T W^{UK}_{(i)} (\mathbf{C}^{KV})^T$$

$$= \mathbf{C}^{Q} (W^{UQ}_{(i)})^T W^{UK}_{(i)} (\mathbf{C}^{KV})^T$$

$$\begin{equation} = \left((W^{UK}_{(i)})^T W^{UQ}_{(i)} (\mathbf{C}^{Q})^T \right)^T (\mathbf{C}^{KV})^T \end{equation}$$

In the above, $W^{UK}_{(i)} \in \mathbb{R}^{d_h \times d_c}$ and $W^{UQ}_{(i)} \in \mathbb{R}^{d_h \times d_c'}$ are weight matrices that were learned separately during training.

Therefore, during inference, we can pre-calculate $(W^{UK}_{(i)})^T W^{UQ}_{(i)} \in \mathbb{R}^{d_c \times d_c'}$ from these weight matrices and combine them into a single new weight matrix.

By doing so, we can calculate $Q_{(i)}K_{(i)}^\top$ in one go without explicitly restoring the keys and queries from the compressed KV representation $\mathbf{C}^{KV}$ and the compressed query representation $\mathbf{C}^{Q}$.

This claim is my personal favorite point! I thought it was very interesting!

This seems to be the meaning of the words "$W^{UK}$ is absorbed into $W^{UQ}$ ($W^Q$)" during inference.

Absorption of $W^{UV}$

The part > $W^{UV}$ can be absorbed into $W^{O}$ refers to the stage where, after creating the Attention map and calculating the matrix product with the values, the output for each head is consolidated into one using $W^{O}$. Since this part is a bit complex, let's organize it step-by-step.

First, let's represent the Attention map part as $Atn_{(i)}$. Specifically, it is as follows:

$$\begin{equation} \text{Atn}_{(i)}(Q_{(i)}, K_{(i)}) = \text{softmax} \left( \frac{Q_{(i)}K_{(i)}^\top}{\sqrt{d_h}} \right) \end{equation}$$

Here, $\text{Atn}_{(i)} \in \mathbb{R}^{T \times T}$.

Then, Scaled Dot-Product Attention can be expressed as follows:

$$\begin{equation} \text{SDPA}_{(i)} = \text{Atn}_{(i)} V_{(i)} \end{equation}$$

In that case, MHA (Multi-Head Attention) is expressed as follows:

$$\begin{equation} \text{MHA}(Q, K, V) = W^{O} \begin{bmatrix} \text{SDPA}_{(1)}^T \\ \text{SDPA}_{(2)}^T \\ \vdots \\ \text{SDPA}_{(n_h)}^T \end{bmatrix} \end{equation}$$

Here, $W^{O} \in \mathbb{R}^{d \times d_hn_h}$ and ${SDPA}_{(i)}^T \in \mathbb{R}^{d_h \times T}$, resulting in:

$$\begin{equation} \begin{bmatrix} \text{SDPA}_{(1)}^T \\ \text{SDPA}_{(2)}^T \\ \vdots \\ \text{SDPA}_{(n_h)}^T \end{bmatrix} \in \mathbb{R}^{d_hn_h \times T} \end{equation}$$

Now, it's a bit tricky, but we will decompose $W^{O}$ as follows:

$$\begin{equation} \text{MHA}(Q, K, V) = [W^{O}_{(1)}, W^{O}_{(2)}, \cdots, W^{O}_{(n_h)}] \begin{bmatrix} \text{SDPA}_{(1)}^T \\ \text{SDPA}_{(2)}^T \\ \vdots \\ \text{SDPA}_{(n_h)}^T \end{bmatrix} \end{equation}$$

In this case, $W^{O}_{(i)} \in \mathbb{R}^{d \times d_h}$.

According to the properties of matrix products, the following holds:

$$\begin{equation} \begin{aligned} \text{MHA}(Q, K, V) &=[W^{O}_{(1)}, W^{O}_{(2)}, \cdots, W^{O}_{(n_h)}] \begin{bmatrix} \text{SDPA}_{(1)}^T \\ \text{SDPA}_{(2)}^T \\ \vdots \\ \text{SDPA}_{(n_h)}^T \end{bmatrix} \\ &= W^{O}_{(1)} \text{SDPA}_{(1)}^T \\ &\quad + W^{O}_{(2)} \text{SDPA}_{(2)}^T \\ &\quad + \cdots + W^{O}_{(n_h)} \text{SDPA}_{(n_h)}^T \end{aligned} \end{equation}$$

Expanding equation (17), we can write it as follows:

$$\begin{equation} \begin{aligned} \text{MHA}(Q, K, V) &= W^{O}_{(1)} \left(\text{Atn}_{(1)}V_{(1)}\right)^T \\ &\quad + W^{O}_{(2)} \left(\text{Atn}_{(2)}V_{(2)}\right)^T \\ &\quad + \cdots + W^{O}_{(n_h)} \left(\text{Atn}_{(n_h)} V_{(n_h)}\right)^T \end{aligned} \end{equation}$$

At this point, from equation (4), we can transform the equation as follows:

Equation (4) is as follows:

$$\mathbf{V} = \left(W^{UV} (\mathbf{C}^{KV})^T\right)^T$$

Note that $\mathbf{C}^{KV} \in \mathbb{R}^{T \times d_c}$, $W^{UV} \in \mathbb{R}^{d_hn_h \times d_c}$, and $\mathbf{V} \in \mathbb{R}^{T \times d_hn_h}$.

When decomposed into each head, it looks like this:

$$\mathbf{V}_{(i)} = \left(W^{UV}_{(i)} (\mathbf{C}^{KV})^T\right)^T$$

$$= \mathbf{C}^{KV} (W^{UV}_{(i)})^T$$

Note that $\mathbf{C}^{KV} \in \mathbb{R}^{T \times d_c}$, $W^{UV}_{(i)} \in \mathbb{R}^{d_h \times d_c}$, and $\mathbf{V}_{(i)} \in \mathbb{R}^{T \times d_h}$.

$$\begin{aligned} \text{MHA}(Q, K, V) &= W^{O}_{(1)} \left(\text{Atn}_{(1)}\mathbf{C}^{KV} (W^{UV}_{(1)})^T \right)^T \\ &\quad + W^{O}_{(2)} \left(\text{Atn}_{(2)}\mathbf{C}^{KV} (W^{UV}_{(2)})^T \right)^T \\ &\quad + \cdots + W^{O}_{(n_h)} \left(\text{Atn}_{(n_h)} \mathbf{C}^{KV} (W^{UV}_{(n_h)})^T \right)^T \end{aligned}$$

$$\begin{equation} \begin{aligned} \text{MHA}(Q, K, V) &= W^{O}_{(1)} W^{UV}_{(1)} (\mathbf{C}^{KV})^T \text{Atn}_{(1)}^T \\ &\quad + W^{O}_{(2)} W^{UV}_{(2)} (\mathbf{C}^{KV})^T \text{Atn}_{(2)}^T \\ &\quad + \cdots + W^{O}_{(n_h)} W^{UV}_{(n_h)} (\mathbf{C}^{KV})^T \text{Atn}_{(n_h)}^T \\ \end{aligned} \end{equation}$$

In the above, $W^{O}_{(i)} \in \mathbb{R}^{d \times d_h}$ and $W^{UV}_{(i)} \in \mathbb{R}^{d_h \times d_c}$ are weight matrices that were learned separately during training.

Therefore, during inference, we can pre-calculate $W^{O}_{(i)} W^{UV}_{(i)} \in \mathbb{R}^{d \times d_c}$ from these weight matrices and combine them into a single new weight matrix.

By doing so, we can calculate $\text{MHA}(Q, K, V)$ in one go without explicitly restoring the values from the compressed KV representation $\mathbf{C}^{KV}$.

This seems to be the meaning of the words "$W^{UV}$ is absorbed into $W^{O}$" during inference.

Summary

Summary of "Absorption of $W^{UK}$"

In the section "Absorption of $W^{UK}$," the following equation holds based on equation (15):

$$Q_{(i)}K_{(i)}^\top = \left((W^{Q}_{(i)})' (\mathbf{C}^{Q})^T \right)^T (\mathbf{C}^{KV})^T$$

However, I have named the consolidated equation (the absorbed equation) as $(W^{Q}_{(i)})'$. $(W^{Q}_{(i)})'$ is defined as follows:

$$(W^{Q}_{(i)})' = (W^{UK}_{(i)})^T W^{UQ}_{(i)}$$

However, $(W^{Q}_{(i)})' \in \mathbb{R}^{d_c \times d_c'}$.
Also, $W^{UK}_{(i)} \in \mathbb{R}^{d_h \times d_c}$ and $W^{UQ}_{(i)} \in \mathbb{R}^{d_h \times d_c'}$.

On that basis, the equation for the Attention map in Scaled Dot-Product Attention is as follows:

$$\text{Atn}_{(i)}(Q_{(i)}, K_{(i)}) = \text{softmax} \left( \frac{\left((W^{Q}_{(i)})' (\mathbf{C}^{Q})^T \right)^T (\mathbf{C}^{KV})^T}{\sqrt{d_h}} \right)$$

Summary of "Absorption of $W^{UV}$"

In the section "Absorption of $W^{UV}$," the following holds based on equation (23):

$$\begin{aligned} \text{MHA}(Q, K, V) &= (W^{O}_{(1)})' (\mathbf{C}^{KV})^T \text{Atn}_{(1)}^T \\ &\quad + (W^{O}_{(2)})' (\mathbf{C}^{KV})^T \text{Atn}_{(2)}^T \\ &\quad + \cdots + (W^{O}_{(n_h)})' (\mathbf{C}^{KV})^T \text{Atn}_{(n_h)}^T \\ \end{aligned}$$

However, I have named the consolidated equation (the absorbed equation) as $(W^{O}_{(i)})'$. $(W^{O}_{(i)})'$ is defined as follows:

$$(W^{O}_{(i)})' = W^{O}_{(i)} W^{UV}_{(i)}$$

However, $(W^{O}_{(i)})' \in \mathbb{R}^{d \times d_c}$.
Also, $W^{O}_{(i)} \in \mathbb{R}^{d \times d_h}$ and $W^{UV}_{(i)} \in \mathbb{R}^{d_h \times d_c}$.

Overall Summary

I have shown all the necessary parts for calculating Multi-head Latent Attention (MLA) below:

$$Q_{(i)}K_{(i)}^\top = \left((W^{Q}_{(i)})' (\mathbf{C}^{Q})^T \right)^T (\mathbf{C}^{KV})^T$$

$$\begin{aligned} \text{MHA}(Q, K, V) &= (W^{O}_{(1)})' (\mathbf{C}^{KV})^T \text{Atn}_{(1)}^T \\ &\quad + (W^{O}_{(2)})' (\mathbf{C}^{KV})^T \text{Atn}_{(2)}^T \\ &\quad + \cdots + (W^{O}_{(n_h)})' (\mathbf{C}^{KV})^T \text{Atn}_{(n_h)}^T \\ \end{aligned}$$

From this, we can see that during inference, by pre-calculating the new weight matrices $(W^{Q}_{(i)})'$ and $(W^{O}_{(i)})'$ and keeping them in the model, it is possible to calculate MLA with fewer computations without explicitly calculating the keys, values, and queries.

This part! It's incredibly interesting, isn't it!

I believe the above is the meaning of the following statement in the paper:

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 compute keys and values out for attention.

My Questions

MLA seems to be a method that places great emphasis on improving memory efficiency, primarily through the compression of the KV cache and queries.

On the other hand, it seems to me that when matrices are absorbed to reduce the number of computations, the number of parameters in the post-absorption matrices becomes quite large.

For example, consider the following equation:

$$(W^{Q}_{(i)})' = (W^{UK}_{(i)})^T W^{UQ}_{(i)}$$

where $(W^{Q}_{(i)})' \in \mathbb{R}^{d_c \times d_c'}$, $W^{UK}_{(i)} \in \mathbb{R}^{d_h \times d_c}$, and $W^{UQ}_{(i)} \in \mathbb{R}^{d_h \times d_c'}$.

In this case, substituting the numbers from the paper ($d_c = 512, d_c' = 1536, d_h = 128$), the number of parameters for $(W^{Q}_{(i)})'$ is $786,432$, for $W^{UK}_{(i)}$ is $65,536$, and for $W^{UQ}_{(i)}$ is $196,608$.

That is, the number of parameters before absorption is $65,536 + 196,608 = 262,144$, while the number of parameters after absorption becomes $786,432$.

This would increase the required parameters during inference, which seems to conflict with MLA's advantage of improving memory efficiency. In that sense, is this absorption actually implemented or not? If it is, is the decrease in memory efficiency negligible or not?

This is a question that I haven't fully settled yet, but I would be happy if someone knowledgeable could explain it to me.

Note that I do not believe the idea is "instead of making a large matrix by absorbing, the multiplication is performed twice with two separate matrices." This is because that approach wouldn't be an issue when RoPE is introduced.

In that case, you would just need to apply the rotation matrix between the two matrix operations, and there would be no need to purposely separate Queries/Keys that have RoPE applied from those that don't, right?

To anyone knowledgeable, I would be very happy if you could let me know!

Summary

Thank you for reading!
I hope this is helpful for anyone who wanted to verify the "be absorbed" part of MLA.

Also, I would be happy if I could get answers to my questions.

Recommended Books on Attention

Transformers for Machine Learning Engineers: Model Development with State-of-the-Art Natural Language Processing Libraries
Since this book is written by developers from Hugging Face, it is reliable.
Naturally, it also contains detailed explanations of the Attention layer, making it a highly recommended book for learning the theory and fundamentals.

Introduction to Large Language Models
Introduction to Large Language Models II: Implementation and Evaluation of Generative LLMs
I introduce these often, but they cover a truly wide range of topics, from fine-tuning LLMs to RLHF, RAG, and distributed training, and I always refer to them.
Especially the first book is recommended as it explains the equations of the Transformer's Attention layer in an easy-to-understand manner.