What is linear attention and how does it attempt to solve quadratic complexity?

Short answer: Linear attention approximates full self-attention by replacing the softmax exponential similarity term $\text{Softmax}(QK^T)$ with kernel feature maps $\phi(Q)\phi(K)^T$, allowing matrix multiplication associativity to re-order computation from $(Q K^T) V \Rightarrow \phi(Q) (\phi(K)^T V)$, reducing time and memory complexity from $O(N^2)$ to $O(N)$.

Detail

Standard self-attention computes an $N \times N$ attention matrix for sequence length $N$:

$\text{Attention}(Q, K, V) = \text{Softmax}\left(\frac{Q K^T}{\sqrt{d}}\right) V \quad \implies \mathcal{O}(N^2 d) \text{ time and memory}$

Standard Attention:   (Q [N x d] @ K^T [d x N]) [N x N Matrix!] @ V [N x d]  ──► O(N^2)
Linear Attention:     phi(Q) [N x d] @ (phi(K)^T [d x N] @ V [N x d])        ──► O(N d^2)

Mathematical Associativity Trick

By decomposing kernel similarity $k(q, k) = \phi(q)^T \phi(k)$:

$\text{LinearAttention}(Q, K, V) = \phi(Q) \left( \sum_{i=1}^N \phi(K_i)^T V_i \right)$

When sequence length $N \gg d$, computing the $d \times d$ matrix $(\phi(K)^T V)$ first scales linearly $O(N)$ with sequence length.

Example

PyTorch implementation concept for kernel-based Linear Attention:

import torch
import torch.nn as nn
import torch.nn.functional as F

class LinearAttention(nn.Module):
    def __init__(self, d_model=512):
        super().__init__()
        # ELU + 1 kernel feature map guarantees positive values
        self.feature_map = lambda x: F.elu(x) + 1.0

    def forward(self, q, k, v):
        Q = self.feature_map(q)
        K = self.feature_map(k)

        # Multiply (K^T @ V) first -> [d x d] matrix instead of [N x N]
        KV = torch.matmul(K.transpose(-2, -1), v) # [d x d]
        out = torch.matmul(Q, KV) # [N x d]
        return out

Interview tips

• Highlight why pure linear attention is rarely used in flagship LLMs: removing non-linear Softmax degrades multi-hop retrieval and long-context needle-in-a-haystack recall.
• Connect linear attention to modern State Space Models (Mamba) and Recurrent Transformers (RWKV).

Related Concepts

• [[How do state space models (SSMs) like Mamba compare to Transformer self-attention?]] (#109): How do state space models (SSMs) like Mamba compare to Transformer self-attention?
• [[How does vLLM PagedAttention manage Virtual Memory Pages to eliminate KV cache fragmentation?]] (#159): How does vLLM PagedAttention manage Virtual Memory Pages to eliminate KV cache fragmentation?
• [[What is an attention mask and why is it needed during batch processing?]] (#104): What is an attention mask and why is it needed during batch processing?

⬅ Back to LLM Fundamentals · All topics