Parallel scan (prefix scan)

Overview

A recurrence computes each state from the one before it, which looks serial. When the step can be written as an associative combination, the steps can be grouped in any order, so the whole sequence is computed as a tree of combines. The invariant is that the operator is associative, so the tree gives the same states as the step-by-step loop 2208.04933 §2.2 2208.04933 app. H. Two papers apply this to different objects. S5 scans a linear state recurrence over time 2208.04933 §3.2. The legged-robot MPC paper scans the value functions of an LQR-type backward pass over the planning horizon 2506.07823 §2.4.

Variants

Scan over a diagonal linear recurrence (S5)

The state matrix is diagonal, so every state dimension is an independent scalar recurrence and the operator acts elementwise on P state channels. The decay term Λ̄ is the same for all steps in the layer, and the drive term B̄u_k is computed for every step before the scan 2208.04933 §3.2. The paper reports that this matches S4's training cost, with the same asymptotic order when the state size is comparable to the feature size 2208.04933 §3.4. Card: Simplified State Space Layers for Sequence Modeling.

The paper also reports measured training speed of 0.7 to 1.9 times S4D across its tasks, and memory of 0.7 to 1.2 times S4D [arxiv:2208.04933 app. C.2]. These numbers are author-reported.

Scan over LQR value functions (primal-dual iLQR)

The backward Riccati recursion of an LQR-type solver runs one node at a time from the end of the horizon. The paper observes that the cost-to-go between any two nodes is again a quadratic of the same form, and that combining two such pieces with a min-plus-style operator is associative [arxiv:2506.07823 §2.4] 2506.07823 eq. 9. A scan over adjacent-node elements then gives the value-function terms P_i and p_i for every node, which the paper states has O(log N) depth 2506.07823 §2.4 2506.07823 eq. 8. The operator contains matrix inverses of the form (I + P̃C̃)⁻¹ [arxiv:2506.07823 eq. 11], so the depth claim assumes a parallel matrix inversion method, which the paper cites rather than derives 2506.07823 §2.4. The forward closed-loop rollout is also written as a scan 2506.07823 eq. 16 2506.07823 eq. 18. Card: Primal-Dual iLQR for GPU-Accelerated Learning and Control in Legged Robots.

The paper's timings combine this scan with the rest of one SQP iteration, a ten-step parallel line search and the simulator, so they are not a measure of the scan alone 2506.07823 §2.5 2506.07823 §4.4.

Not covered: nonlinear recurrences

The sources here are about linear or quadratic recurrences. The operator used by S5 composes affine maps x ↦ a·x + b, and the LQR scan composes quadratic cost-to-go functions. Neither paper says whether a parallel scan applies to a general nonlinear recurrence, so this page makes no claim about it.

When to use which

Implementation sketch

The common core is the associative combine and a scan over it. For the S5 variant, pairs are (a, b), meaning x ↦ a·x + b, applied elementwise over the P channels. In S5, a_k is the same Λ̄ at every step.

# Parallel scan over a linear recurrence  [arxiv:2208.04933 §2.2] [app. H]
# Inputs: a (L, P) decay per step; b (L, P) drive per step
# Output: x (L, P), the state after each step with x_0 = 0
def scan_linear_recurrence(a, b):
    # Each step is the pair (a_k, b_k), meaning x -> a_k * x + b_k  [app. H]
    elems = zip(a, b)
    # Combine pairs with the associative operator, as a tree  [app. H]
    prefixes = associative_scan(combine, elems)     # (L,) pairs
    # With x_0 = 0, the state after step k is the b part of its prefix
    return [b_k for (_, b_k) in prefixes]           # (L, P)


# Compose the earlier pair with the later one  [app. H] [code: s5/ssm.py#L46-L57]
def combine(first, second):
    a1, b1 = first
    a2, b2 = second
    # Applying (a1, b1) then (a2, b2) gives x -> a2*a1*x + (a2*b1 + b2)
    return a2 * a1, a2 * b1 + b2

The LQR variant uses a different element and operator [arxiv:2506.07823 eq. 11] 2506.07823 eq. 12, and its scan runs backward over the horizon. Several symbols in eq. 11 look inconsistent in the paper, so implement it from the official code or derive the rule and test it on a small problem 2506.07823 §2.4.

Contradictions

Between the two papers, none on the scan itself: S5 reports training speed against S4D 2208.04933 app. C.2, and the MPC paper reports solve time against acados and crocoddyl 2506.07823 §4. The two sets of numbers measure different things and cannot be compared.

Within the MPC paper, the rollout bound is written two ways. Section 2.4 gives O(log n log N + log m) 2506.07823 §2.4, and the card's pitfalls record O(log² n log N + log² m) for the same scan 2506.07823 §2.4. Check the paper before quoting either form.

Papers

• Simplified State Space Layers for Sequence Modeling: S5 runs a diagonal, multi-input multi-output state space layer as a parallel scan over time.
• Primal-Dual iLQR for GPU-Accelerated Learning and Control in Legged Robots: runs the LQR value-function backward pass and the forward rollout as parallel scans inside a GPU MPC solver.