Chapter 28

Linear Attention & State-Space Models

Attention as recurrence, DeltaNet, Mamba, and hybrid linear/full-attention stacks.

Softmax attention is powerful, but its cache grows with the whole prefix. Linear attention and state-space models trade exact softmax weights for recurrent states, so decoding can keep a fixed-size summary instead of every key and value. That trade is central to long-context LLMs: most designs keep some full attention, then use linear or state-space layers where constant per-token state is worth the approximation.

28.1 Attention as a feature-map kernel

Causal softmax attention at time $t$ is a normalized weighted average over previous values. Linear attention replaces the exponential score by a positive feature-map dot product:

K(\vq_t, \vk_s) = \vphi(\vq_t)^\T \vphi(\vk_s) .\tag{28.1}

The output is

\vy_t = \frac{\sum_{s\le t} K(\vq_t,\vk_s)\vv_s} {\sum_{s\le t} K(\vq_t,\vk_s)} .\tag{28.2}

Because the kernel factorizes, move the query-dependent term outside the prefix sum:

\mS_t = \mS_{t-1} + \vphi(\vk_t)\vv_t^\T, \qquad \vc_t = \vc_{t-1} + \vphi(\vk_t) .\tag{28.3}

Then $\vy_t = \vphi(\vq_t)^\T\mS_t / (\vphi(\vq_t)^\T\vc_t)$ . The numerator state $\mS_t$ stores key-value outer products; the vector $\vc_t$ stores the normalizer. This is the whole trick from linear Transformers [katharopoulos2020transformers]. It changes the computation from a growing attention matrix to a recurrent update, while preserving a parallel form for training.

Listing 28.1 Causal linear attention: parallel and recurrent forms

def feature_map(x):
    """Positive ELU+1 feature map."""
    x = np.asarray(x)
    return np.where(x > 0, x + 1, np.exp(x))


def parallel_linear_attention(q, k, v, eps=1e-8):
    """Causal linear attention computed from the full lower triangle."""
    q_phi, k_phi = feature_map(q), feature_map(k)
    scores = q_phi @ k_phi.T
    scores *= np.tri(q.shape[0], dtype=scores.dtype)
    numerator = scores @ v
    denominator = scores.sum(axis=-1, keepdims=True)
    return numerator / np.maximum(denominator, eps)


def recurrent_linear_attention(q, k, v, eps=1e-8):
    """Causal linear attention using the recurrent state S and normalizer c."""
    q_phi, k_phi = feature_map(q), feature_map(k)
    state = np.zeros((k_phi.shape[1], v.shape[1]), dtype=q.dtype)
    normalizer = np.zeros(k_phi.shape[1], dtype=q.dtype)
    outputs = []
    for qt, kt, vt in zip(q_phi, k_phi, v):
        state += np.outer(kt, vt)
        normalizer += kt
        numerator = qt @ state
        denominator = qt @ normalizer
        outputs.append(numerator / max(denominator, eps))
    return np.array(outputs)

The chapter tests generate seeded $\vq,\vk,\vv$ arrays and assert that the parallel lower triangle and recurrent state produce the same outputs to float64 precision. That test is the minimum correctness check: if the order of the outer product or the normalizer is wrong, the two forms disagree immediately.

The feature map is doing the approximation work. Softmax attention compares every query with every key and then normalizes those exact scores. Linear attention chooses a representation in which the score already looks like an inner product of transformed vectors. Positivity matters because the denominator should behave like a sum of weights, not like a cancellation between positive and negative terms. Different papers choose different maps, but the implementation pattern is the same once $\vphi$ has been chosen.

28.2 Constant-state decoding

During autoregressive decoding, softmax attention caches every previous $\vk_s$ and $\vv_s$ . A linear-attention layer instead caches only $\mS$ and $\vc$ . For fixed feature and value widths, each new token updates the same two arrays and emits one output. The cache size is independent of the prefix length, so decoding is $O(1)$ memory per token for that layer.

This is not free. The state is a compressed summary: once two different prefixes produce the same $\mS$ and $\vc$ , the layer cannot distinguish them later. Full attention keeps all past keys and values and can revisit a rare token exactly. Linear layers win when the model can use a lossy recurrent memory for much of the stack and reserve exact attention for layers that need retrieval.

