Linear Attention

Attention is quadratic because of the order of two matrix multiplications, not because of the amount of information it moves. $Q K^{\top}$ is $n \times n$; $K^{\top} V$ is $d \times d$. The softmax between them is what forbids the second grouping, because $\softmax(QK^{\top})V$ does not factor.

Remove it. Write the similarity as an inner product of feature maps and the associativity comes back.

The normalisation carries over: divide by $\phi(q_i)^{\top} \sum_j \phi(k_j)$ rather than by a row sum of exponentials. The result is a weighted average with non-negative weights — just not a softmax, and not a sharp one.

It is a linear RNN

With the mask applied, the cumulative sums are a recurrent state: a $d \times d$ matrix $S$ and a $d$-vector $z$, both of fixed size. Decoding is then $O(d^2)$ per token with no cache that grows — the property that makes this family interesting, and the same property state space models arrive at from the other direction.

What is actually lost

A softmax can concentrate almost all its mass on one position. A fixed $d \times d$ state cannot: it is a lossy summary of everything seen so far, and the capacity is $d^2$ regardless of $n$. Exact recall of a specific earlier token — the thing induction heads do — degrades as the sequence exceeds that capacity.

The ungated recurrence in eq. 2 also only ever adds, so old information is never removed, only diluted. Gated variants insert a decay term, $S_t = G_t \odot S_{t-1} + \phi(k_t) v_t^{\top}$, which recovers much of the gap and is what the current generation of these models actually uses.

import torch
from torch import Tensor


def linear_attention_chunked(
    q: Tensor, k: Tensor, v: Tensor, chunk: int = 128  # [B, H, N, D]
) -> Tensor:
    """Chunked parallel form: quadratic within a chunk, recurrent across them."""
    q, k = torch.nn.functional.elu(q) + 1, torch.nn.functional.elu(k) + 1
    b, h, n, d = q.shape

    state = q.new_zeros(b, h, d, v.size(-1))   # running K^T V
    z = q.new_zeros(b, h, d)                   # running sum of phi(k)
    out = []

    for i in range(0, n, chunk):
        qc, kc, vc = q[..., i : i + chunk, :], k[..., i : i + chunk, :], v[..., i : i + chunk, :]

        # Contribution of every earlier chunk, through the summarised state.
        inter = qc @ state
        # Contribution of this chunk, computed the quadratic way and masked.
        mask = torch.ones(qc.size(-2), kc.size(-2), dtype=torch.bool, device=q.device).triu(1)
        intra = (qc @ kc.transpose(-2, -1)).masked_fill(mask, 0.0) @ vc

        denom = qc @ z.unsqueeze(-1) + (qc @ kc.transpose(-2, -1)).masked_fill(mask, 0.0).sum(-1, keepdim=True)
        out.append((inter + intra) / denom.clamp(min=1e-6))

        state = state + kc.transpose(-2, -1) @ vc
        z = z + kc.sum(dim=-2)

    return torch.cat(out, dim=-2)

The chunked form is the one to implement. The purely recurrent version is $n$ serial steps and trains at a fraction of the throughput; the purely parallel version reintroduces the $n^2$ term it was supposed to remove. Chunking pays $O(c^2)$ inside a chunk and $O(d^2)$ between chunks, and $c \approx 128$ is usually where the two balance on current hardware.