Down-projection and On-the-Fly Up-projection in MLA

Hello! Welcome to our next lesson on Multi-head Latent Attention (MLA).

In our previous session, we introduced MLA as a powerful solution to the KV cache memory bottleneck. We established the core concept: instead of reducing the number of K/V heads like in MQA or GQA, MLA compresses the K/V representations themselves into a shared, low-dimensional latent space.

Today, we will move from the "what" to the "how." Our learning goal is to describe the down-projection and on-the-fly up-projection steps in MLA. We will deconstruct the precise mechanism that allows MLA to dramatically shrink the KV cache while preserving the model's expressive power.

1. The MLA Mechanism: A Two-Step Process

At its heart, MLA performs an elegant two-step dance of compression and reconstruction.

Step 1: Down-Projection (Compression)

The first step is to compress the high-dimensional input into the compact latent vector.

Input: The input to the attention block is a sequence of hidden states, $X$ , where each token has a vector of dimension d_model (e.g., 7168 in DeepSeek-R1).
Process: This input $X$ is multiplied by a learned down-projection matrix, $W_{D K V}$ . This matrix maps the d_model dimension down to a much smaller latent dimension, $d_{c}$ (e.g., 512).
$L_{K V} = X \cdot W_{D K V}$
Output: The result is the latent representation, $L_{K V}$ . This single, smaller matrix is the only thing that gets stored in the KV cache. The $W_{D K V}$ matrix is shared across all attention heads in the layer.

This image beautifully visualizes the compression.

Multi-Head Latent Attention: Compression and Decompression — This diagram shows the input word embedding (size 7168) being 'down-projected' by the `KV Down` matrix into a much smaller latent vector (size 512). This smaller vector is what's cached.

This down-projection step is the source of MLA's incredible memory savings. Instead of caching two large matrices (K and V) for each head, we cache just one small latent matrix.

Step 2: On-the-Fly Up-Projection (Reconstruction)

The second step is to reconstruct the full Key and Value matrices from the cached latent vector, but only when they are needed for the attention calculation.

Input: The cached latent representation, $L_{K V}$ .
Process: To get the Key ( $K_{i}$ ) and Value ( $V_{i}$ ) for a specific head $i$ , the cached $L_{K V}$ is multiplied by two separate up-projection matrices, $W_{U K, i}$ and $W_{U V, i}$ . $K_{i} = L_{K V} \cdot W_{U K, i}$ $V_{i} = L_{K V} \cdot W_{U V, i}$
Key Insight: Critically, the up-projection matrices $W_{U K, i}$ and $W_{U V, i}$ are unique to each head. This allows each head to "interpret" or reconstruct information from the shared latent vector in its own specialized way. This is how MLA maintains the representational diversity of Multi-Head Attention, which methods like MQA and GQA sacrifice.

The term "on-the-fly" means these full $K$ and $V$ matrices are never stored in the cache; they are generated dynamically, just for the duration of the attention computation in the forward pass.

2. The Full Workflow and the "Absorption Trick"

Now that we understand the two core steps, let's see how they fit together in a complete forward pass. A naive implementation might seem computationally expensive, but a clever mathematical rearrangement called the "absorption trick" makes it highly efficient.

The following video provides an excellent from-scratch walkthrough of the entire process, showing how the cached latent vector is used and why the absorption trick is so important.

Multi-Head Latent Attention From Scratch | One of the major DeepSeek innovation

This segment from 'Multi-Head Latent Attention From Scratch' by Vizuara walks through the full MLA workflow, detailing the down- and up-projections and explaining the 'absorption trick' that ensures computational efficiency.

Please watch from 37:19 to 51:10. Pay close attention to how the latent matrix C_KV is the only thing cached, and how it's used to compute both the attention scores (via the absorbed query) and the final context vector.

As the video demonstrated, by absorbing the key's up-projection matrix into the query computation, we can calculate attention scores directly using the compressed latent cache. This avoids expanding the entire cache at every step and makes MLA just as fast as standard attention during inference, while using a fraction of the memory.

3. A Look at the Implementation

To connect this theory to practice, let's examine how these steps are implemented in code. Since you have experience with PyTorch, this should help solidify your understanding.

The following resource provides a clear PyTorch implementation of an MLA layer.

Understanding DeepSeek's Multi-Head Latent Attention (MLA)

This article, 'Understanding DeepSeek's Multi-Head Latent Attention (MLA)', provides a clear code implementation that maps directly to the concepts we've discussed.

Please read the section 'Multi-Head Latent Attention : Low-Rank Decomposition of KV matrices' and study the MultiHeadLatentAttention class implementation in the following section. As you read the code, identify the following parts: The down-projection for KV: self.W_dkv. The up-projections for K and V: self.W_uk and self.W_uv. The latent_cache where the compressed L_kv is stored. The lines where the full K and V are reconstructed on-the-fly from L_kv.

In the forward method of the code, you can see the two steps in action:

Down-projection: L_kv = self.W_dkv(x) creates the latent vector.
Up-projection: K = self.W_uk(L_kv) and V = self.W_uv(L_kv) reconstruct the full matrices just before the attention score calculation.

Conclusion

Today we've dissected the mechanics of Multi-head Latent Attention, focusing on the two-step process that enables its efficiency.

Key Takeaways:

Down-Projection: MLA first compresses the input hidden state into a small, shared latent vector ( $L_{K V}$ ) using a down-projection matrix ( $W_{D K V}$ ). This is the compression step, and only $L_{K V}$ is cached.
On-the-Fly Up-Projection: During attention calculation, the cached latent vector is expanded back into full-sized Key and Value vectors using per-head up-projection matrices ( $W_{U K, i}$ and $W_{U V, i}$ ). This is the reconstruction step.
This mechanism allows each head to maintain its own unique representation while benefiting from a dramatically smaller shared cache, giving MLA the "best of both worlds": high memory efficiency and strong model performance.

Preview of the Next Lesson:
We have now examined the mechanisms of MHA, GQA, and MLA in detail. In our next lesson, we will conclude this module by conducting a direct comparison of these three architectures, evaluating their trade-offs in terms of memory, computation, and model quality. This will synthesize our understanding of modern efficient attention.