This distinction also explains why training can still be parallel. During training, the whole sequence is known, so a scan or a masked matrix computation can form all prefix states at once. During decoding, tokens arrive one at a time, and the recurrent form becomes the natural implementation. A useful mental model is therefore not \"linear attention is an RNN instead of attention,\" but \"the attention kernel was chosen so that its causal prefix sums are RNN states.\" The tests exercise both views on the same tensors.

28.3 Delta-rule memories

The delta rule makes the recurrent state an error-correcting associative memory. Let $\mS$ map key features to values. Before writing token $t$ , predict its value from its key, compute the error, and write only the error:

\mS_t = \mS_{t-1} + \beta_t\,\vphi(\vk_t) (\vv_t - \vphi(\vk_t)^\T\mS_{t-1})^\T .\tag{28.4}

If the memory already predicts the value, the update is small. If it is wrong, the update corrects that key. Gated DeltaNet adds a forget gate before the correction:

\mS_t = \gamma_t\mS_{t-1} + \beta_t\,\vphi(\vk_t) (\vv_t - \vphi(\vk_t)^\T\gamma_t\mS_{t-1})^\T .\tag{28.5}

The gate lets the model decide how fast old associations decay, while the delta term still writes prediction error. Gated Delta Networks use this idea to improve Mamba-style sequence models [yang2024gated].

The delta rule differs from the plain linear-attention write in one important way. A plain write adds $\vphi(\vk_t)\vv_t^\T$ every time, even if the state already maps that key to the right value. The delta update first asks what the state would return, then writes the residual. That makes repeated evidence for the same association converge instead of growing without bound. The gated version adds controlled forgetting, so new evidence can overwrite stale associations rather than merely adding to them.

Listing 28.2 Delta and gated-delta updates

def delta_step(state, key_feature, value, beta=1.0):
    """Error-correcting associative-memory update."""
    prediction = key_feature @ state
    error = value - prediction
    next_state = state + beta * np.outer(key_feature, error)
    return next_state, error


def gated_delta_step(state, key_feature, value, beta=1.0, gate=0.95):
    """Forget part of the old state, then write the current prediction error."""
    decayed = gate * state
    prediction = key_feature @ decayed
    error = value - prediction
    next_state = decayed + beta * np.outer(key_feature, error)
    return next_state, error

28.4 State-space models and hybrids

A linear state-space model keeps a hidden state $\vh_t$ . Discretizing a continuous system gives a recurrence of the form

\vh_t = \bar{\mA}_t\vh_{t-1} + \bar{\mB}_t\vx_t, \qquad \vy_t = \mC_t\vh_t .\tag{28.6}

Mamba’s selective SSM makes the discretization and projections input-dependent, so the model can choose what to remember or forget as a function of the current token [gu2023mamba]. The tiny NumPy version below uses a diagonal state: $\bar{\mA}_t$ is an exponential decay and $\bar{\mB}_t,\mC_t$ vary with the input.

Listing 28.3 A tiny selective state-space recurrence

def selective_state_space(x, a, b, c, delta):
    """A tiny diagonal selective SSM recurrence."""
    state = np.zeros_like(a, dtype=x.dtype)
    outputs = []
    for xt, bt, ct, dt in zip(x, b, c, delta):
        decay = np.exp(dt * a)
        state = decay * state + dt * bt * xt
        outputs.append(ct @ state)
    return np.array(outputs)

Hybrid LLMs mix these mechanisms rather than declaring one winner. Jamba is a hybrid Transformer-Mamba model [lieber2024jamba]; MiniMax-01 uses Lightning Attention in a foundation model stack [minimax2025minimax01]; Kimi Linear proposes an efficient linear-attention architecture [zhang2025kimi]. The common recipe is to interleave cheap recurrent layers with occasional full-attention layers, keeping exact retrieval paths while reducing average cache and attention cost.

The state-space view is broader than linear attention, but the engineering motivation is similar: replace an ever-growing table of past activations with a state updated by a recurrence. The selective parameters are what keep that recurrence from being a fixed filter applied to all tokens. A token can ask for a slow decay, a fast decay, or a different input projection. That is why Mamba-like layers are usually discussed as content-dependent sequence models rather than as ordinary convolutions.

