Linear attention and state-space models

Standard softmax attention computes $\text{output} = \text{softmax}(Q @ \frac{K^{T}}{\sqrt{d}}) @ V$ , where the Q @ K^T multiplication produces an N×N matrix — the source of quadratic cost. Linear attention replaces the softmax with a decomposable kernel function, enabling a rewrite that avoids materializing this N×N matrix entirely.

The kernel trick in attention

Katharopoulos et al. (2020) observed that softmax attention can be viewed as a kernel function. Define a feature map φ such that softmax(q_i^T k_j / sqrt(d)) ≈ φ(q_i)^T φ(k_j). With this substitution, the attention output for token i becomes:

o_{i} = \Sigma _j [φ(q_{i})^{T} φ(k_{j})] \cdot \frac{v_{j}}{\Sigma }_j [φ(q_{i})^{T} φ(k_{j})]

The critical insight is associativity of matrix multiplication. The naive computation groups as (φ(Q) @ φ(K)^T) @ V — still O(N²d) because φ(Q) @ φ(K)^T is N×N. But regrouping as φ(Q) @ (φ(K)^T @ V) computes φ(K)^T @ V first — a d × d matrix (assuming φ maps to d dimensions) — then multiplies each query by this matrix. Total cost: O(Nd²), which is linear in N when d << N.

For typical transformer dimensions: d = 128 (head_dim), N = 128K. The quadratic term $N^{2} \times d = 128000^{2} \times 128 ≈ 2.1 \text{trillion} \text{ops}$ . The linear term $N \times d^{2} = 128000 \times 128^{2} ≈ 2.1 \text{billion} \text{ops}$ — a 1000x reduction.

Causal linear attention as a recurrence

For autoregressive generation, causal linear attention maintains a running state. Define:

$S_{t} = \Sigma _{j = 1}^{{}t} φ(k_{j}) @ v_{j}^{T}$ (a d_k × d_v matrix)

$z_{t} = \Sigma _{j = 1}^{{}t} φ(k_{j})$ (a d_k vector, for normalization)

The output at position t is:

o_{t} = \frac{(φ(q_{t})^{T} @ S_{t})}{(φ(q_{t})^{T} @ z_{t})}

For each new token, update the state: $S_{t + 1} = S_{t} + φ(k_{t + 1}) @ v_{t + 1}^{T}$ . This is O(d²) per token — constant with respect to sequence length. The state S is a d_k × d_v matrix: at head_dim = 128, that's $128 \times 128 \times 2 \text{bytes} = 32 \text{KB}$ per head in float16. Compare this to the KV cache in standard attention, which stores all past keys and values: at 128K context, $128000 \times 128 \times 2 \times 2 = 64 \text{MB}$ per head.

python

import torch
import torch.nn.functional as F

def elu_feature_map(x: torch.Tensor) -> torch.Tensor:
    """ELU+1 feature map from Katharopoulos et al. 2020."""
    return F.elu(x) + 1

class CausalLinearAttention(torch.nn.Module):
    """Linear attention with causal masking, using a recurrent state."""

    def __init__(self, d_model: int = 512, n_heads: int = 8):
        super().__init__()
        self.n_heads = n_heads
        self.head_dim = d_model // n_heads
        self.q_proj = torch.nn.Linear(d_model, d_model)
        self.k_proj = torch.nn.Linear(d_model, d_model)
        self.v_proj = torch.nn.Linear(d_model, d_model)
        self.out_proj = torch.nn.Linear(d_model, d_model)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, N, _ = x.shape
        H, D = self.n_heads, self.head_dim

        q = elu_feature_map(self.q_proj(x).view(B, N, H, D))  # (B, N, H, D)
        k = elu_feature_map(self.k_proj(x).view(B, N, H, D))
        v = self.v_proj(x).view(B, N, H, D)

        # Cumulative sum formulation for causal linear attention
        # S_t = Σ_{j<=t} φ(k_j) ⊗ v_j  → maintained as cumsum
        kv = torch.einsum('bnhd,bnhe->bnhde', k, v)  # (B, N, H, D, D)
        S = torch.cumsum(kv, dim=1)  # running state: (B, N, H, D, D)

        # z_t = Σ_{j<=t} φ(k_j)
        z = torch.cumsum(k, dim=1)  # (B, N, H, D)

        # o_t = φ(q_t) @ S_t / (φ(q_t) @ z_t)
        num = torch.einsum('bnhd,bnhde->bnhe', q, S)  # (B, N, H, D)
        den = torch.einsum('bnhd,bnhd->bnh', q, z).unsqueeze(-1)  # (B, N, H, 1)
        out = num / (den + 1e-6)

        return self.out_proj(out.reshape(B, N, H * D))

    def generate_step(self, x_t: torch.Tensor, state: dict) -> tuple:
        """Single-step generation with O(d²) per token."""
        B, H, D = x_t.shape[0], self.n_heads, self.head_dim

        q = elu_feature_map(self.q_proj(x_t).view(B, H, D))  # (B, H, D)
        k = elu_feature_map(self.k_proj(x_t).view(B, H, D))
        v = self.v_proj(x_t).view(B, H, D)

        # Update running state: S += φ(k_t) ⊗ v_t
        state['S'] = state['S'] + torch.einsum('bhd,bhe->bhde', k, v)
        state['z'] = state['z'] + k

        # Compute output: O(d²) — independent of sequence length
        num = torch.einsum('bhd,bhde->bhe', q, state['S'])
        den = torch.einsum('bhd,bhd->bh', q, state['z']).unsqueeze(-1)
        out = num / (den + 1e-6)

        return self.out_proj(out.reshape(B, H * D)), state

