16. Parallel (Associative) Scan for a Linear Recurrence

Problem

A linear recurrence h_t = a_t * h_{t-1} + b_t looks inherently sequential — h_t needs h_{t-1}, which needs h_{t-2}, and so on — so a naive implementation is an O(L) Python loop with L sequential steps that cannot be parallelized across time. But the recurrence is really a composition of affine maps, and affine-map composition is associative, which means the whole sequence can instead be reduced with a Blelloch/Hillis-Steele style parallel scan: O(log L) sequential rounds, each doing O(L) work that is fully vectorizable across the time axis.

Treat each timestep as an affine map h -> a*h + b, represented as a pair (a, b). Composing "apply map1, then map2" gives (a1, b1) . (a2, b2) = (a1 * a2, a2 * b1 + b2). Running an inclusive (Hillis-Steele) scan over this operator, with identity element (a=1, b=0) used whenever a combine would reach before the start of the sequence, produces at position t the composed map from h_{-1}=0 to h_t — whose b-component is exactly h_t itself.

Source: src/16_parallel_associative_scan.py

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

>>> a = torch.sigmoid(torch.randn(2, 16, 4))
>>> b = torch.randn(2, 16, 4)
>>> parallel_scan_linear_recurrence(a, b).shape
torch.Size([2, 16, 4])

Step-by-Step Approach

1. Start with A, B = a.clone(), b.clone(), and a stride = 1 that doubles each round: while stride < seq_len: ... stride *= 2.
2. Each round, build a "look-back" pair shifted by stride positions: A_prev/B_prev, formed by concatenating an identity block (a=1, b=0) of length stride in front of A[:, :-stride, :]/B[:, :-stride, :].
3. Combine every position with its look-back partner using the affine-composition formula: A_new = A_prev * A, B_new = A * B_prev + B — applied to every position at once, not just those >= stride.
4. Notice why no separate boundary case is needed: for positions t < stride, the look-back pair is exactly the identity (1, 0), and composing with identity is a no-op by construction, so those positions come out unchanged automatically.
5. Repeat for ceil(log2(seq_len)) rounds, doubling stride each time, until stride >= seq_len — at that point every position has been combined with the full prefix of maps before it.
6. Return B: after the loop, B[t] holds the b-component of the composed map from h_{-1}=0 to h_t, which is precisely h_t.

The key insight is the identity-element trick: padding out-of-range look-backs with (a=1, b=0) instead of special-casing them means one uniform vectorized combine expression handles every position, including the ones too close to the start of the sequence to have a full-stride partner yet.

Reference solution

import torch


def parallel_scan_linear_recurrence(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    batch, seq_len, hidden = a.shape
    A, B = a.clone(), b.clone()
    stride = 1
    # ceil(log2(seq_len)) rounds total -- each round is O(seq_len) fully
    # vectorized work, not a per-timestep Python loop
    while stride < seq_len:
        # identity element (a=1, b=0): composing with it is a no-op, so
        # padding the "look-back" window with identity for the first
        # `stride` positions makes them correctly pass through unchanged
        # without a separate branch
        ones = torch.ones(batch, stride, hidden, dtype=A.dtype)
        zeros = torch.zeros(batch, stride, hidden, dtype=B.dtype)
        A_prev = torch.cat([ones, A[:, :-stride, :]], dim=1)
        B_prev = torch.cat([zeros, B[:, :-stride, :]], dim=1)
        # combine (A_prev, B_prev) "earlier" with (A, B) "later":
        # (a1,b1).(a2,b2) = (a1*a2, a2*b1 + b2)
        A, B = A_prev * A, A * B_prev + B
        stride *= 2
    # B now holds, at every t, the composed map's b-term starting from
    # h_{-1} = 0 -- which is exactly h_t
    return B

Key Functions & Tricks

• torch.cat([identity_block, A[:, :-stride, :]], dim=1) — builds the shifted "look-back" tensor for the whole batch/time axis in one call, padded with the operator's identity element.
• while stride < seq_len: ... stride *= 2 — the doubling-stride loop that gives O(log L) rounds; note it works correctly for any seq_len, not only powers of two.
• Elementwise */+ implementing the affine-composition formula — the whole "parallel scan" is just this associative combine applied at doubling distances.
• Identity element (a=1, b=0) for an affine map — the same role the number 0 plays for addition or 1 plays for multiplication in a standard prefix-sum/prefix-product scan.

How to Recognize This Pattern