Transcription
This video is sponsored by KiwiCo. More on them later.
In January 2025, the Chinese company DeepSeek shocked the world with the release of R1, a highly competitive language model that requires only a fraction of the compute of other leading models. Perhaps even more shocking is that, unlike most of its American counterparts, DeepSeek has publicly released the R1 model weights, inference code, and extensive technical reports, publishing an average of one report per month in 2024 and detailing any of the innovations that dramatically culminated in the release of R1 in early 2025.
Back in June of 2024, the DeepSeek team introduced a technique that they call multi-head latent attention. Unlike many DeepSeek innovations that occur at the margins of the stack, multi-head latent attention strikes at the core of the Transformer itself. This is the compute architecture that virtually all large language models share. This modification reduces the size of an important bottleneck called the key-value cache by a factor of 57, allowing the model to generate text more than six times faster than a traditional Transformer in DeepSeek's implementation.
But how exactly was the DeepSeek team able to squeeze such a significant improvement out of such a broadly used architecture? Like other modern language models, when you give DeepSeek a prompt, the model generates its response one fragment known as a token at a time. Mathematically, this autoregressive approach means that each new token the model generates is a function of all the tokens that came before it. The interactions between tokens in large language models are handled by a mechanism called attention. Attention works by computing matrices called attention patterns. These are the 144 attention patterns computed by the GPT-2 small model when given the example input text: "The American flag is red, white, and...". This model uses 12 separate attention mechanisms called attention heads per layer and has 12 layers, making for 144 total attention patterns. DeepSeek R1 has 128 attention heads per layer and 61 layers, making for 7,808 total patterns. In both models, the size of the attention pattern is equal to the number of tokens passed into the model. Our example input, "The American flag is red, white, and...", maps to nine tokens, so all of our attention patterns are 9x9 matrices.
Attention patterns are used by attention heads to move information between token positions in the model's residual stream. For example, this first attention pattern in the third layer of GPT-2 has a high value mapping from the input token of "American" to the output token of "flag," meaning this attention head is likely applying the modifier "American" to the noun "flag," creating a unified representation for the concept "American flag." This eighth attention pattern in the 11th layer has high values mapping the words "flag," "red," and "white" to the output of the final token "and," and this attention head has pulled out words in our input that are relevant for predicting the correct next token of "blue," which this GPT-2 small model does correctly predict.
Let's dig a bit deeper into exactly how the standard attention mechanism works in models like GPT-2 and build up a few equations so we can make sense of how the DeepSeek team made such a powerful improvement to compute a given attention pattern. We take the input matrix X. This could be the input to any layer of our model and will have one row for each input token and a number of columns that corresponds to the embedding dimension of the model. This is the length of the vector used to represent each token. GPT-2 Small's embedding dimension is 768, and DeepSeek R1's embedding dimension is 7168. To compute a given attention pattern, we multiply our input matrix X by two separate sets of learned weights, WQ and WK. In GPT-2, these matrices are of dimension 768x64 and result in two new matrices, Q and K, each of dimension 9x64. The rows of our Q matrix are known as queries, and the rows of our K matrix are known as keys. The core idea of attention is that we now search for pairs of tokens that have similar queries and keys, allowing the model to learn various relationships between tokens. For example, a token like "flag" could query for words that modify its meaning, while words like "American" can produce keys in certain attention heads that flag them as modifiers. This modifier query and modifier key should produce similar key and query vectors. Mathematically, to find similar keys and queries, we can take the dot product of the keys and queries for each possible pair of our nine tokens. Similar keys and query vectors will generate high dot products. We can compute all these dot products at once by transposing our key matrix and multiplying by our query matrix, resulting in a new 9x9 matrix where each entry corresponds to the dot product of a single key and query. To compute our attention pattern, we apply a masking operation, effectively zeroing out the upper right portion of our matrix. This step is mostly important in training as it prevents the model from cheating on its task of next-token prediction by just looking at the next token. Finally, we normalize our result by dividing by the square root of our embedding dimension and applying a softmax operation, which forces each of the rows of our matrix to add to one.
Now that we've computed our attention pattern, we need to actually use it to process our data. This involves a couple more matrix multiplies. We first compute what's known as a value matrix by multiplying our input by a third weight matrix, WV. This computation is identical to the way we computed our keys and query matrices, just with a different set of learned weights. We then multiply our attention pattern matrix by our value matrix. This effectively takes a weighted sum of our values following our attention pattern. One way to think about this step is as processing our inputs using a neural network where the weights are controlled by the data itself. Finally, the attention block in each layer has multiple heads. Each head performs the same computations but with different learned weights, resulting in a different set of queries, keys, attention patterns, and values for each head. The idea here is that various attention heads can specialize in various tasks like searching for adjectives or searching for other instances of the same token. To compute the final output of the attention block, we stack the results from each head together and multiply by a final learned weight matrix, WO, giving us the final matrix output of our attention block.
The attention block is a key part of modern language models but requires a significant amount of computation. Since the height and width of our attention pattern are equal to the number of input tokens, the number of entries in this matrix scales as the number of input tokens squared. This is potentially a huge computational problem for large models. OpenAI's ChatGPT models now offer maximum context lengths of over 100,000 tokens. For reference, this is about the length of the first Harry Potter book. So computing each attention pattern for ChatGPT's maximum allowed input size is equivalent to arranging the entire text of the book as a single row and column and then computing dot products for every possible pair of tokens from the entire text. Fortunately, there's a huge computational shortcut that we can take. As large language models generate new text a single token at a time, the attention patterns themselves don't actually change that much. In our "American flag" example, let's say the model generates a new token for the word "blue." Our phrase is now "The American flag is red, white, and blue." To see what the model says next, we now pass this new 10-token input back into the model to get the 11th token, and so on. Our new 10-token input results in key, query, and value matrices each of dimension 10x64. But importantly, since our weight matrices apply the same identical operation to each token, the first nine rows of our key, query, and value matrices are unchanged from our original nine-token input. Transposing our keys and multiplying by our queries to compute our new attention pattern, note that the first nine rows of Q and the first nine columns of K transpose are unchanged. This means that the upper left 9x9 matrix of our attention pattern will also be unchanged, and we only need to compute a new final row and column to arrive at our new 10x10 attention pattern. And further, since we mask out the upper right corner of our attention pattern, we actually only need to compute the new bottom row. The bottom row of our attention pattern results from multiplying the final row of our query matrix by each column of our transposed key matrix. So to compute this final attention pattern row, we need to know all of our keys but only the final new row of our query matrix. Since we already computed nine out of 10 of our keys on the previous call of the model, it's much more computationally efficient to store these keys in memory and just access them when the new 10-token input comes along. The same logic applies to our value matrix. We need our full value matrix to compute our new outputs, but the first nine rows are unchanged, so we can just cache them in memory. Note that there's no need to cache the queries since we only need the new final row of our queries to update our attention pattern. This idea is called key-value or KV caching and is a critical part of large language model infrastructure. Instead of compute growing quadratically as the square of the number of input tokens, key-value caching means that the compute required by the model's attention blocks scales linearly with the number of input tokens.
Now, this compute shortcut does come at a cost, specifically increased memory usage. Our system must now store the keys and values for the full history of the model session for all attention heads across all layers in memory. Given a model with L layers, NH attention heads per layer, a dimension of DH for our key and value matrices, and N input tokens, we must store 2 * N * DH * NH * L unique numbers in our KV cache. Assuming floating-point 16 numbers, the DeepSeek R1 architecture and a context length of 100,000 tokens, we end up needing to retrieve four megabytes per token in the model's context window, resulting in a huge 400 gigabytes of memory reads for each new token we compute.
DeepSeek's solution to this problem is really clever, and it was great to be able to tinker with their inference code to really get my head around it. There's nothing quite like hands-on experimenting like this for developing understanding, which is why I was more than happy to partner again with this video sponsor, KiwiCo. KiwiCo offers hands-on project kits that make learning genuinely fun for kids of all ages. My daughter is obsessed with colors and rainbows right now and loves the Color Discovery Crate. These spinners are such a fun way for us to explore color mixing together. It's amazing to see how the crates progress and build on each other. Last year she was developing fine motor skills as part of the Panda Club, and now in the Sprouts Club, she's creating and experimenting. When she turns six in a few years, she can join the KiwiCo Labs Club where she'll get to work on more complex science and engineering projects like this remote-controlled car. I would have absolutely loved this crate as a kid. My son is working on learning the names of colors; so far, everything is blue. This block puzzle is such a fun interactive way for him to explore different colors at his age. When my kids quickly get bored of or break many of their toys, we find ourselves continually coming back to their KiwiCo crates. The build quality is really great, and the thoughtfulness and multi-purpose design built into each crate really keeps them engaged. If you want your family to experience the awesomeness of KiwiCo, use my code WelchLabs to receive 50% off your first crate for kids three and older or 20% off your first Panda crate for kids under three. Big thanks to KiwiCo for sponsoring this video. Now, back to DeepSeek's solution to the KV cache problem.
Untenable large KV caches are not a new problem. One popular solution is to reuse key and value matrices across multiple attention heads in multi-query attention blocks. Instead of having unique key and value matrices for each attention head, we share a single key and value matrix across all heads. This reduces the required size of our KV cache by a significant factor of the number of heads per layer—128 for the DeepSeek R1 architecture. However, this modification does impact model performance, as forcing all attention heads to use the same keys and values allows for less specialization. A less destructive version of this idea is grouped query attention, where instead of forcing all attention heads in a given layer to share the same key and value matrices, we create multiple groups of attention heads that share the same key and value matrices. Meta Llama 3 models use grouped query attention with groups of eight attention heads sharing the same key and value matrices, reducing the size of the KV cache by a factor of eight. Grouped query attention reduces KV cache size but still takes a performance hit relative to full multi-head attention.
Now, what's really remarkable about DeepSeek's approach, called multi-head latent attention, is that they were able to reduce the needed KV cache size by a factor of 57 while actually improving performance. The key insight is a novel application of a very common idea in machine learning: a latent space. What if the model could learn to efficiently compress its own keys and values? Multi-head latent attention effectively adds an extra step between each attention head's input and the key and value matrices. The idea is to project our input into a compressed latent space that, like multi-query attention, is shared across all attention heads in a given block. However, unlike multi-query attention where each head shares the same exact keys and values, in multi-head latent attention, the compressed latent space is projected back up to keys and values using another set of learned weights, WUK and WUV, where the weights are unique to each attention head. This gives multi-head latent attention more flexibility than multi-query attention or grouped query attention.
Now, at face value, since we've introduced a new matrix multiply, it appears that we've just traded some memory bandwidth for additional compute, and after all, the entire point of KV caching was to reduce the high compute needs of attention blocks. However, as the DeepSeek team points out, with some clever linear algebra, we can rearrange our query computation to absorb the WUK weights and rearrange our final output computation to absorb the WUV weights. Since all these weights are fixed at training time, we only have to compute the absorbed weights once and can avoid any additional compute during inference. So when a new token comes along, we simultaneously compute its query vector and the query's projection into the latent cache space in one step and then compute our attention pattern directly from the latent key-value cache matrix. It's a really elegant solution. With multi-head latent attention, the size of the needed KV cache no longer has any dependence on the number of attention heads per layer and instead just depends on the size of the shared KV cache matrix. For DeepSeek R1, this is equal to the number of input tokens by 576. If implemented with traditional attention blocks, R1 would require four megabytes of KV cache per token. Grouped query attention with a group size of eight would cut this down to 500 kilobytes per token, and multi-head latent attention reduces the needed cache to only 70 kilobytes per token—a 57x reduction. What we're left with is a true improvement to the Transformer architecture, enabling DeepSeek R1 to generate tokens more than six times faster than a vanilla Transformer while actually improving algorithmic performance. Multi-head latent attention allows attention heads to share key and value information in a more optimal way, where the model itself learns how to compress and share this information between attention heads.
The Transformer architecture is one of the most significant breakthroughs in modern AI history, and DeepSeek appears to have just made it work significantly better. It's amazing to see the path that DeepSeek carved through their 2024 papers, systematically making substantial improvements to models that required hundreds of millions of dollars in R&D and infrastructure costs. The stakes have never been higher for neural networks. It will be fascinating to see what new set of ideas unlocks the next level of capabilities as we build more and more intelligent systems.
If you enjoyed the graphics in this video, I think you'll really like the poster version. The poster includes a walkthrough of multi-head latent attention with detailed captions, and I've rearranged the flow a bit to work better as a poster. On the bottom, I've included a detailed comparison between various forms of attention, including the required sizes of the KV caches, and a 3D model of each attention block. The matrix images in the video and poster are actually from the real DeepSeek model. I'm mostly showing the weights from the first layer of DeepSeek V3. The poster looks great in a simple frame that you can pick up on Amazon and is a great way to see how MLA works and just a nice way to decorate your walls. I'm offering free shipping on the poster for a limited time at WelchLabs.com, or you can pick it up as a limited edition bundle with a signed copy of my Imaginary Numbers book. Finally, big thank you to everyone who's purchased from the Welch Lab store; your purchases go a long way to helping me make more great videos.