Linear attention: change the order.

Standard attention constructs pairwise interactions between tokens. As the sequence grows, that matrix grows quadratically. Kernel-based linear attention asks whether we can aggregate the keys and values first, then query the aggregate.

The answer involves changing the attention kernel. You cannot simply move parentheses through the softmax operation.

The associative trick

Replace softmax attention with a positive feature-map kernel \(\phi(q)^\top\phi(k)\). For one query, the normalized output is:

Define a matrix state S and a normalization state z:

For fixed feature and value dimensions, accumulating those states scales linearly with sequence length. The cost still depends on the dimensions: building S requires work proportional to \(n d_\phi d_v\).

An implementation you can check

Here, \(\phi(x)=\operatorname{ELU}(x)+1\). The direct implementation builds every pairwise kernel score; the associative implementation never builds that matrix. They should agree up to floating-point error.

import numpy as np
from numpy.typing import NDArray


def phi(x: NDArray[np.float64]) -> NDArray[np.float64]:
    # exp(min(x, 0)) avoids overflowing the unused branch.
    return np.where(x >= 0, x + 1, np.exp(np.minimum(x, 0)))


def linear_attention(q: NDArray[np.float64], k: NDArray[np.float64], v: NDArray[np.float64]) -> NDArray[np.float64]:
    qf, kf = phi(q), phi(k)
    state = kf.T @ v               # (features, value_dim)
    normalizer = kf.sum(axis=0)    # (features,)
    denominator = qf @ normalizer
    return (qf @ state) / denominator[:, None].clip(1e-8)


rng = np.random.default_rng(7)
q, k = rng.normal(size=(2, 12, 8))
v = rng.normal(size=(12, 6))

pairwise = phi(q) @ phi(k).T
reference = (pairwise @ v) / pairwise.sum(-1, keepdims=True)
output = linear_attention(q, k, v)

np.testing.assert_allclose(output, reference, atol=1e-10)
print(output.shape)  # (12, 6)

This example is non-causal: every query can access all keys. It is a different attention rule from softmax attention, even though both return normalized weighted mixtures.

The causal recurrence

For autoregressive use, maintain prefix states. Add the current key/value pair before answering the current query:

Then evaluate the same normalized expression with \(S_t\) and \(z_t\). The state size does not grow with the number of tokens. This connection between linear attention and recurrent computation is developed in Transformers are RNNs ↗ .

The trade-off

The aggregate compresses previous keys and values into a fixed-size state. That changes what interactions the model can express. Lower asymptotic cost does not guarantee better speed at every sequence length, or the same model quality as softmax attention.

Further reading

See Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention ↗ . Compare this recurrence with a conventional KV cache , which retains individual keys and values.