Attention That You Probably Didnt Know Existed!!

I've been looking at self attention since 2019, since then a lof of variants of self-attention has been proposed. In this article, I breakdown efficient attention mechanisms: linear attention, sparse patterns, low-rank approximations, grouped attention and a lot of other attention as well. Time to pay attention to self-attention xD.

The Quadratic Bottleneck

Standard self-attention computes:

The matrix $QK^T$ has shape $[n \times n]$. For a 100K token sequence with 128 attention heads across 80 layers, you're looking at:

Obviously impossible. Even at 4K tokens this term dominates memory. So the question every method below is answering is: can we get something close to attention without ever materializing the full $n \times n$ matrix?

Approach 1: Sparse Attention Patterns

The simplest idea is to not compute all $n^2$ scores. Each token attends to a subset of positions and ignores the rest.

Fixed Patterns (Longformer, BigBird)

Longformer and BigBird hard-code the sparsity pattern. Every token attends to a local window of $w$ neighbors, a handful of special global tokens attend to (and are attended by) everything, and BigBird adds a sprinkling of random connections so information can hop across the sequence.

def create_longformer_mask(seq_len: int, window_size: int, global_tokens: int):
    """
    Creates Longformer-style attention mask
    Returns: [seq_len, seq_len] boolean mask
    """
    mask = torch.zeros(seq_len, seq_len, dtype=torch.bool)
    
    # Local sliding window
    for i in range(seq_len):
        start = max(0, i - window_size // 2)
        end = min(seq_len, i + window_size // 2 + 1)
        mask[i, start:end] = True
    
    # Global tokens (first k tokens attend to/from everything)
    mask[:global_tokens, :] = True
    mask[:, :global_tokens] = True
    
    return mask

Complexity: $O(n \cdot w)$ where $w$ is the window size. Linear in sequence length.

Attention patterns: Full n² attention, Sliding window, Dilated sliding window, Global+sliding window

The weakness is the assumption baked into the pattern: that locality is what matters. For tasks that need long-range reasoning you are hoping information propagates through the sparse graph, and sometimes it doesn't.

Learned Sparse Patterns (Routing Transformer, Reformer)

The alternative is to let the model decide which tokens to attend to.

Reformer uses Locality-Sensitive Hashing (LSH):

Tokens with the same hash are grouped into buckets and only attend within buckets.

def lsh_attention(Q, K, V, num_hashes: int = 4, num_buckets: int = 32):
    """Simplified LSH attention (conceptual)"""
    batch, seq_len, d = Q.shape
    
    # Generate random projection matrix
    R = torch.randn(d, num_buckets // 2)
    
    # Hash queries and keys
    projections = Q @ R  # [batch, seq_len, num_buckets//2]
    hashes = torch.argmax(
        torch.cat([projections, -projections], dim=-1), dim=-1
    )
    
    # Sort by hash to group similar tokens
    sorted_indices = torch.argsort(hashes, dim=1)
    Q_sorted = torch.gather(Q, 1, sorted_indices.unsqueeze(-1).expand_as(Q))
    K_sorted = torch.gather(K, 1, sorted_indices.unsqueeze(-1).expand_as(K))
    V_sorted = torch.gather(V, 1, sorted_indices.unsqueeze(-1).expand_as(V))
    
    # Attend within chunks (buckets)
    chunk_size = seq_len // num_buckets
    # ... chunked attention within each bucket
    
    return output

Complexity: $O(n \log n)$ for LSH bucketing + $O(n \cdot b)$ for within-bucket attention.

The weakness here is that hashing is probabilistic. Similar tokens can land in different buckets, and the fix, running several hash rounds, eats into the efficiency you were after.

Approach 2: Linear Attention (Kernel Methods)

This is the cleanest approach mathematically: rewrite attention so the $n \times n$ matrix never appears.

The Kernel Trick

Standard softmax attention:

Replace the exponential kernel with a decomposable feature map $\phi$:

Now attention becomes:

What makes this work is that $\sum_j \phi(k_j) v_j^T$ and $\sum_j \phi(k_j)$ don't depend on the query. Compute them once and every query reuses them.

def linear_attention(Q, K, V, feature_map=None):
    """
    Linear attention via kernel feature maps
    Complexity: O(n * d^2) instead of O(n^2 * d)
    """
    if feature_map is None:
        # ELU-based feature map (from "Transformers are RNNs")
        feature_map = lambda x: F.elu(x) + 1
    
    Q = feature_map(Q)  # [batch, seq, d]
    K = feature_map(K)  # [batch, seq, d]
    
    # Compute KV summary: [batch, d, d_v]
    KV = torch.einsum('bnd,bnv->bdv', K, V)
    
    # Compute K summary for normalization: [batch, d]
    K_sum = K.sum(dim=1)
    
    # Compute attention output
    numerator = torch.einsum('bnd,bdv->bnv', Q, KV)
    denominator = torch.einsum('bnd,bd->bn', Q, K_sum).unsqueeze(-1)
    
    return numerator / (denominator + 1e-6)

Complexity: $O(n \cdot d^2)$. Linear in sequence length.

Popular Linear Attention Variants

Performer: Random Feature Maps

Performer provides an unbiased estimator of softmax attention using random Fourier features:

def random_fourier_features(x, num_features: int = 256):
    """
    FAVOR+ feature map from Performer
    """
    d = x.shape[-1]
    
    # Sample random projections (orthogonal for lower variance)
    W = torch.randn(num_features, d) / (d ** 0.5)
    W = torch.linalg.qr(W.T)[0].T[:num_features]  # Orthogonalize
    
    # Compute features
    proj = x @ W.T  # [batch, seq, num_features]
    
    # Normalize for unbiased estimate
    norm_factor = (x ** 2).sum(dim=-1, keepdim=True) / 2
    
    features = torch.exp(proj - norm_factor) / (num_features ** 0.5)
    return features

The problem with linear attention in general is that it doesn't approximate softmax well for every distribution. What you lose is the sharp, selective quality of softmax attention, where a query can put nearly all its weight on one key. For many tasks that loss shows up directly as worse quality.

Approach 3: Low-Rank Approximation (Linformer, Nyströmformer)

If the attention matrix has low intrinsic rank, project it to a smaller space.

Linformer

Project keys and values to a fixed dimension $k \ll n$:

where $E, F \in \mathbb{R}^{k \times n}$ are learned projection matrices.

class LinformerAttention(nn.Module):
    def __init__(self, d_model: int, seq_len: int, k: int = 256):
        super().__init__()
        self.k = k
        
        # Projection matrices
        self.E = nn.Parameter(torch.randn(k, seq_len) / seq_len ** 0.5)
        self.F = nn.Parameter(torch.randn(k, seq_len) / seq_len ** 0.5)
        
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
    
    def forward(self, x):
        Q = self.W_q(x)  # [batch, n, d]
        K = self.W_k(x)  # [batch, n, d]
        V = self.W_v(x)  # [batch, n, d]
        
        # Project K and V: [batch, k, d]
        K_proj = torch.einsum('kn,bnd->bkd', self.E, K)
        V_proj = torch.einsum('kn,bnd->bkd', self.F, V)
        
        # Standard attention on reduced dimensions
        # [batch, n, d] @ [batch, d, k] = [batch, n, k]
        scores = torch.matmul(Q, K_proj.transpose(-2, -1)) / (Q.shape[-1] ** 0.5)
        attn = F.softmax(scores, dim=-1)
        
        # [batch, n, k] @ [batch, k, d] = [batch, n, d]
        return torch.matmul(attn, V_proj)

Complexity: $O(n \cdot k)$ where $k$ is the projection dimension.

Two weaknesses. A fixed projection dimension is a bet that the attention matrix is always low-rank, which holds for some tasks and not others. And because the projection matrices have a sequence-length dimension baked in, variable-length inputs are awkward.

Nyströmformer

Uses the Nyström method to approximate the full attention matrix from a subset of "landmark" tokens:

where $\tilde{K}$ are landmark keys selected via pooling.

Approach 4: State Space Models (Mamba, S4)

The drastic option: drop attention entirely and replace it with a structured state space model that processes the sequence in $O(n)$ time.

The State Space Formulation

A continuous-time state space model:

Discretize and unroll:

This is an RNN. The difference is the structure imposed on $A$, which allows the same model to be trained in parallel as a convolution and run at inference as a linear-time recurrence.

Mamba: Selective State Spaces

Mamba makes the state space parameters input-dependent:

That selectivity lets the model decide, per token, what to remember and what to forget. It recovers some of the content-based addressing that attention gives you for free.

class SimplifiedMambaBlock(nn.Module):
    """Conceptual Mamba block (actual impl uses CUDA kernels)"""
    
    def __init__(self, d_model: int, d_state: int = 16, d_conv: int = 4):
        super().__init__()
        self.d_model = d_model
        self.d_state = d_state
        
        # SSM parameters
        self.A_log = nn.Parameter(torch.randn(d_model, d_state))
        
        # Input-dependent projections
        self.x_proj = nn.Linear(d_model, d_state * 2 + 1)  # B, C, delta
        
        # Convolution for local context
        self.conv = nn.Conv1d(d_model, d_model, d_conv, padding=d_conv-1, groups=d_model)
        
    def forward(self, x):
        batch, seq_len, d = x.shape
        
        # Local convolution
        x_conv = self.conv(x.transpose(1, 2))[:, :, :seq_len].transpose(1, 2)
        x_conv = F.silu(x_conv)
        
        # Compute input-dependent SSM parameters
        x_dbl = self.x_proj(x_conv)
        B = x_dbl[..., :self.d_state]
        C = x_dbl[..., self.d_state:2*self.d_state]
        delta = F.softplus(x_dbl[..., -1:])
        
        # Discretize A
        A = -torch.exp(self.A_log)
        A_bar = torch.exp(delta * A)
        
        # Recurrent scan (parallel via associative scan in practice)
        h = torch.zeros(batch, d, self.d_state, device=x.device)
        outputs = []
        for t in range(seq_len):
            h = A_bar[:, t:t+1] * h + delta[:, t:t+1] * B[:, t:t+1].unsqueeze(1) * x_conv[:, t:t+1].unsqueeze(-1)
            y_t = (h * C[:, t:t+1].unsqueeze(1)).sum(-1)
            outputs.append(y_t)
        
        return torch.cat(outputs, dim=1)

Complexity: $O(n)$ time and memory.

The trade is that the quadratic cost has been swapped for a fixed-size hidden state, and every long-range dependency has to survive being compressed into it. On tasks that need precise retrieval from far back in the context (the needle-in-a-haystack family), SSMs still lose to attention.

Approach 5: IO-Aware Algorithms (FlashAttention)

Maybe the algorithm was never the problem. Maybe the implementation was.

The Memory Hierarchy Insight

Standard attention is memory-bound, not compute-bound. The $O(n^2)$ attention matrix is written out to slow HBM (GPU main memory), then read back for the softmax, then read back again to compute the output.

FlashAttention keeps everything in fast SRAM by:

1. Tiling the computation into blocks
2. Recomputing during backward pass instead of storing
3. Fusing operations to minimize memory transfers

Standard Attention Memory Access:
Q, K, V (HBM) → Compute QK^T → Store S (HBM) → Load S → Softmax → Store P (HBM) → Load P, V → Output
 
FlashAttention Memory Access:  
Q, K, V (HBM) → [Tiled compute entirely in SRAM] → Output (HBM)

Complexity: Still $O(n^2)$ compute, but $O(n)$ memory, and 2-4x faster in practice because far less data moves between HBM and SRAM.

FlashAttention doesn't change the asymptotic cost. It makes the quadratic algorithm practical by respecting the memory hierarchy, and for a lot of workloads that turns out to be all you needed.

Approach 6: Grouped-Query Attention (GQA)

This one gets less airtime than it deserves given that it's in most modern LLMs. The premise: what if the bottleneck isn't computing attention but storing the key-value cache?

The KV Cache Problem

During autoregressive generation, we cache key and value tensors from previous tokens to avoid recomputation. With standard Multi-Head Attention (MHA), each head has its own K and V projections:

For a 70B model with 80 layers and 64 heads, generating 8K tokens requires ~40GB just for KV cache. That's often more memory than the model weights themselves.

Multi-Query Attention (MQA)

The blunt fix: all query heads share a single set of keys and values.

Each query head attends using the same K and V:

KV cache reduction: $h \times$ smaller (64× for a 64-head model).

The cost is quality. Different heads want to learn different attention patterns, and forcing them onto shared keys and values gets in the way.

Grouped-Query Attention (GQA): The Sweet Spot

GQA is the Goldilocks solution: instead of 1 KV set (MQA) or $h$ KV sets (MHA), use $g$ groups where $1 < g < h$.

Each group of $h/g$ query heads shares one set of keys and values:

Multi-Head, Grouped-Query, and Multi-Query Attention: Trading KV cache size for quality

class GroupedQueryAttention(nn.Module):
    def __init__(self, d_model: int, n_heads: int, n_kv_heads: int):
        """
        n_heads: number of query heads
        n_kv_heads: number of key-value heads (groups)
        """
        super().__init__()
        self.n_heads = n_heads
        self.n_kv_heads = n_kv_heads
        self.n_rep = n_heads // n_kv_heads  # queries per KV group
        self.head_dim = d_model // n_heads
        
        self.W_q = nn.Linear(d_model, n_heads * self.head_dim)
        self.W_k = nn.Linear(d_model, n_kv_heads * self.head_dim)
        self.W_v = nn.Linear(d_model, n_kv_heads * self.head_dim)
        self.W_o = nn.Linear(d_model, d_model)
    
    def forward(self, x):
        batch, seq_len, _ = x.shape
        
        # Project
        Q = self.W_q(x).view(batch, seq_len, self.n_heads, self.head_dim)
        K = self.W_k(x).view(batch, seq_len, self.n_kv_heads, self.head_dim)
        V = self.W_v(x).view(batch, seq_len, self.n_kv_heads, self.head_dim)
        
        # Repeat K, V to match query heads
        # [batch, seq, n_kv_heads, head_dim] -> [batch, seq, n_heads, head_dim]
        K = K.repeat_interleave(self.n_rep, dim=2)
        V = V.repeat_interleave(self.n_rep, dim=2)
        
        # Standard attention
        Q, K, V = Q.transpose(1, 2), K.transpose(1, 2), V.transpose(1, 2)
        scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.head_dim ** 0.5)
        attn = F.softmax(scores, dim=-1)
        out = torch.matmul(attn, V)
        
        out = out.transpose(1, 2).reshape(batch, seq_len, -1)
        return self.W_o(out)

The Math: Memory Savings

Llama 2 70B uses GQA with 8 KV heads and 64 query heads, an 8× reduction in KV cache with minimal quality loss.

GQA doesn't touch computational complexity (still $O(n^2)$). What it cuts is memory bandwidth during generation, and when you're producing tokens one at a time, loading the KV cache is usually the bottleneck. In practice that makes inference 2-4× faster.

Used by: Llama 2/3, Mistral, Gemma, most modern open-source LLMs.

The Hybrid Future

No single approach dominates. The trend is hybrid architectures:

Sliding Window + Global (Mistral, Gemma 2)

Layer 1-4:   Sliding window attention (local)
Layer 5:    Global attention (every 4th layer)
Layer 6-9:   Sliding window attention (local)
Layer 10:   Global attention
...

This gets $O(n)$ for most layers while preserving some global context mixing.

Attention + SSM (Jamba, Zamba)

Interleave attention layers with Mamba layers. The Mamba layers carry information across long distances cheaply; the attention layers do precise retrieval where it's needed.

class HybridLayer(nn.Module):
    def __init__(self, d_model: int, use_attention: bool):
        super().__init__()
        if use_attention:
            self.mixer = MultiHeadAttention(d_model, num_heads=8)
        else:
            self.mixer = MambaBlock(d_model)
        self.norm = nn.LayerNorm(d_model)
        self.ffn = FeedForward(d_model)
    
    def forward(self, x):
        x = x + self.mixer(self.norm(x))
        x = x + self.ffn(self.norm(x))
        return x
 
# Example: 1 attention layer per 4 Mamba layers
layers = []
for i in range(32):
    use_attention = (i % 4 == 0)
    layers.append(HybridLayer(d_model=2048, use_attention=use_attention))

Comparison: What Actually Works?

The uncomfortable truth: FlashAttention plus sliding-window hybrids currently win for most practical applications. Pure linear attention and SSMs give up too much quality for their theoretical efficiency.

Final Thoughts: The Quadratic Wall Is Softer Than It Looks

FlashAttention reset the debate: $O(n^2)$ compute is fine as long as memory stays $O(n)$. Hybrids are the practical answer, mixing cheap local attention with occasional global layers. SSMs are complements rather than replacements; they buy efficiency, but attention still wins whenever you need to retrieve something specific. And hardware may settle the argument anyway, since a GPU with a lot more SRAM would make much of this moot.

The search for $O(n)$ attention isn't finished, but the practical ceiling has already moved from 4K to 1M+ tokens. For most applications, that's enough.

Next up: How Mamba's selective state spaces actually work, and why the "selectivity" mechanism is the key innovation.

References

1. Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., & Polosukhin, I. (2017). Attention Is All You Need. NeurIPS. arXiv:1706.03762

2. Beltagy, I., Peters, M. E., & Cohan, A. (2020). Longformer: The Long-Document Transformer. arXiv preprint. arXiv:2004.05150

3. Zaheer, M., Guruganesh, G., Dubey, A., Ainslie, J., Alberti, C., Ontanon, S., Pham, P., Ravula, A., Wang, Q., Yang, L., & Ahmed, A. (2020). Big Bird: Transformers for Longer Sequences. NeurIPS. arXiv:2007.14062

4. Kitaev, N., Kaiser, Ł., & Levskaya, A. (2020). Reformer: The Efficient Transformer. ICLR. arXiv:2001.04451

5. Katharopoulos, A., Vyas, A., Pappas, N., & Fleuret, F. (2020). Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention. ICML. arXiv:2006.16236

6. Choromanski, K., Likhosherstov, V., Dohan, D., Song, X., Gane, A., Sarlos, T., Hawkins, P., Davis, J., Mohiuddin, A., Kaiser, L., Belanger, D., Colwell, L., & Weller, A. (2021). Rethinking Attention with Performers. ICLR. arXiv:2009.14794

7. Qin, Z., Sun, W., Deng, H., Li, D., Wei, Y., Lv, B., Yan, J., Kong, L., & Zhong, Y. (2022). cosFormer: Rethinking Softmax in Attention. ICLR. arXiv:2202.08791

8. Wang, S., Li, B., Khabsa, M., Fang, H., & Ma, H. (2020). Linformer: Self-Attention with Linear Complexity. arXiv preprint. arXiv:2006.04768

9. Xiong, Y., Zeng, Z., Chakraborty, R., Tan, M., Fung, G., Li, Y., & Singh, V. (2021). Nyströmformer: A Nyström-Based Algorithm for Approximating Self-Attention. AAAI. arXiv:2102.03902

10. Gu, A., Goel, K., & Ré, C. (2022). Efficiently Modeling Long Sequences with Structured State Spaces. ICLR. arXiv:2111.00396

11. Gu, A., & Dao, T. (2023). Mamba: Linear-Time Sequence Modeling with Selective State Spaces. arXiv preprint. arXiv:2312.00752

12. Dao, T., Fu, D. Y., Ermon, S., Rudra, A., & Ré, C. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS. arXiv:2205.14135

13. Dao, T. (2023). FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. arXiv preprint. arXiv:2307.08691

14. Peng, B., Alcaide, E., Anthony, Q., Albalak, A., Arcadinho, S., Biderman, S., Cao, H., Cheng, X., Chung, M., et al. (2023). RWKV: Reinventing RNNs for the Transformer Era. EMNLP. arXiv:2305.13048

15. Lieber, O., Lenz, B., Ratner, E., et al. (2024). Jamba: A Hybrid Transformer-Mamba Language Model. arXiv preprint. arXiv:2403.19887

16. Jiang, A. Q., Sablayrolles, A., Roux, A., et al. (2024). Mixtral of Experts. arXiv preprint. arXiv:2401.04088

17. Ainslie, J., Lee-Thorp, J., de Jong, M., Zemlyanskiy, Y., Lebrón, F., & Sanghai, S. (2023). GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. EMNLP. arXiv:2305.13245

18. Shazeer, N. (2019). Fast Transformer Decoding: One Write-Head is All You Need. arXiv preprint. arXiv:1911.02150