49. Sequential vs. Parallel Scan for a Linear Recurrence

Problem

State-space and linear-recurrent sequence models (S4, Mamba, and gated linear-recurrence variants like GateLoop) are all built on the same core primitive: a per-channel linear recurrence h_t = a_t * h_(t-1) + b_t. The textbook way to compute it is a Python for-loop over time — correct, but sequential, so it can't use a GPU's parallelism across the time dimension. The alternative is to unroll the recurrence into a closed form using cumulative products, which lets every timestep be computed in parallel. This is the "associative scan" formulation underlying why SSMs can be trained on full sequences roughly as fast as a Transformer, despite being recurrent at inference time.

Given the sequential reference implementation, implement the parallel (no time-loop) version and confirm it matches numerically.

Source: src/49_sequential_vs_parallel_scan.py

def linear_recurrence_parallel(
    a: torch.Tensor,
    b: torch.Tensor,
    h0: torch.Tensor,
) -> torch.Tensor:
    ...

Examples:
>>> a = torch.rand(2, 5, 3) * 0.4 + 0.5  # gate values in [0.5, 0.9]
>>> b = torch.randn(2, 5, 3)
>>> h0 = torch.randn(2, 3)
>>> h = linear_recurrence_parallel(a, b, h0)
>>> h.shape
torch.Size([2, 5, 3])

Step-by-Step Approach

1. Unroll the recurrence by substitution: h_t = A_t * h0 + Σ_(j=1..t) (A_t / A_j) * b_j, where A_t = ∏_(i=1..t) a_i is the cumulative product of gates up to time t.
2. Compute A in one call with torch.cumprod(a, dim=1) — this replaces the sequential state-carrying loop entirely.
3. Compute the h0 contribution: A * h0.unsqueeze(1), broadcasting h0 across the time dimension.
4. Factor A_t out of the sum over j: Σ_j (A_t/A_j)*b_j = A_t * Σ_j (b_j/A_j), so the inner sum becomes torch.cumsum(b / A, dim=1).
5. Multiply that cumulative sum back by A to get the b-driven contribution, and add it to the h0 contribution.
6. Compare against the given sequential implementation with a moderate tolerance (e.g. atol=1e-3) — the two should agree closely for reasonable sequence lengths and gate ranges, but not bit-for-bit, because of the division-based reformulation.

The key insight, and the actual engineering tradeoff being tested, is that dividing by A_j is numerically fragile: since gates a_t are typically in (0,1), A_j shrinks toward zero as j grows, so b_j/A_j can become very large before the outer multiply by A_t rescales it back down — for long sequences or gates near 0 this loses precision or overflows, which is exactly why production SSM kernels use chunked/blocked scans (splitting the sequence into blocks and combining block-level results associatively) instead of one global cumulative-product ratio.

Reference solution

def linear_recurrence_parallel(a, b, h0):
    # h_t = A_t*h0 + sum_j (A_t/A_j)*b_j,  A_t = prod_{i<=t} a_i
    A = torch.cumprod(a, dim=1)  # (batch, seq_len, dim)

    h0_contrib = A * h0.unsqueeze(1)

    # factor A_t out of the sum: A_t * cumsum_j(b_j / A_j)
    # dividing by A_j (shrinking toward 0) is the numerically fragile step
    # for long sequences -- the reason real kernels use chunked scans
    b_scaled = b / A
    cumsum_b_scaled = torch.cumsum(b_scaled, dim=1)
    b_contrib = A * cumsum_b_scaled

    return h0_contrib + b_contrib

Key Functions & Tricks

• torch.cumprod(a, dim=1) — computes every prefix product along the time axis in one vectorized call, replacing the sequential state carry.
• torch.cumsum(..., dim=1) — the additive counterpart used on the rescaled b/A term.
• tensor.unsqueeze(1) — broadcasts h0 (batch, dim) up to (batch, 1, dim) so it aligns against the (batch, seq_len, dim) cumulative product.
• Factoring a shared term out of a sum before vectorizing — the general algebraic move that turns a recurrence into a closed-form scan.
• Restricting gate values to a bounded range (e.g. via a sigmoid in a real model) — the standard mitigation that keeps A_t from decaying to numerically-unusable magnitudes too fast.

How to Recognize This Pattern