LLM Learning Roadmap
✨ Awesome-Anything
🖼️ Digital Image Processing
🍃 LLM Components
🌱 LLM Pre-training
☘️ LLM Post-Training
🍀 LLM Popular Models
🪴 LLM Applications
🌿 LLM Optimization
🌾 LLM Compression
🌵 LLM Hands-on Practice
🌴 LLM Must-read Papers
🌳 LLM Q&A
🐝 VLM Image Encoders
📝 MISC.
1.3 Attention
1.3 Attention
1.3.1 Introduction
Basic Idea
The core idea of the Attention mechanism when handling sequential data is that networks should focus more on the important parts of the input and ignore the less important parts. By learning weights for different parts, the model highlights important parts of the sequence, thus allowing the model to better focus on the input features that matter. Traditional Recurrent Neural Networks (RNNs) or Convolutional Neural Networks (CNNs) have difficulty capturing the varying importance of different positions in a sequence, which may lead to information distortion, especially when dealing with long sequences.
Understanding Attention
- Query (Q): This represents the “question” (like: What do I need to know about the rest of the sentence to understand this word?)
- Key (K): Think of this as a way for each word to say, "Here’s what I might be useful for.”
- Value (V): The content that gets passed along if the corresponding Key is relevant to the Query.
In a sequence, a token embedding at a certain position in the Query is dot-producted to the Keys at all positions to generate a set of corresponding
weights. Those resulting values of weights reflect how closely they match to each other. This helps the model understand which tokens are more relevant to the current token in the Query, rather than treating all tokens equally.Finally, the
weights represent how much information from other tokens is combined into this token's representation and will be used to multiply with the Value to get the final weighted outputs.Mathematical Formula
Scaled Dot-Product


“Scaled” in Scaled Dot-Product refers to scaling down the attention weights to ensure numerical stability, which is achieved by dividing by , which is the embedding size for a single head.
When the dimensions of Query and Key are relatively large, the dot product between the two vectors can have a large variance. The largest value dominates the softmax output, making it closer to 1, and the others closer to 0. This results in a sharper distribution, which may cause the model to over-focus because the softmax results become closer to one-hot, which may let the model focus on one uninformative token. This is undesirable in the early stages of training. Besides, locating at two extreme positions (0 or 1) leads to small gradients, which is also not desirable for the model to converge to a better state.
Scaling helps mitigate this issue by reducing the impact of the initial large variance.
In the early stages of training, a smaller gradient is often not beneficial. Since at the beginning, the model has not yet found good parameters. It's important that the model reaches a relatively good point through training. When the parameters of the model are already very close to the optimal solution, the gradient may be small at this time because the model convergence curve is already in a relatively flat area. At this time, the small gradient can help the model converge to the optimal solution more stably, avoiding oscillation and overfitting.
<ins/>
1.3.2 Attention in Transformers

Self-Attention in Encoder & Decoder
Q/K/V come from the same original sequence, only expecting to exchange information among tokens of the same input sequence. The main goal is to capture the dependency relationships within the sequence. In the Transformer’s encoder and decoder, every layer has self-attention mechanisms. It can focus each part of the input sequence on other parts. The difference lies in:
In the encoder, self-attention means each token in the current position attends to all tokens in the sequence.
In the decoder, self-attention means each token in the current position only attends to the previous tokens (Masked Attention / Causal Attention) to avoid information leakage during the decoding process.
Cross-Attention in Decoder
Q/K/V come from different sequences. Sometimes, we want to integrate information from another input into our current tokens. Queries come from the input sequence, and keys/values come from another sequence (outputs of the encoder). This allows the decoder to attend to the output of the encoder.
In many scenarios, this can be useful — for example, in machine translation, where the decoder (generate target language) attends to the encoder’s output (encoded source language), which is beneficial for aligning different languages.
1.3.3 Computation Complexity Optimization
Self Attention has a complexity of . It needs to compute the correlation between any two token embeddings in the sequence, resulting in a correlation matrix of size .

Sparse Attention
A mixture of dilated attention and local attention.

