---
kind: concept
slug: parallel-scan
title: Parallel scan (prefix scan)
summary: >-
  A parallel scan evaluates a recurrence over a whole sequence at once by combining steps with an associative operator, so the result matches the step-by-step loop. Two variants: S5, a diagonal state space layer run in the time domain at S4's cost, and a scan over LQR value functions in GPU model predictive control for legged robots. Both are author-reported.
tags: [sequence-modeling, parallel-computation, state-space-models, optimal-control]
papers: ["2208.04933", "2506.07823"]
related: [state-space-model, recurrent-neural-network, model-predictive-control, linear-quadratic-regulator]
status: draft
updated_at: 2026-10-25
---

## 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 [arxiv:2208.04933 §2.2] [arxiv:2208.04933 app. H].
Two papers apply this to different objects. S5 scans a linear state recurrence
over time [arxiv:2208.04933 §3.2]. The legged-robot MPC paper scans the
value functions of an LQR-type backward pass over the planning horizon
[arxiv: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 [arxiv: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
[arxiv:2208.04933 §3.4]. Card: [Simplified State Space Layers for Sequence Modeling](/papers/2208.04933).

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] [arxiv: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 [arxiv:2506.07823 §2.4] [arxiv: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 [arxiv:2506.07823 §2.4]. The forward
closed-loop rollout is also written as a scan [arxiv:2506.07823 eq. 16]
[arxiv:2506.07823 eq. 18]. Card: [Primal-Dual iLQR for GPU-Accelerated Learning and Control in Legged Robots](/papers/2506.07823).

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 [arxiv:2506.07823 §2.5] [arxiv: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

| Situation | Use | Why |
|---|---|---|
| Long sequence, linear recurrence, training should parallelise over time | Parallel scan over a diagonal SSM (S5) | Reported to run at S4's cost without S4's convolution kernel [arxiv:2208.04933 §3.4] [arxiv:2208.04933 §2.3] |
| Need the S4 HiPPO-LegS initialisation as-is | Not this scan directly | HiPPO-LegS cannot be diagonalised stably, which blocks a direct parallel-scan version [arxiv:2208.04933 §2.3] [arxiv:2208.04933 §4.2]. S5 uses a diagonalised HiPPO-N instead [arxiv:2208.04933 app. B.1.1] |
| Memory is tight | Budget for the stored states | The scan keeps states of size O(P·L) [arxiv:2208.04933 app. C.1]. The paper reports memory of 0.7 to 1.2 times S4D [arxiv:2208.04933 app. C.2] |
| Sequence sampled at irregular times | Per-step timescale Δ in the scan | The paper reports this for its pendulum task [arxiv:2208.04933 §3.3] [arxiv:2208.04933 §6.3]. The pendulum code is not in the official repository, so this cannot be reproduced from that code |
| Implementation in JAX | Official S5 code | The paper's listing is in JAX [arxiv:2208.04933 app. A]. The official repository has the scan operator [code: s5/ssm.py#L46-L57] |
| Optimal control over a long horizon, GPU available, re-planning at a fixed rate | Scan over LQR value functions (primal-dual iLQR) | Reported to solve faster than acados and crocoddyl at long horizons, with 16 centralised SRBD robots under 25 ms [arxiv:2506.07823 §4.3] [arxiv:2506.07823 §4.2] |
| Short horizon, small problem, or CPU only | A CPU solver | Crocoddyl is faster for whole-body runs below 80 nodes, and the GPU pays a host-device overhead of about 1 ms [arxiv:2506.07823 §4.2] [arxiv:2506.07823 §4.1]. acados is competitive for small batches [arxiv:2506.07823 §4.4] |

## 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.

```pseudocode
# 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] [arxiv: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
[arxiv:2506.07823 §2.4].

## Contradictions

Between the two papers, none on the scan itself: S5 reports training speed
against S4D [arxiv:2208.04933 app. C.2], and the MPC paper reports solve time
against acados and crocoddyl [arxiv: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) [arxiv:2506.07823 §2.4], and the card's pitfalls
record O(log² n log N + log² m) for the same scan [arxiv:2506.07823 §2.4].
Check the paper before quoting either form.

## Papers

- [Simplified State Space Layers for Sequence Modeling](/papers/2208.04933): 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](/papers/2506.07823): runs the LQR value-function backward pass and the forward rollout as parallel scans inside a GPU MPC solver.