26. Linear (Kernelized) Attention: A Softmax-Free Alternative

Problem

Softmax attention computes an explicit (seq, seq) score matrix — O(seq²) time and memory, with no way around it, because softmax's normalization couples every key to every query. Linear attention replaces softmax(QKT) with a positive feature map phi applied to Q and K, then regroups phi(Q) phi(K)T V by associativity into phi(Q) (phi(K)T V) — computing the (d_k, d_v) matrix phi(K)T V first, never the full (seq, seq) matrix.

For causal attention specifically, this regrouping becomes an explicit linear recurrence: maintain a running (d_k, d_v) state matrix S and a running normalizer Z, updating both by one rank-1 outer-product term per new token — O(1) work per step, with a fixed-size state that never grows, unlike a softmax-attention KV-cache. This recurrent-state view is the same shape of tradeoff that state-space / linear-recurrent sequence models (an alternative model family to attention entirely) are built around: a fixed-size state cheap to run sequentially at inference, versus a parallel/batched form of the same recurrence used at training time.

Source: src/26_linear_attention.py

def causal_linear_attention(
    q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
    feature_map=None,
) -> torch.Tensor: ...

>>> q = torch.randn(1, 4, 3)
>>> k = torch.randn(1, 4, 3)
>>> v = torch.randn(1, 4, 2)
>>> causal_linear_attention(q, k, v).shape
torch.Size([1, 4, 2])

Step-by-Step Approach

1. Apply a positive elementwise feature map to q and k (default elu(x) + 1) — positivity is required so that the accumulated "attention weight" terms can never go negative, the same role softmax's exponential plays.
2. Initialize a running state matrix S of shape (batch, d_k, d_v) and a running normalizer Z of shape (batch, d_k), both zero.
3. At each step t, update the state with a rank-1 outer-product term: S += outer(phi(k_t), v_t), and the normalizer with Z += phi(k_t) — this is the entire "new information" contributed by token t.
4. Compute this step's output as phi(q_t) @ S / (phi(q_t) . Z): the numerator is a (d_k,) @ (d_k, d_v) -> (d_v,) contraction, the denominator a scalar dot product (add a small epsilon to avoid division by zero).
5. Loop over all sequence positions accumulating S and Z in place and stacking each step's output — note the state never grows with sequence length, unlike a softmax KV-cache that grows by one key/value pair per token.
6. Sanity-check position 0 (no history yet): S after step 0 is exactly outer(phi(k_0), v_0), so phi(q_0) @ S = (phi(q_0) . phi(k_0)) * v_0 and the normalizer is phi(q_0) . phi(k_0) — those cancel, so output 0 must equal v_0 regardless of what q_0 is.

The key insight is that dropping softmax's nonlinearity for a plain feature-map dot product is what makes attention associative, and associativity is what turns an O(seq²) all-pairs computation into an O(seq) sequential state update -- exactly the same algebraic move that lets linear-recurrent sequence models maintain a fixed-size hidden state instead of an ever-growing cache.

Reference solution

import torch
import torch.nn.functional as F


def _default_feature_map(x: torch.Tensor) -> torch.Tensor:
    return F.elu(x) + 1.0  # smooth, always positive


def causal_linear_attention(q, k, v, feature_map=None) -> torch.Tensor:
    if feature_map is None:
        feature_map = _default_feature_map

    batch, seq_len, d_k = q.shape
    d_v = v.shape[-1]

    phi_q = feature_map(q)
    phi_k = feature_map(k)

    state = torch.zeros(batch, d_k, d_v, device=q.device, dtype=q.dtype)
    normalizer = torch.zeros(batch, d_k, device=q.device, dtype=q.dtype)

    outputs = []
    for t in range(seq_len):
        state = state + phi_k[:, t, :].unsqueeze(-1) * v[:, t, :].unsqueeze(1)
        normalizer = normalizer + phi_k[:, t, :]

        numerator = torch.einsum("bd,bdv->bv", phi_q[:, t, :], state)
        denominator = torch.einsum("bd,bd->b", phi_q[:, t, :], normalizer).unsqueeze(-1)
        outputs.append((numerator / (denominator + 1e-6)).unsqueeze(1))

    return torch.cat(outputs, dim=1)

Key Functions & Tricks

• F.elu(x) + 1.0 — a standard positive feature map for linear attention; simpler alternatives like F.relu(x) + eps also appear in the literature, but ELU's smoothness avoids a dead-zone gradient at exactly zero.
• phi_k[:, t, :].unsqueeze(-1) * v[:, t, :].unsqueeze(1) — builds the rank-1 outer product (d_k, 1) * (1, d_v) -> (d_k, d_v) via broadcasting, no explicit torch.outer needed once a batch dim is present.
• torch.einsum("bd,bdv->bv", ...) — contracts the query's feature vector against the state matrix per batch element; equivalent to a batched matrix-vector product but reads more directly as "sum over d."
• Deferred, per-step normalization by phi(q_t) . Z — the linear-attention analog of softmax's denominator, computed from the same running accumulator as the numerator rather than a separate pass.
• Equivalence check against a directly-masked kernelized score matrix (phi(Q) phi(K)T with a causal mask, no recurrence) — confirms the sequential state recurrence and the "batch" all-pairs kernelized form compute the identical function, just at different memory/compute tradeoffs.

How to Recognize This Pattern