The quality gap

Linear attention consistently underperforms softmax attention on language modeling perplexity. On WikiText-103 at 125M parameters, Katharopoulos et al. (2020) reported a perplexity gap of ~5 points (24.2 vs 19.1). The Performer (Choromanski et al. 2021) reduced this gap using random orthogonal features as the kernel, but still fell 2–3 perplexity points short of standard transformers.

The root cause is the softmax's sharpening effect. Softmax concentrates probability mass on a few highly-relevant tokens — in a typical attention head, 80–90% of the attention weight falls on 5–10% of the positions (Clark et al. 2019). The kernel approximation φ(q)^T φ(k) produces smoother, more diffuse attention distributions. It cannot represent the near-one-hot patterns that softmax attention uses for precise information retrieval (e.g., a closing bracket attending almost exclusively to its matching opening bracket).

Linformer: a different linear approach

Linformer (Wang et al. 2020) takes a different route: project K and V down to a fixed size k << N before computing attention.

K' = E_K @ K (k × N projection applied to N × d_k → k × d_k)

V' = E_V @ V

Then compute standard softmax attention with the projected keys/values: softmax(Q @ K'^T / sqrt(d)) @ V'. The attention matrix is now N × k instead of N × N — linear in N if k is fixed.

The projection matrices E_K and E_V are learned. Linformer's insight: the attention matrix in trained transformers is empirically low-rank. With k = 256, Linformer matches full attention quality on sequences up to 4K tokens. Beyond that, the fixed projection size becomes a bottleneck.

State-space models: Mamba

Mamba (Gu & Dao 2023) replaces attention entirely with a selective state-space model. The core operation at each position:

$h_{t} = A_{t} \cdot h_{t - 1} + B_{t} \cdot x_{t}$ (state update)

$y_{t} = C_{t} \cdot h_{t}$ (output)

Where h_t is a hidden state of size d_state (typically 16), and crucially, the matrices A_t, B_t, C_t are input-dependent — computed from x_t through learned projections. This selectivity is what distinguishes Mamba from classical SSMs (S4, which used fixed A, B, C): the model can choose what to remember and what to forget at each step based on the input.

Computational characteristics

Training: parallel scan over the sequence — O(N) work, O(log N) depth on GPU. Processed as a single fused kernel (the "selective scan" CUDA kernel).
Inference: pure recurrence — O(1) per new token, state size fixed at d_model × d_state × 2 bytes. For Mamba-2.8B with d_model = 2560 and d_state = 16: $2560 \times 16 \times 2 = 80 \text{KB}$ per layer, versus multi-GB KV caches in transformers at long contexts.
Throughput: Mamba-3B generates at 4–5x the tokens/sec of a similarly-sized transformer at 64K context (Gu & Dao 2023, Figure 6), because it never reads from or writes to a growing cache.

Mamba-2 and hybrid architectures

Mamba-2 (Dao & Gu 2024) reformulates the SSM as a structured matrix operation that maps to tensor cores, achieving 2–8x faster training than Mamba-1. It introduces a connection between SSMs and attention: the selective scan can be viewed as a form of linear attention with a specific decay structure.

Jamba (AI21, Lieber et al. 2024) interleaves transformer attention layers with Mamba layers in a single model: the first 7B-parameter configuration uses a ratio of 1 attention layer per 7 Mamba layers. The attention layers provide precise long-range retrieval (which SSMs struggle with), while Mamba layers handle the bulk of sequence processing efficiently. On standard language modeling benchmarks, Jamba matches pure transformer quality while using 2–3x less KV cache memory.

python

class SimpleMambaBlock(torch.nn.Module):
    """Minimal Mamba-style selective SSM block (simplified for clarity)."""

    def __init__(self, d_model: int = 512, d_state: int = 16, d_conv: int = 4):
        super().__init__()
        self.d_state = d_state

        # Input projections (expand to 2x for gating)
        self.in_proj = torch.nn.Linear(d_model, d_model * 2, bias=False)

        # Short convolution before SSM
        self.conv = torch.nn.Conv1d(
            d_model, d_model, kernel_size=d_conv,
            padding=d_conv - 1, groups=d_model
        )

        # Selectivity projections: input-dependent A, B, C
        self.dt_proj = torch.nn.Linear(d_model, d_model)
        self.B_proj = torch.nn.Linear(d_model, d_state)
        self.C_proj = torch.nn.Linear(d_model, d_state)

        # Learnable log of diagonal A matrix
        self.A_log = torch.nn.Parameter(
            torch.log(torch.arange(1, d_state + 1).float().repeat(d_model, 1))
        )

        self.out_proj = torch.nn.Linear(d_model, d_model, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, N, D = x.shape

        # Project and split into main path + gate
        xz = self.in_proj(x)
        x_main, z = xz.chunk(2, dim=-1)

        # Short convolution (local context mixing)
        x_conv = self.conv(x_main.transpose(1, 2))[:, :, :N].transpose(1, 2)
        x_conv = F.silu(x_conv)

        # Compute input-dependent SSM parameters
        dt = F.softplus(self.dt_proj(x_conv))  # (B, N, D) — discretization step
        B_t = self.B_proj(x_conv)  # (B, N, d_state)
        C_t = self.C_proj(x_conv)  # (B, N, d_state)
        A = -torch.exp(self.A_log)  # (D, d_state) — negative for stability

        # Selective scan (sequential for clarity; real impl uses parallel scan)
        h = torch.zeros(B, D, self.d_state, device=x.device)
        outputs = []
        for t in range(N):
            # Discretize: A_bar = exp(dt * A)
            A_bar = torch.exp(dt[:, t, :].unsqueeze(-1) * A)  # (B, D, d_state)
            B_bar = dt[:, t, :].unsqueeze(-1) * B_t[:, t, :].unsqueeze(1)  # (B, D, d_state)

            # State update: h_t = A_bar * h_{t-1} + B_bar * x_t
            h = A_bar * h + B_bar * x_conv[:, t, :].unsqueeze(-1)

            # Output: y_t = C_t @ h_t
            y_t = (h * C_t[:, t, :].unsqueeze(1)).sum(dim=-1)  # (B, D)
            outputs.append(y_t)

        y = torch.stack(outputs, dim=1)  # (B, N, D)

        # Gate and project
        y = y * F.silu(z)
        return self.out_proj(y)

Ring Attention: scaling full attention to millions of tokens

Ring Attention (Liu et al. 2023) does not approximate attention — it computes exact full attention by distributing the computation across GPUs. Each of P GPUs holds a contiguous chunk of N/P tokens (their Q, K, V). Computation proceeds in P rounds:

Round 1: each GPU computes attention between its local Q and its local KV.
Round 2: each GPU passes its KV block to the next GPU in a ring topology. Now each GPU computes attention between its Q and the neighbor's KV.
After P rounds, every GPU has computed attention against every KV block.

The key optimization: the KV communication (send to next neighbor) overlaps with the attention computation of the current block. If the compute time per block exceeds the communication time — which holds when blocks are large enough — the communication is fully hidden. Total memory per GPU: O(N/P) for activations + O(N/P) for local KV. Total compute: O(N²/P) per GPU — the same total work as full attention, just distributed.

Ring Attention with FlashAttention integration enabled training with 1M+ context at research labs. The Llama 3.1 405B long-context extension to 128K tokens used a variant of this approach for its continued pre-training phase (Dubey et al. 2024).

The landscape in 2026

The hierarchy of attention mechanisms by sequence length and quality:

N < 8K: full softmax attention, FlashAttention kernel. No reason to approximate.
8K < N < 128K: sliding window (Mistral) or GQA with FlashAttention handles most tasks. Use full attention only if long-range retrieval accuracy is critical.
N > 128K: Ring Attention for exact computation across GPUs, or hybrid Mamba-attention (Jamba) for single-GPU deployment. Pure Mamba for throughput-critical workloads where slight quality regression is acceptable.
Inference at any N: linear attention / Mamba for constant-memory generation. The recurrent state enables infinite-length generation without growing memory, at the cost of reduced retrieval precision for information stored many thousands of tokens earlier.

← Previous