Looking at the attention matrix, all attention weights for tokens whose relative distance is beyond a threshold k and between nk and (n+1)k are set to 0. In this way, Attention has the characteristic of "dense attention for nearby tokens and sparse attention for distant tokens". For many cases, this might be a reasonable assumption, because tasks that truly require dense long-range dependencies are rare.
However, this approach clearly has two shortcomings:
- How to choose the threshold k is determined manually, which involves a high degree of arbitrariness.
- It often requires special design and optimization at the implementation level to achieve efficient performance, making it hard to generalize.
Linear Attention
The main idea is to remove the softmax, and first compute , so the computation complexity between Q and K changes from to approximately linear .
Softmax is used to get non-negative attention weights from . If there are other ways to achieve this, then we can remove softmax and compute the matrix product first. Note that the dimensions of Q, K, and V are all in the self-attention case.
We can use a kernel function to linearize the attention form. Using two activation functions, the formula becomes:
This is equivalent to:
Here, we still do the normalization by dividing by the sum of the weights, but this is different from the softmax normalization because the softmax first applies exponential operation and then normalizes the weights. This only works if the function maps to positive values to preserve numerical stability and meaningful normalization.
Another way is we can still adopt softmax to keep the normalization and assure non-negative attention weights. Note that is already normalized if Q is normalized in dimension and K is normalized in dimension , then we can use the following formula to calculate attention:
Here, , denote softmax operations across the first and second dimensions, respectively. In other words, now we apply softmax separately to Q and K, rather than applying softmax to the product . This design has already been used in computer vision multiple times, such as in A2-Nets.

One implementation, full code at: https://github.com/lucidrains/linear-attention-transformer
Prove that normalizing Q in dimension (rows) and K in dimension (columns) results in a normalized .
- Normalization Assumptions:
- For Q (shape ): Apply normalization (e.g., softmax) along dimension . This ensures each row of Q sums to 1.
- For K (shape ): Apply normalization along dimension . This ensures each column of K sums to 1.
- Matrix Multiplication :
- The product has shape .
- The element in is the dot product of the i-th row of Q and the j-th row of K (or equivalently, the j-th column of ).
- Row-wise Normalization of :
- Compute the sum over the j-th column of for the i-th row:
- Swap the order of summation:
- Since K is column-normalized, for all . Thus:
- Conclusion:
- Each row of sums to 1. This guarantees that the attention weights are row-normalized without additional softmax.
L2 Norm or Sum Norm?
Above derivation is based on the fact that we are doing sum norm (e.g. softmax) instead of L2 norm.
- Sum Norm
- then let
- , ,
- So a + b + c = 1
Note that softmax is also one form of sum norm. The sum norm is usually used for probabilities like in Attention.
- L2 Norm
- then let , ,
- So
The L2 norm is usually used for unit vectors in geometry or machine learning
<ins/>
1.3.4 KV Cache
In one sentence:
K and V refer to the key and value states in the attention mechanism. The KV cache only appears in the autoregressive decoder of the transformer architecture, and it's used to avoid redundant computation in the scaled dot-product attention process.Principles
KVCache:
In mainstream autoregressive LLMs using Causal Attention, when generating tokens recursively token by token, each time the attention for the current token is computed, it requires all previous tokens’ K and V.
To avoid recalculating this part repeatedly, the key and value computed from previous tokens can be cached — this is known as the KVCache technique.
As shown in the figure below, the third row of is only determined by the third row of Q and first three columns and rows of K and V, respectively.