In practice

Use linear attention when long prefixes make exact attention too expensive and the task can benefit from a compressed recurrent memory. Use full attention when exact copying, retrieval, or cross-token comparison matters. Modern long-context stacks usually combine them: recurrent or linear layers carry most tokens cheaply, while full-attention layers refresh global access. Mamba and DeltaNet-style layers should be read as sequence-memory layers, not as drop-in exact softmax replacements [gu2023mamba][yang2024gated].

Key equations

K(\vq,\vk) = \vphi(\vq)^\T\vphi(\vk)

\mS_t = \mS_{t-1} + \vphi(\vk_t)\vv_t^\T, \quad \vc_t = \vc_{t-1} + \vphi(\vk_t)

\vy_t = \frac{\vphi(\vq_t)^\T\mS_t} {\vphi(\vq_t)^\T\vc_t}

\mS_t = \mS_{t-1} + \beta_t\vphi(\vk_t) (\vv_t - \vphi(\vk_t)^\T\mS_{t-1})^\T

\vh_t = \bar{\mA}_t\vh_{t-1} + \bar{\mB}_t\vx_t, \quad \vy_t = \mC_t\vh_t

28.5 Teach it

The one-sentence version. Linear attention factorizes the attention score so the whole prefix can be summarized by a recurrent key-value state.

An analogy. Softmax attention keeps every note you ever wrote; linear attention keeps a running ledger. The ledger is compact and fast to update, but it cannot recover a note that was summarized away.

At the board.

Replace $e^{\vq^\T\vk}$ by $\vphi(\vq)^\T\vphi(\vk)$ .
Pull $\vphi(\vq_t)$ outside the prefix sum and define $\mS_t$ and $\vc_t$ .
Show the one-token decoding update: add one outer product, then read with the query.
Contrast a plain write with the delta rule: write the prediction error, not the whole value.

Misconceptions to address. Linear attention is not exact softmax attention. Constant-state decoding saves memory per layer, not necessarily all model memory. Mamba is a selective SSM, not just attention with a different kernel.

Check for understanding. What information is lost when a layer keeps only $\mS$ and $\vc$ instead of all past keys and values?

28.6 Exercises

Exercise 28.1 ★ Kernel replacement

State the condition a feature map must satisfy for (28.2) to be a normalized weighted average, and explain why positivity matters.

Solution in Appendix D

Exercise 28.2 ★★ Derive the recurrent form

Starting from (28.2), derive (28.3) and the readout $\vphi(\vq_t)^\T\mS_t / (\vphi(\vq_t)^\T\vc_t)$ .

Solution in Appendix D

Exercise 28.3 ★★ Error-correcting write

For one key feature $\vk = (1,0)$ and value $\vv$ , show what repeated delta-rule updates with $\beta = 1/2$ do to the prediction error.

Solution in Appendix D

Exercise 28.4 ★★★ Decode one token

Implement a one-token decoder update that receives $\vq_t,\vk_t,\vv_t,\mS,\vc$ and returns the output and updated state. Compare a full prefix processed this way with recurrent_linear_attention.

Solution in Appendix D

References

[gu2023mamba] A. Gu and T. Dao. Mamba: Linear-Time Sequence Modeling with Selective State Spaces. 2023. arXiv:2312.00752
[katharopoulos2020transformers] A. Katharopoulos et al. Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention. 2020. arXiv:2006.16236
[lieber2024jamba] O. Lieber et al. Jamba: A Hybrid Transformer-Mamba Language Model. 2024. arXiv:2403.19887
[minimax2025minimax01] MiniMax et al. MiniMax-01: Scaling Foundation Models with Lightning Attention. 2025. arXiv:2501.08313
[yang2024gated] S. Yang, J. Kautz, and A. Hatamizadeh. Gated Delta Networks: Improving Mamba2 with Delta Rule. 2024. arXiv:2412.06464
[zhang2025kimi] Y. Zhang et al. Kimi Linear: An Expressive, Efficient Attention Architecture. 2025. arXiv:2510.26692