Parallel Scan

An RNN trains slowly because $h_t$ needs $h_{t-1}$. That dependency looks like a law of nature. For a linear recurrence it is not.

with $a_t$ and $b_t$ computed from the input alone — not from $h$. Every element of the state evolves independently, and the update is affine.

Composition

Two consecutive steps compose into a single step of the same shape.

Substituting confirms it: applying $(a_1,b_1)$ then $(a_2,b_2)$ gives $a_2(a_1 h + b_1) + b_2$, which is the composed pair applied once. And the operator is associative, because function composition always is.

That is the whole trick. Associativity means the operations can be grouped however you like, and a different grouping is a different schedule — a balanced binary tree instead of a chain.

Blelloch’s two sweeps

Up-sweep. Pair adjacent elements and compose, giving $n/2$ pairs each representing two steps. Repeat. After $\log n$ rounds the root holds the composition of the entire sequence. This is a reduction tree.

Down-sweep. Walk back down, pushing each node’s accumulated prefix into its children, so every leaf ends up holding the composition of everything to its left. Another $\log n$ rounds.

For $n = 4096$: 24 dependent steps rather than 4096. The total arithmetic is about twice the sequential version, which is an excellent trade on hardware with thousands of idle cores.

Why every recurrent model since 2021 is linear

The scan requires eq. 1’s form. A $\tanh$ inside the recurrence destroys associativity — $\tanh(a_2\tanh(a_1 h + b_1) + b_2)$ does not collapse into one affine step — and no amount of engineering recovers it.

So the design question inverted. Rather than asking what recurrence is most expressive and accepting slow training, the field asked what can be trained in parallel and found expressiveness elsewhere: nonlinearities between layers, and input-dependent $a_t$ within them. Selective SSMs, RWKV and gated linear attention are all eq. 1 with different parameterisations of $a$ and $b$, and they are all trained by this algorithm.

import torch
from torch import Tensor


def scan(a: Tensor, b: Tensor) -> Tensor:
    """h_t = a_t * h_{t-1} + b_t, over dim 1. a, b: [B, T, D]."""
    a, b = a.clone(), b.clone()
    t, shift = a.size(1), 1

    while shift < t:
        # Compose each element with the one `shift` places behind it. After
        # round k every element holds the composition of 2^k steps.
        a_prev = torch.nn.functional.pad(a[:, :-shift], (0, 0, shift, 0), value=1.0)
        b_prev = torch.nn.functional.pad(b[:, :-shift], (0, 0, shift, 0), value=0.0)

        b = a * b_prev + b
        a = a * a_prev
        shift *= 2

    return b


def scan_reference(a: Tensor, b: Tensor) -> Tensor:
    """The loop it replaces. Same answer, 4096 launches instead of 12."""
    h = torch.zeros_like(b[:, 0])
    return torch.stack([(h := a[:, t] * h + b[:, t]) for t in range(b.size(1))], 1)

Hillis–Steele is $O(n\log n)$ work against Blelloch’s $O(n)$, and is nonetheless what most implementations use: it has half the depth, no down-sweep, and a memory access pattern the hardware likes. Asymptotically worse and reliably faster at the sizes that occur.

The real kernels do not run the scan on the full sequence in one go. They tile it: load a chunk into SRAM, scan it there, carry one composed pair across the chunk boundary, and never write intermediate states to HBM. That is the same IO-aware argument as FlashAttention, and it is why Mamba’s paper spends as much space on the memory hierarchy as on the model — the algorithm was settled in 1980, and what was new was making it bandwidth-efficient.