- When generating tokens without KV Cache
If the model outputs are
“I show speed” , when we generate the first token “I” :
Next we generate the second token
“show” , the input is “<s>I” :
In the above, the
Mask is for causal inference, which means the token “<s>” cannot see tokens “I”, “show”, and “speed”, and the token “I” cannot see token “show” and “speed”, and so on.You can see that, in the second generation, we still calculate the , V, and attention for the first row (i.e. token
“<s>”), although they are already computed during the first generation. Therefore, the key idea of KV Cache is to cache all previous K and V, so that we can focus on only generating the attention for the current token.Without KV Cache, we cannot drop the compute for previous tokens, because in the next layer of Attention, all previous tokens’ attentions/outputs in the current Attention are used as K and V. As long as we cache all KV for each layer of Attention, we can drop these redundant computation and save a lot of inference resource and time!
- A simplified snippet of KV Cache implementation in Huggingface
KV Cache GPU Memory Usage Analysis
In transformer-based models (like GPT), KV cache stores key and value tensors during autoregressive generation to speed up inference.
KV Cache Shape and Formula (for a single layer of Attention)
For
kv_seq_len KV values, the shape is:Where:
b: batch size
head_num (h): number of attention heads
kv_seq_len (s + n): total sequence length (input + output)
head_dim (d): dimension per head
Assuming FP16 precision (2 bytes per element), the peak KV cache memory usage is:
Explanation for the number 4:
2for keys and values (K/V)
2for FP16 (2 bytes)
- So total multiplier:
4
Example: GPT-3 Scale (175B Parameters)
To illustrate, take GPT-3 with:
h = 96(number of heads)
d = 128(embedding dimension per head)
L = 96(layers of Attention)
- Assume
b = 1,s + n = 2048(context length)
Then:
Meanwhile, GPT-3 model weights in FP16 take up approximately 350 GB.
KV Cache Optimization Methods
- Shared KV Cache: MQA, GQA
- Window Optimization:
- Fixed Window Length (Figure b):
- KV Recomputation (Figure c):
- Λ-shaped Attention Window:
The purpose of the KV cache is to compute attention. When the text length during inference (T) is greater than the maximum context length during training (L), a natural idea is a sliding window. There are three methods here:
A typical example is Longformer. This method is easy to implement, with a space complexity of O(TL), but may sacrifice some precision.
This method recalculates the KV cache every time. Thanks to the recomputation, accuracy can be guaranteed, though it may be less efficient.
In (c), the tokens within context length L will be used for recomputing K/V by treating them as in position starting from 0, 1, …; but in (b), these cached tokens’ K/V are computed previously by treating them as in position, for example, 4, 5, … This is the most important difference. Because the text with length T is greater than training context window L, all tokens should be rearranged to start from position 0, otherwise the cached K/V involves more wide information outside the context window L, which forms out-of-distribution features and confuse the model, leading to much higher perplexity (PPL).
Already proposed in LM-Infinit, and its principle aligns with StreamingLLM. By involving the first few tokens, the model can expect the same distributions of features (not much information at first and more and more information later on) and acquire more information from the start (more aligned with the “attention sink” behavior of the model), this largely improve the performance in terms of PPL. Also, this design didn’t add up too much computation burden, so the time complexity is still .

- Quantization and Sparsity:
Currently, mainstream inference frameworks are increasingly supporting KV Cache quantization. A typical example is lmdeploy.
- Storage and Computation Optimization:
- vLLM’s PagedAttention: Simply put, it allows storing the K and V values in non-contiguous memory spaces.
- FlashDecoding is an inference-time optimization based on FlashAttention, mainly divided into three steps:
- Split long sequences of KV into smaller chunks that are easier to process in parallel
- For each chunk’s KV and Q, perform FlashAttention as before to get the result for that chunk
- Reduce the results of all chunks

<ins/>
1.3.5 MHA to MLA: trade-off between efficiency and performance

Multi-Head Attention (MHA)
In the original Transformer, the attention mechanism is MHA. Each head has a
head_size, which is equal to the original embed_size divided by num_head. It's similar to group convolution, like establishing many communication channels, where each channel focuses on different details of the information. In other words, each head can attend to features in different subspaces of the sequence.Multi-Query Attention (MQA)
All queries share KV. The compression of KV Cache by MQA is too aggressive (from
hidden_size to head_dim), which may impact the learning efficiency and final performence of the model.Grouped-Query Attention (GQA)
Divide the query into g groups, with queries in each group sharing KV to balance performance and memory usage.
<ins/>
Multi-Head Latent Attention (MLA)
Core issue addressed: When using KV cache, the sequence length increases over time, leading to insufficient memory.

