How Mistral 7B Made Attention Efficient

Mistral 7B gave almost 2x speed-up as a result of modifications to attention mechanism. And it's not like this speed-up came with a performance drop. It achieved this speed-up while delivering performance comparable to and on several benchmarks even outperforming Llama 2 models much larger than itself. So, what exactly changed inside the attention mechanism? In standard self-attention used in transformers and let's focus on auto-regressive models for now. Each token attends to all tokens before it, including itself. So, given a sequence of n tokens, at every attention layer, we compute dot products between the queries and keys of all valid token pairs, giving us roughly n squared attention computations. This means that as context length grows, both the computational cost and memory usage scales quadratically as well. But in a long sequence of tokens, information that is really far away is often less important than nearby context. So, why compute attention scores with all previous tokens, no matter how distant they are from the current token? Mistral 7B uses this idea to move from full attention to sliding window attention. Instead of each token attending to all previous tokens, it only attends to tokens within a fixed window size. In this example on screen, each token attends only to itself and the two tokens immediately before it. Now, if you think about the number of attention scores that we actually need to compute for sliding window attention, you'll notice that each token only attends to w nearby tokens. So, instead of computing n squared attention scores, we now only need roughly n times w attention scores, which means the attention mechanism now scales linearly with sequence length, obviously assuming w is smaller than n. So, how do we actually compute this efficiently? One obvious approach would be to first materialize the full n cross n attention matrix by computing attention scores between all queries and keys, and then simply mask out all positions where the query and key are outside the sliding window boundary. While this would give the correct result, we still ended up doing n squared computations anyway. So, we did not really gain that much. Instead, I'll show the implementation used in Longformer, which uses chunked matrix multiplication. Given n tokens and a window size of w, and in this example, each token attends to two tokens before it and two tokens after it. We first compute Q and K representations. We then divide them into overlapping chunks of size 2w with an overlap of w tokens between two consecutive chunks. Now, for each chunk, we simply compute QK transpose. And if you look carefully, all the three chunked attention matrix already contain all the local attention scores that we need for sliding window attention. Let's talk about the computational complexity of this approach. Since the chunks are created with a stride of w, we end up with roughly n by w chunks. Each chunk involves a matrix multiplication of size proportional to w squared. So, the total computation becomes roughly O of n times w, and when w is smaller than n, this scales approximately linearly with sequence length. Similarly for memory, instead of materializing the full n cross n attention matrix, we now only store these local chunked attention matrix, which are much smaller, significantly reducing the memory requirement as well. Now, given all the tokens within the window, the attention mechanism can build contextual relationships between nearby tokens. But, what about global context? Any relationship outside the window boundary is something the model cannot directly attend to. So, doesn't sliding window attention completely lose access to long-range information? Well, the authors of Longformer argue that this is not entirely true. At a single layer, each token obviously only has direct access to tokens within a distance W. But, those neighboring tokens themselves had access to their own neighboring tokens in the previous layer. And this process continues across multiple layers. So, after multiple stacked layers of windowed attention, tokens can indirectly exchange information with tokens that are much farther away than the original window size. In fact, after L layers, the effective receptive field can grow to roughly L * W. This is very similar to how stacking convolution layers with small kernels in a CNN eventually produces a much larger receptive field, covering almost the entire image. However, in practice, this information propagation is not perfect. Tokens that are very far away only contribute indirectly through many intermediate layers. And the amount of information that survives tends to decay as it passes through repeated attention operations, softmax layers, residual connections, and normalization layers. So, while the theoretical receptive field becomes large, the effective influence of distant tokens can still become weak in practice. Because of this, a common modification is to combine local attention layers with occasional global attention layers. That way, information can still propagate efficiently across the entire sequence. For example, Gemma 2 uses a 1:1 ratio between local and global attention layers together with a 4K token sliding window. Gemma 3 moved to a 5:1 ratio while using a smaller 1024 token window. So, by interleaving local and global attention layers, we reduce computation using the local layers while still maintaining access to global context through the full attention layers. And for those full attention layers, techniques like group query attention can further reduce KV cache size. Speaking of KV cache, in standard auto regressive attention, each token attends to all previous tokens. So, as the context length grows, the KV cache also grows linearly with sequence length because we need to store the keys and values for all previous tokens. But with sliding window attention, each token only attends to previous W tokens. So, we only need to keep the W most recent keys and values in memory. And all the older KV cache entries can simply be evicted. Mistral 7B implements this using a rolling buffer cache. Basically, until we reach the window size, new keys and values are continuously added to the cache. But once we exceed the window size, older entries in the cache start getting overwritten. More specifically, the keys and values for time step I are stored at position I mod W inside the cache. So, once I becomes larger than W, earlier cache entries get overwritten in a circular fashion. As a result, the cache size never grows beyond W tokens while still storing all keys and values needed for sliding window attention computation. This leads to major memory savings while also reducing memory bandwidth bottlenecks during inference. Now, remember that KV cache memory scales with number of layers times number of heads times head dimension times sequence length times two times bytes per value. With sliding window attention, instead of storing KV values for all T previous tokens, we only store them for W tokens. So, on a 32K token sequence with a window size of 4K, sliding window attention reduces KV cache memory usage by roughly 8X while still maintaining strong model performance. For sliding window attention, that's all that you really need to know for understanding the evolution of attention mechanisms that we are going to be covering in this series. In the next video, we'll look at Deep Seek's multi-head latent attention and how it further reduces memory and computation costs. So, I'll see you there.