Associative Scan

30 July 2025 · Download PDF

Recently a plethora of methods within machine learning have made use of state-space models (SSMs) to model sequences [1], [2], [3], [4]. While the early S4 model is computed as a convolution [1], a big reason for the success of later models such as S5 [2] and Mamba [3] is their usage of the associative scan operation to parallelize the state computation on the GPU. In my opinion this is a neat trick that was poorly explained in the original papers (although Appendix B in [4] does explain it for the scalar case), and I will therefore explain it in more detail in this blog post. The aim is to show that composing the transitions of a (linear) state-space model

is associative¹¹ Associativity simply means that repeated applications of the operator (denoted by $•$) can be grouped arbitrarily, i.e., that $𝑎•\left(𝑏•𝑐\right)=\left(𝑎•𝑏\right)•𝑐$. The order of the operands must still be kept, as the operator need not be commutative. Examples of associative operators are addition and multiplication, including matrix multiplication.. In short this means that the states of a (linear) state-space model up until time step $𝑇$ can be computed in parallel using an associative scan operation. For details on the associative operation open the example box below.

>>> lax.associative_scan(jnp.add, jnp.arange(0, 4))
Array([0, 1, 3, 6], dtype=int32)

julia> accumulate(+, 0:3)
4-element Vector{Int64}:
  0
  1
  3
  6

While the above computation is trivial and easily computed by looping from the first to last element in the vector, the main idea of the associative scan is that the operator can be applied pairwise in parallel. That means if enough processors are available the computation can be done in $\mathrm{𝒪︀}\left(log𝑇\right)$ time rather than the trivial implementation of $\mathrm{𝒪︀}\left(𝑇\right)$.

The first step in showing the associativity of the state-space model is to define the transition of the state-space models using matrix multiplication (which is associative) by embedding the transition into a larger matrix ${𝐬}_{𝑘}$ as follows

Using the definition of ${𝐬}_{𝑘}$ the state transition from state $𝑘-1$ to state $𝑘$ can be computed using matrix multiplication as

Using this we can compute the $𝑖$th state of the state-space model as

Given that the cumulative product can be computed using the associative scan operator the full dynamics can be computed as

While the above works, it can be simplified slightly. As we are really only interested in the top row of blocks (the accumulated product ${𝐩}_{𝑖}$ has ${𝐡}_{𝑖}$ in its top-right block), we can instead define elements by just the top row, i.e., ${𝐬}_{𝑘}=\left[\begin{array}{cc}{𝐀}_{𝑘}& {𝐁}_{𝑘}{𝐱}_{𝑘}\end{array}\right]$, and then define the operator (denoted by $•$) by how the top row propagates. For general pairs this reads

which is the top block row of the product $\left[\begin{array}{cc}{𝐀}^{″}& {𝐛}^{″}\\ \mathrm{𝟎}& 1\end{array}\right]\left[\begin{array}{cc}{𝐀}^{\prime }& {𝐛}^{\prime }\\ \mathrm{𝟎}& 1\end{array}\right]$. The operator is therefore associative, because matrix multiplication is. The SSM states can then be computed as ${𝐩}_{𝑖}={𝐬}_{𝑖}•{𝐩}_{𝑖-1}$ with ${𝐩}_0={𝐬}_0=\left[\begin{array}{cc}\mathrm{𝟎}& {𝐡}_0\end{array}\right]$, and ${𝐡}_{𝑖}$ is the right block of ${𝐩}_{𝑖}$.

As a final remark note that while the associative scan is parallelizable it performs matrix-matrix products of the form ${𝐀}_{𝑘}{𝐀}_{𝑘-1}$ which will be computationally prohibitive unless ${𝐀}_{𝑘}$ has structure (e.g., diagonal or low-rank). This is one of the reasons why e.g. Mamba-2 utilizes a scaled identity as its ${𝐀}_{𝑘}$ [4].

@misc{paltorp2025associativescan,
  author = {Paltorp, Mikkel},
  title = {{Associative Scan}},
  year = {2025},
  month = jul,
  url = {https://mipals.github.io/pdf/2025/07-30-associative-scan/associative-scan.pdf},
  note = {Blog post}
}