Generally, MHA splits QKV into heads:
During inference, all keys and values need to be cached for acceleration. Therefore, MHA needs to cache elements per token, where is the number of attention layers. When deploying the model, this large KV cache becomes a bottleneck, limiting batch size and sequence length.
MLA improvements:
- Low-rank KV compression:
- is the hidden state projected to a lower-dimensional latent space, where , is the compression dimension for KV.
- is the down-projection matrix.
- are projection matrices for key and value:

During inference, MLA only needs to cache , only elements in total for all attention layers. Also, can be fused into (query and output matrices), so we don't need to separately compute Q/K/V anymore.
Additionally, MLA performs low-rank projection on query as well to reduce training-time activation memory:
- Decoupling RoPE
- RoPE is incompatible with low-rank KV Cache. During compression, it’s possible to fuse into , but after adding RoPE, we cannot do it.
- MLA's Solution: Use extra multi-head queries and shared keys to integrate RoPE
- : Dimension of a single RoPE head
- : Used to generate extra multi-head queries and shared keys
- During inference, the shared key needs caching, so the total cached elements is
In RoPE: . There is a relative position matrix between the query Q and key K. Therefore, cannot be fused to form a fixed projection matrix.

Where:

- MLA computation during training and inference
- Training
- Inference
During the inference stage, the Head Size of Q and K becomes , and the Head Size of V (now it’s ) becomes . According to the settings in the original paper, this is 4 times the size of and . So, in fact, what MLA does during inference is this absorbing matrix, which can effectively reduce the KV Cache, but the amount of computation during inference increases.
Then why can it still improve inference efficiency? We can divide LLM inference into two parts: Prefilling and token-by-token Generation.
The prefill stage involves parallel computation of all input tokens and then storage of the corresponding KV cache. In this part, both computation and memory bandwidth are bottlenecks. MLA indeed increases the computation, but the reduction of the KV Cache also eases the memory and bandwidth pressure—it's a tradeoff.
In the generation/inference stage, since only one current token is computed per step, the computation is not a headache. Slightly increasing the computation would not cause too much problem. Therefore, because of the reduction of KV Cache, MLA can significantly speed up generation theoretically.
- MLA Implementation
Core Configuration Parameters:
Note that, currently, the official implementation of Deepseek V2 does not store latent states when saving the KV Cache. Instead, it decompresses the latents into standard MHA KV Cache format. Therefore, this does not save any memory at all (it's just how it's implemented in V2 Attention).
<ins/>
1.3.6 Dual Chunk Attention (DCA)
The core idea of DCA is to divide a long text into multiple smaller “chunks”, and then apply attention mechanisms both within and between these chunks. The specific steps are as follows:
- Chunking:
Divide the long text into several small chunks, with each chunk containing a portion of the text. For example, a 2000-word text can be split into 4 chunks of 500 words each.
- Intra-chunk Attention:
Apply attention mechanisms within each chunk independently. This means that the words in each chunk only attend to other words within the same chunk. This significantly reduces the computational cost.
- Inter-chunk Attention:
After calculating intra-chunk attention, apply attention mechanisms between the chunks. This means each chunk can attend to others to exchange global information and capture broader context across the entire text.
By using this approach, DCA can efficiently handle long texts while maintaining relatively high computational efficiency and low memory usage.

1.3.7 Shifted Sparse Attention (S2-Attention)
Source code: https://github.com/dvlab-research/LongLoRA
- Motivation
The traditional dense Attention has a computational complexity of O(n²), which is extremely costly. To address this issue, S2-Attention adopts a mechanism of local computation combined with information flow, improving efficiency while maintaining information interaction.
- Details
The context length is divided into several groups, and attention is computed independently within each group (sequence dimension). In half of the attention heads, tokens are shifted by half a group size to ensure information flow between neighboring groups (attention head dimension). For example, using S2-Attn with a group size of 2048 simulates training with a context length of approximately 8192 tokens (like stacking multiple CNN layers to broaden the receptive field).

Although shifting may introduce some information leakage (i.e. past tokens see future tokens), this can be mitigated by slightly adjusting the attention mask.

Prev
1.2 Embedding
Next
1.4 FFN, Residual Addition, and LN
Loading...