Attention Mechanisms - All You Needed To Know

Abstract

Table of Contents

1. Introduction

The concept of Attention, was first introduced in the 2017 paper “Attention Is All You Need”, that actually revolutionized how neural networks process sequencial data¹. Before it sequential data modeling relied heavily on Recurrent Neural Networks (RNNs) and LSTMs. They are effective onto certain low level task but these architectures process data sequentially that creates bottleneck to make parallelization impossible and also caused them to “forget” earlier parts of a information.

With introduction of the attention mechanism, how we handle sequence data has totally changed. It represents the projection of input data into specific subspaces to study the importance of each token relative to every other token in a sequencial data which answers a critical question “How does a given token contribute to learning the embedding with respect to all other tokens in the sequence ?”

2. Fundamental Attention Mechanisms

2.1 Self Attention

Consider that our sequence of token has a numerical representation is of a matrix form

where, $x_i \in \mathbb{R}^{d_m}$, $T$ be the number of tokens in sequence and $d_m$ be the dimension of each token vector. For each token in our sequence, we compute three distinct subspace representations called the Query, Key and Value denoted by $Q, K \; \text{and} \; V$ resp. We generate these by projecting our input matrix $X$ using three separate, learnable operator matrices ($W_q$, $W_k$, and $W_v$) as follows

and each row of $Q, K \; \text{and} \; V$ corresponds to the subspace representation of a single input token from $X$.

1. Query : What a token is “looking for” in the sequence ?
2. Key : It’s the “identity” of the token that others can match with.
3. Value : It Contains the actual semantic information of the token.

Remark : If we assume $d_q = d_k = d_v$, the projection will be to the same dimension, but into completely different conceptual subspaces.

Similarity Score

To fully understand how tokens interact or What is the similarity of a particular token with all other tokens in the sequence $X$? or what is the amount of information there in a query regarding the other tokens in sequence ? Let $q \in Q$ be a query and $k_i \in K$ be $i^{th}$ key. We compute the dot product of query $q$ with the all keys. For the $i$-th key, the similarity $s_i$ is $s_i = q k_i^T$ and for the entire sequence it will be

Quest : How do we prevent the numbers from blowing up during the computation of Attention? Because we are doing lots of multiplications, multiplying numbers greater than 1 produces even larger numbers. This pushes the subsequent Softmax function into regions with extremely small gradients.

Ans : We must use a normalized value rather than the original multiplied values. We do this by scaling the dot product by the square root of the key dimension $\sqrt{d_k}$:

Did normalized values make any sense to us on their own ? Not quite. But we can convert them into a probability distribution via the Softmax function $\alpha : \mathbb{R}^T \longrightarrow \mathbb{R}^{[0,1]}$ defined as follows

For an individual token interaction, this looks like

Here, $0 \le \alpha_i \le 1$ and $\sum_{i=1}^{T} \alpha_i = 1$. The value $\alpha$ represents the scaled similarity between a query and all keys in the sequence.

Attention Weights

Now, we use these probabilities to calculate Attention as a weighted sum of the values $v_i$

and for all query tokens, the output will be a set of vectors as a mapping $A : \mathbb{R^{T}} \times \mathbb{R^{T \times d_v}} \to \mathbb{R^{d_v}}$ defined as follows

Hence : Attention is a vectorial representation of tokens that dynamically captures how one token interacts with others. The actual interaction happens due to the Value ($V$) subspace; the Query ($Q$) and Key ($K$) subspaces exist purely to compute the vectorial weights.

## 1. Standard Self-Attention Algorithm (Encoder-Style)

This algorithm represents the unmasked bidirectional attention, where every token can perfectly "see" every other token in the sequence (past, present, and future). This is the core operation inside a Transformer Encoder (like BERT).

**Input:**
* $X \in \mathbb{R}^{T \times d_{model}}$ (Input sequence embeddings)

**Learnable Parameters:**
* $W_Q \in \mathbb{R}^{d_{model} \times d_k}$ (Query projection matrix)
* $W_K \in \mathbb{R}^{d_{model} \times d_k}$ (Key projection matrix)
* $W_V \in \mathbb{R}^{d_{model} \times d_v}$ (Value projection matrix)
* $W_O \in \mathbb{R}^{d_v \times d_{model}}$ (Output linear projection matrix)

**Algorithm:**
1.  **Project Subspaces:**
    * $Q \leftarrow X W_Q$  *(Shape: $T \times d_k$)*
    * $K \leftarrow X W_K$  *(Shape: $T \times d_k$)*
    * $V \leftarrow X W_V$  *(Shape: $T \times d_v$)*
2.  **Compute Raw Interaction Scores:**
    * $S_{raw} \leftarrow Q K^T$ *(Shape: $T \times T$)*
3.  **Scale Scores (Variance Control):**
    * $S_{scaled} \leftarrow \frac{S_{raw}}{\sqrt{d_k}}$
4.  **Compute Alignment Probabilities:**
    * $P \leftarrow \text{Softmax}(S_{scaled})$ *(Softmax applied row-wise. Shape: $T \times T$)*
5.  **Aggregate Contextual Values:**
    * $A \leftarrow P \cdot V$ *(Shape: $T \times d_v$)*
6.  **Final Linear Projection:**
    * $\text{Output} \leftarrow A W_O$
7.  **Return** $\text{Output} \in \mathbb{R}^{T \times d_{model}}$

2.2 Mask Attention

Transformer consists of two primary components that handle this sequence data differently:

Because the Decoder operates auto-regressively, we must ensure its Attention mechanism strictly depends on the past and present, not the future. We achieve this by modifying our soft attention mechanism slightly using a Mask Matrix.

Consider a mask matrix $M \in \mathbb{R}^{T \times T}$ with entries defined as:

We add this mask to our scaled similarities before applying the Softmax:

This is called Masked Attention.

How it works: The mask matrix does nothing to the past and present tokens (adding $0$ leaves the value unchanged). However, for future tokens, it forces the value to $-\infty$. When passed through the Softmax function, $e^{-\infty}$ becomes $0$. This ensures the Attention weights for all future tokens are exactly zero. Hence, the generation step successfully depends only on the past and present, preserving the integrity of auto-regressive generation.

## 2. Masked Self-Attention Algorithm (Decoder-Style)

This algorithm introduces the causal boundary required for auto-regressive generation. By enforcing a strict lower-triangular mask, we guarantee that the prediction of token $t+1$ depends exclusively on tokens $1$ through $t$, preventing any data leakage from the future. This is the core operation inside a Transformer Decoder (like GPT).

**Input:**
* $X \in \mathbb{R}^{T \times d_{model}}$ (Input sequence embeddings)

**Learnable Parameters:**
* $W_Q \in \mathbb{R}^{d_{model} \times d_k}$ 
* $W_K \in \mathbb{R}^{d_{model} \times d_k}$ 
* $W_V \in \mathbb{R}^{d_{model} \times d_v}$ 
* $W_O \in \mathbb{R}^{d_v \times d_{model}}$ 

**Algorithm:**
1.  **Construct the Causal Mask ($M$):**
    * Initialize $M \in \mathbb{R}^{T \times T}$ as a matrix of zeros.
    * **For** $i = 1$ **to** $T$ **do:** *(Rows / Queries)*
        * **For** $j = 1$ **to** $T$ **do:** *(Columns / Keys)*
            * **If** $j > i$ **then:** *(If the Key is in the future relative to the Query)*
                * $M[i, j] \leftarrow -\infty$
    * *(Result: A lower-triangular matrix of $0$s with $-\infty$ in the strictly upper triangle).*
2.  **Project Subspaces:**
    * $Q \leftarrow X W_Q$ 
    * $K \leftarrow X W_K$ 
    * $V \leftarrow X W_V$ 
3.  **Compute and Scale Interaction Scores:**
    * $S \leftarrow \frac{Q K^T}{\sqrt{d_k}}$ *(Shape: $T \times T$)*
4.  **Apply Causal Mask:**
    * $S_{masked} \leftarrow S + M$ 
    * *(Effect: All future interaction scores are driven to $-\infty$)*
5.  **Compute Alignment Probabilities:**
    * $P \leftarrow \text{Softmax}(S_{masked})$ 
    * *(Effect: $e^{-\infty} = 0$, completely nullifying future tokens' weights. Probabilities of valid past/present tokens still sum to 1).*
6.  **Aggregate Contextual Values:**
    * $A \leftarrow P \cdot V$ 
7.  **Final Linear Projection:**
    * $\text{Output} \leftarrow A W_O$
8.  **Return** $\text{Output} \in \mathbb{R}^{T \times d_{model}}$

2.3 Cross Attention

What happens when we are dealing with a Sequence-to-Sequence (Seq2Seq) task, like translating a Hindi sentence into an English sentence ? Or a multi-modal task, like generating a text caption for an image ? The generated sequence (English) needs a way to dynamically “look back” at the source sequence (French) at every step of generation to figure out what to translate next. Self-attention cannot do this because the source and target are two completely different sequences.

So, with the help of cross attention we able to uses the exact same mathematical engine as Self-Attention, but it intentionally decouples the origin of the Queries from the Keys and Values. How we actually do this ? In cross attention, there are always two distinct sequences involved

1. The Source Sequence (Encoder Output): This is the fully processed input data (e.g., the French sentence or the Image embeddings). It provides the Keys ($K$) and Values ($V$). It represents “what information is available.”
2. The Target Sequence (Decoder State): This is the sequence currently being generated. It provides the Queries ($Q$). It represents “what information I am currently looking for.”

By projecting $Q$ from the target and $K, V$ from the source, the resulting attention matrix maps the relationship between two entirely different spaces, allowing the decoder to selectively extract relevant context from the encoder.

Let our processed Source sequence (Encoder output) be $X_{enc} \in \mathbb{R}^{T_e \times d_{model}}$, where $T_e$ is the length of the source sequence. Let our current Target sequence (Decoder hidden states) be $X_{dec} \in \mathbb{R}^{T_d \times d_{model}}$, where $T_d$ is the length of the generated sequence so far.

(Note: $T_e$ and $T_d$ can be completely different lengths. This is a crucial feature of Cross Attention).

Step 1: Decoupled Subspace Projections Instead of projecting all three matrices from the same input, we project them from their respective origins:

Step 2: The Cross-Attention Matrix We compute the scaled dot-product attention just as before:

Let’s look at the dimensions of this matrix multiplication:

This is the beautiful part of Cross Attention. The resulting score matrix $S$ is of size $T_d \times T_e$. Each row corresponds to a single token in the generated output, and the columns represent its similarity to every single token in the source input. ### Step 3: Probabilities and Value Extraction We apply the Softmax function to convert these scores into probabilities. Crucial Note on Masking: Unlike Self-Attention in the decoder, Cross Attention typically does not use a causal/auto-regressive mask ($-\infty$ upper triangle). Why? Because the entire source sequence (Encoder output) is already fully known. The decoder is allowed to look at any part of the source sequence at any time. (It may, however, use a padding mask to ignore <PAD> tokens in the source).

Finally, we multiply by the Source Values:

The output maps perfectly back to the length of the Decoder sequence ($T_d$), having successfully infused it with the exact, dynamically weighted information it needed from the Encoder sequence.

## Cross Attention Algorithm(Encoder-Decoder Style)
It highlights the dual-input nature of the Cross Attention layer. Today, Cross Attention is not just for text translation anymore. It is the mathematical bridge powering modern AI. In **Stable Diffusion**, Cross Attention is how the image noise (Queries) attends to our text prompt (Keys/Values). In **Vision-Language Models**, it is how the text generator (Queries) attends to the image patches (Keys/Values) to answer questions about a picture.

**Input:**
* $X_{enc} \in \mathbb{R}^{T_e \times d_{model}}$ (Source/Encoder embeddings)
* $X_{dec} \in \mathbb{R}^{T_d \times d_{model}}$ (Target/Decoder embeddings)
* Optional: $M_{pad} \in \mathbb{R}^{T_d \times T_e}$ (Padding mask for source sequence)

**Learnable Parameters:**
* $W_Q \in \mathbb{R}^{d_{model} \times d_k}$ (Query projection weight)
* $W_K \in \mathbb{R}^{d_{model} \times d_k}$ (Key projection weight)
* $W_V \in \mathbb{R}^{d_{model} \times d_v}$ (Value projection weight)
* $W_O \in \mathbb{R}^{d_v \times d_{model}}$ (Output linear projection weight)

**Algorithm:**
1.  **Project Target into Query Space:**
    * $Q \leftarrow X_{dec} W_Q$
2.  **Project Source into Key/Value Space:**
    * $K \leftarrow X_{enc} W_K$
    * $V \leftarrow X_{enc} W_V$
3.  **Compute Cross-Similarity Scores:**
    * $S \leftarrow \frac{Q K^T}{\sqrt{d_k}}$ *(Shape will be $T_d \times T_e$)*
4.  **Apply Padding Mask (If applicable):**
    * $S \leftarrow S + M_{pad}$ *(Masks out padded tokens in the source)*
5.  **Compute Alignment Probabilities:**
    * $P \leftarrow \text{Softmax}(S)$ *(Softmax applied across the $T_e$ dimension)*
6.  **Extract Source Information:**
    * $O_{cross} \leftarrow P \cdot V$ *(Shape returns to $T_d \times d_v$)*
7.  **Final Linear Projection:**
    * $\text{Output} \leftarrow O_{cross} W_O$
8.  **Return** $\text{Output} \in \mathbb{R}^{T_d \times d_{model}}$

3. Advanced Attention Mechanisms

3.1 Multi Head Attention(MHA)

A single self-attention function computes the interactios between tokens based on set of projects($Q$, $K$, $V$) which limits the model’s ability to focus on different types of relationships simultaneously. For example, in a sentence, one token might need to attend to a grammatical subject, another to an adjective, and another to a structural dependency. If we only have one set of attention weights (a single “head”), the model is forced to average these distinct relationships into a single subspace. Multi-Head Attention allows the model to jointly attend to information from different representation subspaces at different positions.

Let our input sequence be represented by $X \in \mathbb{R}^{T \times d_{model}}$, where $d_{model}$ is the embedding dimension of the network.

Instead of performing a single attention function with dimension $d_{model}$, we linearly project the queries, keys, and values $h$ times (where $h$ is the number of heads) using different, learnable weight matrices.

For each head $i \in {1, 2, \dots, h}$, we define the projection matrices

To keep computational cost similar to single-head attention with full dimensionality, the dimensions are typically reduced such that $d_k = d_v = d_{model} / h$.

Step 1: Independent Subspace Projections For each head $i$, we compute the distinct queries, keys, and values:

Step 2: Parallel Attention Computation We apply the Scaled Dot-Product Attention function to each head independently in parallel

Where the output of each head is

Step 3: Concatenation and Final Projection Once all $h$ heads have computed their outputs, we concatenate them along the feature dimension.

Since $h \cdot d_v = d_{model}$, the concatenated matrix is back to $\mathbb{R}^{T \times d_{model}}$.

Finally, we project this concatenated matrix back into the expected output space using a final learnable weight matrix $W_O$:

Where $W_O \in \mathbb{R}^{d_{model} \times d_{model}}$.

## 2. Mathematical Pseudo-Code Algorithm

This algorithm represents the exact operations required to implement the forward pass of a Multi-Head Attention layer.

**Input:**
* $X \in \mathbb{R}^{T \times d_{model}}$ (Input sequence embeddings)
* $h \in \mathbb{Z}^+$ (Number of attention heads)
* Optional: $M \in \mathbb{R}^{T \times T}$ (Mask matrix for auto-regressive decoding)

**Learnable Parameters:**
* $\{W_Q^{(i)}, W_K^{(i)}, W_V^{(i)}\}_{i=1}^h$ (Subspace projection weights)
* $W_O \in \mathbb{R}^{d_{model} \times d_{model}}$ (Output linear projection weight)

**Algorithm:**
1.  **Initialize** $d_k = d_v = d_{model} / h$
2.  **For** $i = 1$ **to** $h$ **do in parallel:**
    * **Project Queries:** $Q_i \leftarrow X W_Q^{(i)}$
    * **Project Keys:** $K_i \leftarrow X W_K^{(i)}$
    * **Project Values:** $V_i \leftarrow X W_V^{(i)}$
    * **Compute Scores:** $S_i \leftarrow \frac{Q_i K_i^T}{\sqrt{d_k}}$
    * **Apply Mask (If Decoder):** $S_i \leftarrow S_i + M$
    * **Compute Probabilities:** $P_i \leftarrow \text{Softmax}(S_i)$
    * **Compute Head Output:** $\text{head}_i \leftarrow P_i V_i$
3.  **End For**
4.  **Concatenate Heads:** $H_{\text{concat}} \leftarrow [\text{head}_1 \,\|\, \text{head}_2 \,\|\, \dots \,\|\, \text{head}_h]$ 
    *(Note: $\|$ denotes column-wise concatenation)*
5.  **Final Projection:** $\text{Output} \leftarrow H_{\text{concat}} W_O$
6.  **Return** $\text{Output} \in \mathbb{R}^{T \times d_{model}}$

3.2 Multi Query Attention(MQA)

Multi-Head Attention (MHA) drastically improves a model’s ability to capture diverse linguistic relationships, it introduces a massive computational bottleneck during auto-regressive decoding(inference) called kv-cache bottleneck.

During generation, a decoder must predict the next token one step at a time. To avoid recomputing the Keys and Values for all past tokens at every step, Transformers use a KV Cache—storing the previously computed $K$ and $V$ tensors in memory.

• In standard Multi-Head Attention (MHA), for $h$ heads, we must store $h$ distinct Keys and $h$ distinct Values for every single token in the sequence.
• The Problem: As the sequence length and batch size grow, the memory bandwidth required to load this massive KV Cache from VRAM to the compute cores becomes the primary bottleneck. The GPU spends more time waiting for data than performing matrix multiplications.

The Solution: Multi-Query Attention (introduced by Noam Shazeer in 2019) elegantly solves this memory bottleneck by projecting $h$ different sets of Queries, but sharing a single set of Keys and Values across all heads.

Let our input sequence be $X \in \mathbb{R}^{T \times d_{model}}$. In MQA, we still have $h$ different query projections, but only one global key projection and one global value projection.

We define the projection matrices as

(Note: Just like in MHA, $d_k = d_v = d_{model} / h$, but we only instantiate one $W_K$ and one $W_V$.)

Step 1: Subspace Projections (The Difference)

We compute $h$ distinct Queries, but only one Key and one Value matrix for the entire layer:

Step 2: Parallel Attention Computation

We apply the Scaled Dot-Product Attention function to each head $i$ in parallel, but they all broadcast against the exact same $K$ and $V$:

Where the output of each head remains:

Step 3: Concatenation and Final Projection

Just as in MHA, we concatenate the $h$ heads and project them back to the model dimension:

Where $W_O \in \mathbb{R}^{d_{model} \times d_{model}}$.

## 3. Mathematical Pseudo-Code Algorithm

Notice how the Key and Value projections are moved *outside* the parallel loop, drastically reducing the memory footprint of the KV Cache.

**Input:**
* $X \in \mathbb{R}^{T \times d_{model}}$ (Input sequence embeddings)
* $h \in \mathbb{Z}^+$ (Number of query heads)
* Optional: $M \in \mathbb{R}^{T \times T}$ (Mask matrix for auto-regressive decoding)

**Learnable Parameters:**
* $\{W_Q^{(i)}\}_{i=1}^h$ (Query subspace projection weights)
* $W_K \in \mathbb{R}^{d_{model} \times d_k}$ (Single Key projection weight)
* $W_V \in \mathbb{R}^{d_{model} \times d_v}$ (Single Value projection weight)
* $W_O \in \mathbb{R}^{d_{model} \times d_{model}}$ (Output linear projection weight)

**Algorithm:**
1.  **Initialize** $d_k = d_v = d_{model} / h$
2.  **Project Shared Keys:** $K \leftarrow X W_K$
3.  **Project Shared Values:** $V \leftarrow X W_V$
4.  **For** $i = 1$ **to** $h$ **do in parallel:**
    * **Project Queries:** $Q_i \leftarrow X W_Q^{(i)}$
    * **Compute Scores:** $S_i \leftarrow \frac{Q_i K^T}{\sqrt{d_k}}$  *(Note: $K$ is shared)*
    * **Apply Mask (If Decoder):** $S_i \leftarrow S_i + M$
    * **Compute Probabilities:** $P_i \leftarrow \text{Softmax}(S_i)$
    * **Compute Head Output:** $\text{head}_i \leftarrow P_i V$ *(Note: $V$ is shared)*
5.  **End For**
6.  **Concatenate Heads:** $H_{\text{concat}} \leftarrow [\text{head}_1 \,\|\, \text{head}_2 \,\|\, \dots \,\|\, \text{head}_h]$ 
7.  **Final Projection:** $\text{Output} \leftarrow H_{\text{concat}} W_O$
8.  **Return** $\text{Output} \in \mathbb{R}^{T \times d_{model}}$

By sharing $K$ and $V$, the size of the KV Cache is reduced by a factor of $h$. This allows for vastly larger batch sizes during inference and dramatically faster decoding speeds, at the cost of a very minor degradation in model capacity (which led to the later development of Grouped-Query Attention as a middle ground).

3.3 Group Quary Attention(GQA)

Multi-Query Attention (MQA) successfully solved the KV Cache memory bottleneck by sharing a single Key and Value head across all queries, it introduced a new problem: capacity degradation. By forcing the model to compress all Key and Value representations into a single subspace, MQA slightly degrades the model’s ability to capture complex, multi-faceted nuances compared to standard Multi-Head Attention (MHA).

The Motivation: We need a middle ground. Can we retain the high inference speed and low memory footprint of MQA, while keeping the high representational quality of MHA?

Grouped-Query Attention (GQA) is the solution. It smoothly interpolates between MHA and MQA by dividing the query heads into $g$ distinct groups. Each group shares a single Key and Value head.

• If $g = h$ (groups equal to number of heads), GQA becomes MHA.
• If $g = 1$ (one single group), GQA becomes MQA.
• If $1 < g < h$, we achieve the optimal balance of speed and quality.

Let our input sequence be $X \in \mathbb{R}^{T \times d_{model}}$. Let $h$ be the total number of query heads, and $g$ be the number of groups. The number of queries per group is $k = \frac{h}{g}$.

We define our projection matrices as:

• $h$ distinct Query projections: $W_Q^{(i)} \in \mathbb{R}^{d_{model} \times d_k} \quad \text{for } i \in {1, \dots, h}$
• $g$ distinct Key projections: $W_K^{(j)} \in \mathbb{R}^{d_{model} \times d_k} \quad \text{for } j \in {1, \dots, g}$
• $g$ distinct Value projections: $W_V^{(j)} \in \mathbb{R}^{d_{model} \times d_v} \quad \text{for } j \in {1, \dots, g}$

(Again, $d_k = d_v = d_{model} / h$)

Step 1: Grouped Subspace Projections We compute $h$ Query matrices, but only $g$ Key and Value matrices:

Step 2: Parallel Attention Computation with Group Mapping To compute the attention for a specific query head $i$, we must map it to its corresponding group $j$. The index mapping is simply $j = \lfloor \frac{i-1}{k} \rfloor + 1$ (assuming 1-based indexing).

Therefore, multiple query heads in the same group will broadcast against the same group-specific $K_j$ and $V_j$:

Where the output of each head remains:

Step 3: Concatenation and Final Projection As with all multi-head variants, we concatenate all $h$ computed heads and project them back to the original model dimension.

Where $W_O \in \mathbb{R}^{d_{model} \times d_{model}}$.

####  Mathematical Pseudo-Code Algorithm

This algorithm highlights the two-tiered looping structure: one loop for the groups to fetch the shared KV matrices, and an inner loop for the queries belonging to that specific group.

**Input:**
* $X \in \mathbb{R}^{T \times d_{model}}$ (Input sequence embeddings)
* $h \in \mathbb{Z}^+$ (Total number of query heads)
* $g \in \mathbb{Z}^+$ (Total number of KV groups, where $h$ is divisible by $g$)
* Optional: $M \in \mathbb{R}^{T \times T}$ (Mask matrix)

**Learnable Parameters:**
* $\{W_Q^{(i)}\}_{i=1}^h$ (Query projections)
* $\{W_K^{(j)}, W_V^{(j)}\}_{j=1}^g$ (Key and Value group projections)
* $W_O \in \mathbb{R}^{d_{model} \times d_{model}}$ (Output linear projection weight)

**Algorithm:**
1.  **Initialize** $d_k = d_v = d_{model} / h$
2.  **Initialize** $k = h / g$ *(Number of queries per group)*
3.  **Initialize empty list:** $H_{\text{list}} = []$
4.  **For** $j = 1$ **to** $g$ **do in parallel:**
    * **Project Group Keys:** $K_j \leftarrow X W_K^{(j)}$
    * **Project Group Values:** $V_j \leftarrow X W_V^{(j)}$
    * **For** $m = 1$ **to** $k$ **do in parallel:** *(Iterate over queries in this group)*
        * **Determine global query index:** $i \leftarrow (j - 1) \cdot k + m$
        * **Project Query:** $Q_i \leftarrow X W_Q^{(i)}$
        * **Compute Scores:** $S_i \leftarrow \frac{Q_i K_j^T}{\sqrt{d_k}}$ *(Note: $K_j$ is shared within group)*
        * **Apply Mask (If Decoder):** $S_i \leftarrow S_i + M$
        * **Compute Probabilities:** $P_i \leftarrow \text{Softmax}(S_i)$
        * **Compute Head Output:** $\text{head}_i \leftarrow P_i V_j$ *(Note: $V_j$ is shared within group)*
        * **Append to list:** $H_{\text{list}}[i] \leftarrow \text{head}_i$
    * **End For**
5.  **End For**
6.  **Concatenate Heads:** $H_{\text{concat}} \leftarrow [H_{\text{list}}[1] \,\|\, H_{\text{list}}[2] \,\|\, \dots \,\|\, H_{\text{list}}[h]]$ 
7.  **Final Projection:** $\text{Output} \leftarrow H_{\text{concat}} W_O$
8.  **Return** $\text{Output} \in \mathbb{R}^{T \times d_{model}}$

3.4 Multi-Head Latent Attention(MLA)

While Grouped-Query Attention (GQA) struck a strong balance between the speed of MQA and the quality of MHA, it still suffers from a fundamental scaling law: the KV Cache grows linearly with the sequence length. When modern LLMs push context windows to 128K, 256K, or even 1M tokens, storing even a “grouped” KV cache requires hundreds of gigabytes of VRAM per batch. GQA mitigates the bottleneck, but it does not eliminate it.

The Motivation: To achieve truly massive context lengths without running out of GPU memory, we cannot just share Keys and Values; we must fundamentally compress them. Multi-Head Latent Attention (MLA) achieves this by projecting the Key-Value states into a low-dimensional latent space, only decompressing them when necessary.

Latent Compression and RoPE Decoupling

Instead of computing and caching full, high-dimensional Key and Value matrices for every token, MLA computes a single compressed latent vector ($c^{KV}$) for each token.

However, Attention mechanisms rely heavily on Rotary Position Embeddings (RoPE) to understand the relative distances between tokens. RoPE operations are highly sensitive to matrix dimensions and cannot easily survive being squashed into a tiny latent vector. To solve this, MLA uses RoPE Decoupling: it splits the Queries and Keys into a compressed semantic part (which goes through the latent bottleneck) and a separate, uncompressed positional part (which carries the RoPE information).

Let our input sequence be $X \in \mathbb{R}^{T \times d_{model}}$. Let $h$ be the number of heads. We define new dimensionalities:

• $d_c$ = The dimension of the KV latent space ($d_c \ll d_{model}$).
• $d_c’$ = The dimension of the Query latent space.
• $d_r$ = The decoupled RoPE dimension.

Step 1: Down-Projection (Compression) First, we compress the input $X$ into our latent spaces using down-projection matrices:

Crucially, during inference, $C^{KV}$ is the only semantic tensor we need to store in the KV cache!

Step 2: Up-Projection (Decompression) for Semantic Info For each head $i \in {1, \dots, h}$, we up-project the latent vectors back into the multi-head dimensions ($d_k$ and $d_v$):

Step 3: RoPE Decoupling for Positional Info In parallel, we compute the positional Queries and Keys directly from $X$, bypassing the latent bottleneck, and apply the RoPE function ($\mathcal{R}$):

(Notice that $K^r$ can be shared across all heads, acting like an MQA mechanism specifically for positional data).

Step 4: Recombination and Attention Computation We concatenate the semantic and positional components to form our final Queries and Keys:

Finally, we compute standard Self-Attention for each head:

Step 5: Final Projection Just like previous architectures, we concatenate the output of all heads and project back:

Why is this a breakthrough? (The Absorption Trick)

Looking at the math above, you might think: “Wait, if we up-project $C^{KV}$ back into $K_i^c$ and $V_i$ during decoding, aren’t we still doing the same amount of computation?”

Here is the mathematical magic of MLA during inference. By leveraging the associative property of matrix multiplication, we can absorb the decompression weights into the Query and Output projection weights.

Look at the semantic dot product: $Q_i^c (K_i^c)^T$

Because $W_{UQ}^{(i)}$ and $W_{UK}^{(i)}$ are just fixed learnable weight matrices, we can pre-multiply them offline into a single absorbed weight matrix $W_{absorbed}^{(i)}$. Result: We can compute the attention scores directly against the compressed $C^{KV}$ vector! We literally never instantiate the massive $K_i$ or $V_i$ matrices in memory during generation.

This pseudo-code illustrates the forward pass during **training** (where full materialization happens for parallel processing).

**Input:**
* $X \in \mathbb{R}^{T \times d_{model}}$ 
* Optional: $M \in \mathbb{R}^{T \times T}$ (Mask matrix)

**Algorithm:**
1.  **Compress to Latent Space:**
    * $C^{KV} \leftarrow X W^{DKV}$
    * $C^Q \leftarrow X W^{DQ}$
2.  **Compute Decoupled RoPE Keys:**
    * $K^r \leftarrow \mathcal{R}(X W_{KR})$
3.  **Initialize empty list:** $H_{\text{list}} = []$
4.  **For** $i = 1$ **to** $h$ **do in parallel:**
    * **Decompress Semantic Query:** $Q_i^c \leftarrow C^Q W_{UQ}^{(i)}$
    * **Decompress Semantic Key:** $K_i^c \leftarrow C^{KV} W_{UK}^{(i)}$
    * **Decompress Value:** $V_i \leftarrow C^{KV} W_{UV}^{(i)}$
    * **Compute Positional Query:** $Q_i^r \leftarrow \mathcal{R}(X W_{QR}^{(i)})$
    * **Concatenate:** $Q_i \leftarrow [Q_i^c \parallel Q_i^r]$
    * **Concatenate:** $K_i \leftarrow [K_i^c \parallel K^r]$
    * **Compute Scores:** $S_i \leftarrow \frac{Q_i K_i^T}{\sqrt{d_k + d_r}}$
    * **Apply Mask:** $S_i \leftarrow S_i + M$
    * **Compute Probabilities:** $P_i \leftarrow \text{Softmax}(S_i)$
    * **Compute Head Output:** $\text{head}_i \leftarrow P_i V_i$
    * **Append to list:** $H_{\text{list}}[i] \leftarrow \text{head}_i$
5.  **End For**
6.  **Concatenate Heads:** $H_{\text{concat}} \leftarrow [H_{\text{list}}[1] \,\|\, \dots \,\|\, H_{\text{list}}[h]]$ 
7.  **Final Projection:** $\text{Output} \leftarrow H_{\text{concat}} W_O$
8.  **Return** $\text{Output} \in \mathbb{R}^{T \times d_{model}}$

3.5 Sliding Window Attention(SWA)

While techniques like MQA, GQA, and MLA successfully optimize the memory required for the KV Cache during decoding, they do not solve the original mathematical bottleneck of the Self-Attention equation: time and compute complexity.

The Motivation: The $O(T^2)$ Bottleneck In standard Self-Attention, every token computes a dot product with every other token in the sequence to form the $T \times T$ attention matrix.

• Compute complexity: $O(T^2 \cdot d)$
• Memory complexity: $O(T^2)$

For short sequences, this is fine. But if we want a model to process an entire book (e.g., $T = 100,000$ tokens), $T^2$ becomes $10,000,000,000$ operations per head, per layer. It absolutely shatters both compute limits and VRAM.

The Solution: Sliding Window Attention (popularized by models like Longformer and Mistral) recognizes that language is highly local. A word at the end of a book rarely needs a direct, unmediated mathematical connection to a word on page one. Instead of computing the full global attention matrix, SWA restricts each token to only attend to a fixed-size local “window” of adjacent tokens.

Local Windows and Receptive Fields

By defining a window size $w$, we restrict the attention span. For an auto-regressive decoder, instead of token $i$ attending to all previous tokens from $1$ to $i$, it only attends to tokens from $i-w$ to $i$.

• The complexity drops from $O(T^2)$ to $O(T \times w)$. Since $w$ is a fixed constant (e.g., $w = 4096$), the complexity becomes linear with respect to sequence length!

But how does the model understand the global context? It does this through the Receptive Field across multiple layers. If layer 1 allows a token to see $w$ tokens back, and layer 2 allows the token to see $w$ tokens back from those previous tokens’ representations, the effective receptive field grows with each layer.

For a model with $L$ layers, a token at the final layer has an effective receptive field of $L \times w$. Information flows forward through the network dynamically, much like a stacked Convolutional Neural Network (CNN).

Let our input sequence be $X \in \mathbb{R}^{T \times d_{model}}$. Let $w \in \mathbb{Z}^+$ be our fixed window size.

The subspace projections for Queries, Keys, and Values remain identical to standard attention:

The Sliding Window Mask ($M_{\text{SWA}}$)

The entire magic of SWA happens in the modification of the Mask matrix. In standard auto-regressive attention, the mask $M$ blocks the future:

In Sliding Window Attention, we modify this mask to block both the future and the distant past. We define the new banded mask matrix $M_{\text{SWA}} \in \mathbb{R}^{T \times T}$ as:

When we apply this mask to our scaled scores:

The Softmax function drives the probabilities of all tokens outside the window $w$ to exactly $0$.

Sparse Matrix Computation

Mathematically, writing out the full $T \times T$ matrix just to mask it with $-\infty$ defeats the purpose of saving compute. In practice (using custom CUDA kernels like FlashAttention), we do not instantiate the full matrix. We only compute the dot products for the valid band of width $w$:

This pseudo-code demonstrates the constrained inner loop, ensuring that the algorithm only computes similarity scores within the localized window, achieving linear $O(T)$ complexity.

**Input:**
* $X \in \mathbb{R}^{T \times d_{model}}$ (Input sequence embeddings)
* $w \in \mathbb{Z}^+$ (Sliding window size)

**Learnable Parameters:**
* $W_Q, W_K \in \mathbb{R}^{d_{model} \times d_k}$ 
* $W_V \in \mathbb{R}^{d_{model} \times d_v}$ 
* $W_O \in \mathbb{R}^{d_{model} \times d_{model}}$ 

**Algorithm:**
1.  **Project Subspaces:**
    * $Q \leftarrow X W_Q$
    * $K \leftarrow X W_K$
    * $V \leftarrow X W_V$
2.  **Initialize empty output list:** $H_{\text{out}} = []$
3.  **For** $i = 1$ **to** $T$ **do:** *(Iterate over each token)*
    * **Determine Window Bounds:** * $start \leftarrow \max(1, i - w)$
        * $end \leftarrow i$
    * **Slice Local Keys and Values:**
        * $K_{local} \leftarrow K[start : end]$ *(Shape: $\leq w \times d_k$)*
        * $V_{local} \leftarrow V[start : end]$ *(Shape: $\leq w \times d_v$)*
    * **Extract current Query:** $q_i \leftarrow Q[i]$ *(Shape: $1 \times d_k$)*
    * **Compute Local Scores:** $S_i \leftarrow \frac{q_i K_{local}^T}{\sqrt{d_k}}$ 
    * **Compute Probabilities:** $P_i \leftarrow \text{Softmax}(S_i)$
    * **Compute Output Vector:** $o_i \leftarrow P_i V_{local}$
    * **Append to list:** $H_{\text{out}}[i] \leftarrow o_i$
4.  **End For**
5.  **Stack Outputs:** $H_{stack} \leftarrow \text{Stack}(H_{\text{out}})$ *(Shape: $T \times d_v$)*
6.  **Final Projection:** $\text{Output} \leftarrow H_{stack} W_O$
7.  **Return** $\text{Output} \in \mathbb{R}^{T \times d_{model}}$

3.6 DeepSeek Sparse Attention(DSA)

While Sliding Window Attention (SWA) successfully reduces the computational complexity of attention from $O(T^2)$ to $O(T)$, it introduces a severe architectural blind spot. Because SWA strictly enforces a local boundary, a token on page 100 of a document cannot directly attend to a crucial instruction on page 1. The model relies entirely on information trickling through the hidden states layer-by-layer, which often leads to poor performance on complex “needle-in-a-haystack” retrieval tasks.

The Motivation: We need a mechanism that is as computationally cheap as SWA ($O(T \times k)$), but capable of attending to the most important tokens regardless of where they appear in the sequence.

The Solution: DeepSeek Sparse Attention (DSA) replaces the rigid, fixed-distance window with a dynamic, two-stage indexer. Instead of assuming that “recent tokens are the most important,” DSA employs a tiny, lightweight “Lightning Indexer” to rapidly scan the entire sequence and dynamically select the top $k$ most relevant tokens for each query. Full attention is then computed only on that sparsely selected subset.

1. The Core Concept: The Two-Stage Pipeline

DSA breaks the attention calculation into two phases:

1. The Lightning Indexer: A highly compressed, low-dimension attention mechanism. It calculates rough similarity scores between a query and all past keys using a fast activation function (like ReLU) instead of an expensive Softmax.
2. Fine-Grained Token Selection (Top-$k$): Based on the Lightning Indexer’s scores, the model identifies the indices of the $k$ most important tokens. It then fetches the full-precision, high-dimensional Keys and Values only for those $k$ tokens to compute the final Attention output.

By doing this, DSA maintains performance parity with dense $O(T^2)$ attention in long-context scenarios, while massively reducing per-token GPU costs.

Let our input sequence be $X \in \mathbb{R}^{T \times d_{model}}$. Let $k \in \mathbb{Z}^+$ be our sparsity budget (the maximum number of tokens a query is allowed to attend to). Let $d_{light} \ll d_k$ be the dramatically reduced dimension of our Lightning Indexer.

Step 1: The Lightning Indexer Projections

We project the input into a highly compressed query and key space for the indexer:

For a given token $i$, we compute the raw indexer scores against all past tokens $j \le i$. Instead of scaling and applying Softmax, DSA uses a fast, non-linear activation (often ReLU) to zero out weak connections instantly:

Step 2: Dynamic Top-$k$ Masking

We sort the scores in $S_{light}^{(i)}$ to find the indices of the $k$ largest values. Let $\mathcal{J}_i$ be the set of selected token indices for query $i$, where $|\mathcal{J}_i| = k$.

We define our dynamic sparse mask $M_{\text{DSA}} \in \mathbb{R}^{T \times T}$ as:

Step 3: Full-Precision Sparse Attention

Finally, we compute the standard attention using the full-dimensional $Q$, $K$, and $V$ matrices, but masked by our dynamic selections:

Note: Mathematically, applying the mask $M_{\text{DSA}}$ makes the probabilities of all non-selected tokens $0$. In hardware implementation, the model does not actually compute the full $QK^T$ matrix; it uses gather operations to only compute dot products for the indices in $\mathcal{J}_i$, making the operation strictly $O(T \cdot k)$.

This algorithm highlights the two-tiered approach: a cheap, global scan followed by an expensive, localized computation.

**Input:**
* $X \in \mathbb{R}^{T \times d_{model}}$ (Input sequence embeddings)
* $k \in \mathbb{Z}^+$ (Number of sparse tokens to select)

**Learnable Parameters:**
* $W_{Q, light}, W_{K, light} \in \mathbb{R}^{d_{model} \times d_{light}}$ (Indexer weights)
* $W_Q, W_K \in \mathbb{R}^{d_{model} \times d_k}$ (Full precision Q, K weights)
* $W_V \in \mathbb{R}^{d_{model} \times d_v}$ (Full precision V weights)
* $W_O \in \mathbb{R}^{d_{model} \times d_{model}}$ (Output linear projection weight)

**Algorithm:**
1.  **Project Indexer Subspaces:**
    * $Q_{light} \leftarrow X W_{Q, light}$
    * $K_{light} \leftarrow X W_{K, light}$
2.  **Project Full-Precision Subspaces:**
    * $Q \leftarrow X W_Q$
    * $K \leftarrow X W_K$
    * $V \leftarrow X W_V$
3.  **Initialize empty output list:** $H_{\text{out}} = []$
4.  **For** $i = 1$ **to** $T$ **do:** *(Iterate over each token)*
    * **Extract current Indexer Query:** $q_{i, light} \leftarrow Q_{light}[i]$ 
    * **Slice Valid Past Indexer Keys:** $K_{past, light} \leftarrow K_{light}[1 : i]$
    * **Compute Indexer Scores:** $S_{indexer} \leftarrow \text{ReLU}(q_{i, light} K_{past, light}^T)$
    * **Select Top-k Indices:** $\mathcal{J}_i \leftarrow \text{ArgTopK}(S_{indexer}, k)$
    * **Gather Sparse Full-Precision Keys/Values:**
        * $K_{sparse} \leftarrow K[\mathcal{J}_i]$ *(Shape: $k \times d_k$)*
        * $V_{sparse} \leftarrow V[\mathcal{J}_i]$ *(Shape: $k \times d_v$)*
    * **Extract current Full Query:** $q_i \leftarrow Q[i]$ *(Shape: $1 \times d_k$)*
    * **Compute Sparse Scores:** $S_i \leftarrow \frac{q_i K_{sparse}^T}{\sqrt{d_k}}$ 
    * **Compute Probabilities:** $P_i \leftarrow \text{Softmax}(S_i)$
    * **Compute Output Vector:** $o_i \leftarrow P_i V_{sparse}$
    * **Append to list:** $H_{\text{out}}[i] \leftarrow o_i$
5.  **End For**
6.  **Stack Outputs:** $H_{stack} \leftarrow \text{Stack}(H_{\text{out}})$ *(Shape: $T \times d_v$)*
7.  **Final Projection:** $\text{Output} \leftarrow H_{stack} W_O$
8.  **Return** $\text{Output} \in \mathbb{R}^{T \times d_{model}}$

3.7 Gated Attention

While techniques like Sliding Window and Sparse Attention solve the computational bottlenecks of long sequences, they don’t address a critical vulnerability in the core Attention equation itself: Softmax Information Leakage and Scalar Weighting.

The Motivation: Standard Self-Attention relies entirely on the Softmax function to filter information. This creates two distinct problems:

1. Dense Leakage: Softmax rarely outputs absolute zero. Even highly irrelevant tokens get a tiny fractional weight (e.g., $0.001$). Over thousands of tokens, this “tail noise” accumulates, degrading the exactness of the context.
2. Scalar Bottleneck: Attention assigns a single scalar value ($\alpha$) to an entire token’s Value vector ($v_i$). What if token $A$ is syntactically relevant to the query but semantically irrelevant? Standard attention scales the entire vector equally, unable to filter out specific noisy dimensions within the token.

The Solution: Gated Attention introduces a dynamic, parameterized Gate parallel to the attention calculation. By combining Attention with the logic of Gated Linear Units (GLUs) or LSTM-style gates, the model can apply dimension-wise filtering, squelching irrelevant features after the tokens have interacted.

1. The Core Concept: The Element-wise Gate

In Gated Attention, the input sequence is projected into an additional subspace: the Gate ($G$).

The Gate uses a non-linear activation function bounded between 0 and 1 (like Sigmoid) or a smooth approximation (like SiLU/Swish). After the standard Attention mechanism computes the weighted sum of the Values, the Gate is applied element-wise to the output.

This means that even if the Softmax attention pulls in noisy information from a distant token, the Gate can act as a final firewall, multiplying the noisy dimensions by $0$ and letting the useful dimensions pass through by multiplying them by $1$.

1. Mathematical Formulation

Let our input sequence be $X \in \mathbb{R}^{T \times d_{model}}$.

We compute our standard Queries, Keys, and Values:

Step 1: The Gate Projection In parallel, we project the input $X$ into a Gating subspace using a separate weight matrix $W_G$. The dimension of the gate must perfectly match the dimension of the Value vectors ($d_v$) so they can be multiplied element-wise. We apply an activation function $\sigma$ (typically Sigmoid or SiLU).

Step 2: Standard Attention Computation We compute the standard scaled dot-product attention to aggregate the sequence context. Let’s call this intermediate aggregated matrix $A$:

Step 3: The Element-wise Gating (Hadamard Product) We now apply the Gate $G$ to the aggregated attention context $A$ using the Hadamard product ($\odot$), which represents element-wise multiplication.

For a specific token representation $i$ and feature dimension $m$, the calculation is:

Step 4: Final Projection As usual, the gated output is projected back into the model dimension:

1. Modern Evolution: The Gated Attention Unit (GAU)

In 2022, researchers pushed this concept to its absolute limit with the Gated Attention Unit (GAU) (Hua et al.). They realized that if the gating mechanism is strong enough, the Attention mechanism itself can be drastically simplified.

In a GAU, the heavy Multi-Head Attention blocks and the Feed-Forward Network (FFN) blocks are collapsed into a single, highly efficient layer.

The math for a GAU simplifies to:

1. Project $U$ and $V$: Compute an un-gated Value matrix ($V$) and a Gate matrix ($U$) using the SiLU activation.

   \[U = \text{SiLU}(X W_U)\] \[V = \text{SiLU}(X W_V)\]

2. Cheap Attention ($Z$): Compute a very low-dimensional, single-head attention using simplified $Q$ and $K$ projections, and apply it to $V$.

   \[Z = \text{Softmax}\left( \frac{QK^T}{\sqrt{d_k}} \right) V\]

3. Gate the Output:

   \[O = (U \odot Z) W_O\]

Because the element-wise gate ($U$) handles so much of the representational heavy lifting, the attention calculation becomes much cheaper to run, often yielding higher quality in linear or near-linear time.

This pseudo-code demonstrates the standard Gated Attention mechanism, highlighting the parallel computation of the Gate and the final element-wise filtering step.

**Input:**
* $X \in \mathbb{R}^{T \times d_{model}}$ (Input sequence embeddings)
* Optional: $M \in \mathbb{R}^{T \times T}$ (Mask matrix)

**Learnable Parameters:**
* $W_Q, W_K \in \mathbb{R}^{d_{model} \times d_k}$ 
* $W_V \in \mathbb{R}^{d_{model} \times d_v}$ 
* $W_G \in \mathbb{R}^{d_{model} \times d_v}$ (Gate projection weight)
* $W_O \in \mathbb{R}^{d_v \times d_{model}}$ (Output projection weight)

**Algorithm:**
1.  **Project Standard Subspaces:**
    * $Q \leftarrow X W_Q$
    * $K \leftarrow X W_K$
    * $V \leftarrow X W_V$
2.  **Project and Activate the Gate:**
    * $G_{raw} \leftarrow X W_G$
    * $G \leftarrow \text{SiLU}(G_{raw})$ *(or Sigmoid)*
3.  **Compute Attention Scores:**
    * $S \leftarrow \frac{Q K^T}{\sqrt{d_k}}$ 
    * **If Masked:** $S \leftarrow S + M$
4.  **Compute Aggregated Context:**
    * $P \leftarrow \text{Softmax}(S)$
    * $A \leftarrow P \cdot V$ *(Matrix multiplication)*
5.  **Apply Dimension-wise Gating:**
    * $O_{gated} \leftarrow G \odot A$ *(Element-wise multiplication)*
6.  **Final Linear Projection:**
    * $\text{Output} \leftarrow O_{gated} W_O$
7.  **Return** $\text{Output} \in \mathbb{R}^{T \times d_{model}}$

3.8 Linear Attention

While Sliding Window, Sparse, and Hybrid attentions reduce the compute cost of processing long sequences, they all still fundamentally compute the standard Softmax Attention equation.

The Motivation: The true root of the $O(T^2)$ complexity bottleneck is the Softmax function itself. Look at the standard equation: $A = \text{Softmax}(QK^T)V$. Because Softmax is a non-linear operation applied over the rows of the $QK^T$ matrix, we are mathematically forced to compute the entire $T \times T$ matrix first before we can multiply by $V$.

The Solution: What if we could remove the Softmax function? If we remove the non-linearity, we can use the associative property of matrix multiplication. Instead of computing $(Q K^T) V$, which costs $O(T^2 \cdot d)$, we could compute $Q (K^T V)$, which costs $O(T \cdot d^2)$. Since the sequence length $T$ is usually much larger than the embedding dimension $d$, this drops the complexity to strictly linear $O(T)$.

1. The Core Concept: The Kernel Trick

To remove the Softmax while still maintaining a valid probability-like weighting (where scores are positive and sum to 1), Linear Attention (Katharopoulos et al., 2020) uses a generalized kernel feature map.

In standard attention, the output of the $i$-th token is the weighted average of all values $v_j$, where the similarity is defined by an exponential function:

Linear Attention replaces the $\exp(\cdot)$ similarity with a decomposable kernel feature map $\phi(\cdot)$. Instead of $\text{sim}(q_i, k_j) = \exp(q_i k_j^T)$, we define:

Where $\phi(\cdot)$ is a function that applies row-wise to ensure all vectors have strictly positive values (commonly elu(x) + 1 or ReLU(x)).

1. Mathematical Formulation: The Associativity Magic

Let our projected matrices be $Q, K \in \mathbb{R}^{T \times d_k}$ and $V \in \mathbb{R}^{T \times d_v}$.

Step 1: Substitute the Feature Map We substitute our new similarity function into the attention equation for the $i$-th query:

Step 2: Factor out the Query Because $\phi(q_i)$ depends only on $i$ and not on the summation index $j$, we can pull it outside the sum! This is the mathematical breakthrough that Softmax prevented.

Step 3: Compute the Global Context (The $K^T V$ part) Look closely at the summation terms. Neither of them depends on the query $Q$. They only depend on $K$ and $V$. We can define two new global context matrices:

1. The Numerator Matrix ($S$): $S = \sum_{j=1}^{T} \phi(k_j)^T v_j \implies S \in \mathbb{R}^{d_k \times d_v}$
2. The Denominator Vector ($Z$): $Z = \sum_{j=1}^{T} \phi(k_j)^T \implies Z \in \mathbb{R}^{d_k}$

(Note: $\phi(k_j)^T$ is a column vector of size $d_k \times 1$, and $v_j$ is a row vector of size $1 \times d_v$. Their outer product creates a $d_k \times d_v$ matrix).

Step 4: Final Linear Computation Now, to find the output for any query, we just do a simple matrix multiplication with our pre-computed $S$ and $Z$:

Complexity Result: Computing $S$ takes $O(T \cdot d_k \cdot d_v)$. Multiplying by $Q$ takes $O(T \cdot d_k \cdot d_v)$. The dreaded $T^2$ has completely vanished.

1. The RNN Connection (Auto-Regressive Inference)

The most powerful feature of Linear Attention is what happens during causal/auto-regressive decoding. In standard attention, generating token $T+1$ requires keeping the entire KV cache of all previous tokens to compute the $T \times T$ matrix.

In Linear Attention, the summations for $S$ and $Z$ can be expressed as a running cumulative sum (just like the hidden state of an RNN). For step $i$:

Result: The memory footprint for the KV cache becomes $O(1)$ with respect to sequence length. The model only needs to store a fixed-size $d_k \times d_v$ matrix ($S$) and a $d_k$ vector ($Z$), regardless of whether it has generated 10 tokens or 100,000 tokens!

This algorithm highlights the auto-regressive (causal) implementation of Linear Attention, demonstrating how it operates exactly like an RNN using a running state.

**Input:**
* $X \in \mathbb{R}^{T \times d_{model}}$ (Input sequence embeddings)

**Learnable Parameters:**
* $W_Q, W_K \in \mathbb{R}^{d_{model} \times d_k}$ 
* $W_V \in \mathbb{R}^{d_{model} \times d_v}$ 
* $W_O \in \mathbb{R}^{d_{model} \times d_{model}}$

**Algorithm:**
1.  **Project Subspaces:**
    * $Q \leftarrow X W_Q$
    * $K \leftarrow X W_K$
    * $V \leftarrow X W_V$
2.  **Apply Feature Map ($\phi$):**
    * $Q_{\phi} \leftarrow \text{ELU}(Q) + 1$ *(Ensures strictly positive values)*
    * $K_{\phi} \leftarrow \text{ELU}(K) + 1$
3.  **Initialize Running States:**
    * $S_0 \leftarrow \text{Zeros}(d_k, d_v)$ *(The cumulative Numerator state)*
    * $Z_0 \leftarrow \text{Zeros}(d_k, 1)$ *(The cumulative Denominator state)*
4.  **Initialize empty output list:** $H_{\text{out}} = []$
5.  **For** $i = 1$ **to** $T$ **do:** *(Iterate causally)*
    * **Extract current vectors:**
        * $q_i \leftarrow Q_{\phi}[i]$ *(Shape: $1 \times d_k$)*
        * $k_i \leftarrow K_{\phi}[i]$ *(Shape: $1 \times d_k$)*
        * $v_i \leftarrow V[i]$ *(Shape: $1 \times d_v$)*
    * **Update Running States (The RNN step):**
        * $S_i \leftarrow S_{i-1} + (k_i^T \cdot v_i)$ *(Outer product: adds to $d_k \times d_v$ matrix)*
        * $Z_i \leftarrow Z_{i-1} + k_i^T$ *(Vector addition)*
    * **Compute Output for current token:**
        * $\text{Num}_i \leftarrow q_i \cdot S_i$ *(Vector-Matrix mult: $1 \times d_v$)*
        * $\text{Den}_i \leftarrow q_i \cdot Z_i$ *(Dot product: scalar)*
        * $o_i \leftarrow \frac{\text{Num}_i}{\text{Den}_i}$ *(Element-wise division)*
    * **Append to list:** $H_{\text{out}}[i] \leftarrow o_i$
6.  **End For**
7.  **Stack Outputs:** $H_{stack} \leftarrow \text{Stack}(H_{\text{out}})$ 
8.  **Final Projection:** $\text{Output} \leftarrow H_{stack} W_O$
9.  **Return** $\text{Output} \in \mathbb{R}^{T \times d_{model}}$

3.9 Hybrid Attention

While Sliding Window Attention (SWA) limits computation by restricting tokens to a local neighborhood, and DeepSeek Sparse Attention (DSA) attempts to fix this dynamically by searching for top-$k$ tokens, both have limitations. SWA loses global context entirely. DSA requires dynamic routing, which introduces memory-gathering overheads and makes hardware optimization (like custom CUDA kernels) extremely difficult.

The Motivation: Is there a way to achieve the $O(T)$ linear complexity of local attention while maintaining a deterministic, guaranteed pathway for global information to flow across the entire sequence?

The Solution: Hybrid Attention (pioneered by models like Longformer, BigBird, and Sparse Transformers) fundamentally combines two distinct attention patterns: Local Windowed Attention and Global Attention. It does this deterministically by assigning special “global” status to a small subset of tokens, allowing them to bypass the local window restrictions.

1. The Core Concept: Deterministic Sparsity

Hybrid Attention achieves its balance through two complementary mechanisms operating within the same layer:

1. Local Sliding Window: The vast majority of tokens in the sequence are treated as “local.” They can only attend to their immediate neighbors within a fixed window $w$. This captures local grammar, syntax, and immediate context.
2. Global Tokens: A highly restricted subset of tokens (e.g., the first prompt token, special [CLS] tokens, or periodic tokens spaced every $k$ steps) are designated as “global.”
   • Global-to-All: A global token can attend to every token in the sequence.
   • All-to-Global: Every token in the sequence can attend to the global tokens.

This creates a “hub-and-spoke” information flow. If a token on page 100 needs information from page 1, the information flows from page 1 into a Global Token, and then from the Global Token directly to the token on page 100. The maximum path length between any two tokens is reduced to just 2 steps, solving the isolation problem of SWA while remaining computationally cheap.

Let our input sequence be $X \in \mathbb{R}^{T \times d_{model}}$. Let $w$ be the sliding window size. Let $\mathcal{G} \subset {1, 2, \dots, T}$ be the set of indices designated as global tokens. The size of this set is very small ($|\mathcal{G}| = c \ll T$).

We compute standard subspace projections:

Step 1: The Hybrid Auto-Regressive Mask

The defining feature of Hybrid Attention is its mathematically unified Mask matrix $M_{hybrid} \in \mathbb{R}^{T \times T}$.

For an auto-regressive decoder (where future tokens are strictly blocked), the mask combines the local band and the global rows/columns. We define the mask entries $m_{ij}$ as:

Step 2: Dual-Projection (Advanced Capacity)

In standard attention, projecting global and local interactions into the same dense subspace can cause feature collapse (the keys get confused trying to represent both local syntax and global summary).

To solve this, advanced Hybrid Attention uses Dual-Projections. We define a second set of weights specifically for global interactions:

If a token is interacting locally, we use the standard $W_K$. If it is interacting globally, we use $W_K^{global}$. The final attention scores are computed by summing the unmasked portions of both the local $QK^T$ matrix and the global $Q_{global} K_{global}^T$ matrix.

For simplicity in our core formulation below, we will assume a single shared subspace, which is often sufficient for interleaved layer approaches.

Step 3: Layer-Wise Hybridization (Alternative)

Instead of putting a complex hybrid mask inside a single layer, modern architectures (like GPT-3’s sparse variants or Mistral) often achieve Hybrid Attention by interleaving layers.

• Layer 1: Pure Sliding Window Attention (SWA)
• Layer 2: Pure Global Dense Attention
• Layer 3: Pure Sliding Window Attention (SWA)

Mathematically, this forces the Dense layers to act as the “Global Hubs” while keeping the overall model FLOPs massively reduced.

This algorithm represents token-wise Hybrid Attention (hub-and-spoke). Note how the context for any given token is formed by taking the union of the global keys and the local sliding window keys.

**Input:**
* $X \in \mathbb{R}^{T \times d_{model}}$ (Input sequence embeddings)
* $w \in \mathbb{Z}^+$ (Sliding window size)
* $\mathcal{G}$ (Set of global token indices)

**Learnable Parameters:**
* $W_Q, W_K \in \mathbb{R}^{d_{model} \times d_k}$ 
* $W_V \in \mathbb{R}^{d_{model} \times d_v}$ 
* $W_O \in \mathbb{R}^{d_{model} \times d_{model}}$

**Algorithm:**
1.  **Project Subspaces:**
    * $Q \leftarrow X W_Q$
    * $K \leftarrow X W_K$
    * $V \leftarrow X W_V$
2.  **Isolate Global States (Pre-compute):**
    * $K_{\mathcal{G}} \leftarrow K[\mathcal{G}]$ *(Shape: $c \times d_k$)*
    * $V_{\mathcal{G}} \leftarrow V[\mathcal{G}]$ *(Shape: $c \times d_v$)*
3.  **Initialize empty output list:** $H_{\text{out}} = []$
4.  **For** $i = 1$ **to** $T$ **do:** *(Iterate over each token)*
    * **Extract current Query:** $q_i \leftarrow Q[i]$ 
    * **Determine Local Window Bounds:** * $start \leftarrow \max(1, i - w)$
        * $end \leftarrow i$
    * **Slice Valid Past Local Keys/Values:**
        * $K_{local} \leftarrow K[start : end]$
        * $V_{local} \leftarrow V[start : end]$
    * **Form Hybrid Context (Union):**
        * *If $i \in \mathcal{G}$ (Query is Global):*
            * Context is ALL past tokens: $K_{ctx} \leftarrow K[1 : i]$, $V_{ctx} \leftarrow V[1 : i]$
        * *Else (Query is Local):*
            * Filter $K_{\mathcal{G}}$ to only include past indices ($j \le i$)
            * Remove any overlap between local window and global set to prevent double-counting.
            * $K_{ctx} \leftarrow [K_{\mathcal{G}, past} \parallel K_{local}]$
            * $V_{ctx} \leftarrow [V_{\mathcal{G}, past} \parallel V_{local}]$
    * **Compute Scores:** $S_i \leftarrow \frac{q_i K_{ctx}^T}{\sqrt{d_k}}$ 
    * **Compute Probabilities:** $P_i \leftarrow \text{Softmax}(S_i)$
    * **Compute Output Vector:** $o_i \leftarrow P_i V_{ctx}$
    * **Append to list:** $H_{\text{out}}[i] \leftarrow o_i$
5.  **End For**
6.  **Stack Outputs:** $H_{stack} \leftarrow \text{Stack}(H_{\text{out}})$ 
7.  **Final Projection:** $\text{Output} \leftarrow H_{stack} W_O$
8.  **Return** $\text{Output} \in \mathbb{R}^{T \times d_{model}}$

4. Current Research Directions

While the fundamental scaled dot-product attention has driven the AI revolution over the past decade, the relentless push for infinite context windows and edge-device deployment has exposed its ultimate limits. The research frontier has now shifted from basic structural tweaks to fundamental architectural reorganizations.

Here are the primary directions shaping the future of Attention Mechanisms:

1. Sub-Quadratic Hybrids: Attention Meets State Space Models (SSMs)

For years, Linear Attention variants attempted to replace the Softmax bottleneck, but often suffered from performance degradation on complex reasoning tasks. The current state-of-the-art approach does not eliminate Attention; it hybridizes it.

Researchers are actively combining local/sparse Attention layers with State Space Models (SSMs) like Mamba.

• The Paradigm: SSMs are mathematically equivalent to continuous-time convolutions and RNNs, offering true $O(1)$ inference memory and $O(T)$ training scaling. However, they occasionally struggle with exact “needle-in-a-haystack” retrieval.
• The Direction: Modern architectures interleave layers (e.g., 3 layers of Mamba for every 1 layer of Dense Attention). The SSM handles the smooth, linear accumulation of global context, while the Dense Attention layer acts as an exact retrieval mechanism.

2. Dynamic KV-Cache Eviction (Selective Attention)

As models scale to context windows of $1,000,000+$ tokens, storing the Keys and Values for every single token—even using Multi-Head Latent Attention (MLA)—becomes unsustainable.

• The Paradigm: Not all tokens are created equal. Words like “the,” “a,” or filler sentences stop contributing meaningful semantic weight after they have been processed by the early layers.
• The Direction: “Selective Attention” or “Preserve-then-Select” mechanisms. These algorithms introduce parameter-free pruning algorithms that dynamically evaluate the attention scores during the forward pass. If a token’s cumulative attention weight falls below a critical threshold $\epsilon$, it is completely evicted from the KV cache. This dynamically shrinks the matrix sizes during auto-regressive generation, keeping memory footprints strictly bounded.

3. Hardware-Algorithm Co-Design (Extreme Block Sparsity)

The math of Attention is no longer developed in a vacuum; it is designed explicitly for the memory hierarchy of GPUs and custom TPUs (SRAM vs. HBM).

• The Paradigm: Moving data from High Bandwidth Memory (HBM) to the compute cores (SRAM) is infinitely slower than doing the math itself.
• The Direction: Expanding on kernels like FlashAttention, modern research focuses on Training-aware Block Sparse Attention. Instead of masking out individual tokens mathematically (which hardware hates), the attention matrix is partitioned into massive $N \times N$ blocks. Neural routing layers determine which entire blocks of the matrix can be dropped before they are ever loaded into SRAM. This allows for massive context parallelization across clusters.

4. Energy-Based Routing and Non-Probabilistic Mixing

The core assumption of standard Attention is that information must be mixed based on a probability distribution (the Softmax output summing to 1).

• The Paradigm: The Softmax constraint artificially forces competition between tokens. If a query finds ten highly relevant tokens, it must dilute its attention across all ten, assigning them a score of $0.1$.
• The Direction: Moving toward Energy-based Attention models (e.g., GD-Attention). These mechanisms ditch the strict probability constraint and instead evaluate the “Semantic Energy” between tokens. If ten tokens are important, they can all pass through the routing gate with full magnitude. This shifts the mechanism from “Weighted Mixing” to strict “Selection and Protection.”

5. Conclusion

The evolution of the Attention Mechanism is the defining mathematical story of modern artificial intelligence.

Prior to Attention, sequence modeling was fundamentally trapped by the Information Bottleneck of recurrent networks. Every piece of input data, no matter the length, had to be ruthlessly compressed into a single, fixed-length hidden vector. The network was forced to act as a lossy funnel, leading to catastrophic forgetting over long dependencies.

The radical genius of the Attention Mechanism was to shatter this funnel. By retaining the explicit representations of all tokens and computing a dynamic $T \times T$ similarity matrix, Attention allowed networks to “look back” globally. It transformed neural networks from sequential processors into dynamic, parallel graph-routing engines, where information flows based on semantic relevance rather than temporal proximity.

However, this breakthrough introduced its own poison pill: quadratic complexity. The equation $\text{Softmax}(QK^T)V$ meant that doubling the sequence length quadrupled the computational cost.

As we have explored in this notebook, the history of architecture since 2017 has been a relentless war against this $O(T^2)$ scaling:

1. Multi-Head Attention (MHA): Broadened the representational capacity of the model.
2. Multi-Query & Grouped-Query Attention (MQA/GQA): Solved the KV Cache memory limits during auto-regressive decoding.
3. Sliding Window & Sparse Attention: Restricted the mathematical receptive field to achieve linear $O(T)$ compute.
4. Multi-Head Latent Attention (MLA): Compressed the semantic dimensions fundamentally via RoPE decoupling.
5. Linear Attention: Altered the foundational kernel to exploit matrix associativity.

Attention is no longer just a “mechanism” inside a neural network; it is a fundamental mathematical paradigm for information retrieval. It teaches us that intelligence in deep learning is not just about having powerful features, but about the ability to dynamically align, route, and weigh those features relative to a specific query. Regardless of how the underlying hardware or scaling laws evolve in the future, the foundational principle of Attention—dynamic, data-dependent context aggregation—will remain the bedrock of artificial reasoning.

References