S5: Simplified State Space Layers for Sequence Modeling

Can a recurrent layer process a whole sequence in parallel?

A recurrence seems to force a choice: either process one observation at a time, or replace the recurrence with an operation that sees the entire sequence. S5 shows why that choice is too restrictive. A linear state update is an affine map, and affine maps can be composed associatively. A parallel scan can therefore compute every prefix state without changing the recurrence into an FFT convolution.

The challenge is making each composition cheap. Dense state matrices would make a parallel scan expensive. S5 combines a shared multi-input, multi-output state-space model with diagonal state dynamics, so composition uses elementwise products. Dense input and output projections still let channels communicate. The simplification is a coordinated choice of state size, parameterization and execution strategy, rather than a claim that every recurrent network can be scanned efficiently. Paper v3, §§2–3

This review follows the ICLR 2023 paper by Jimmy T. H. Smith, Andrew Warrington and Scott W. Linderman. The first preprint appeared in August 2022; the equations and results below use the March 2023 revision, v3. The header date belongs to this write-up. The official repository identifies the venue and supplies the implementation at the pinned revision used here.

The six explorers are authored diagrams and analytical calculations, not trained model outputs. Published benchmark results appear separately. Start with the architectural change that makes a small diagonal recurrence useful.

From independent SISO memories to one MIMO memory

An S4 layer receives an input vector with H features at each time. It applies one single-input, single-output (SISO) state-space model to each feature. If every SISO model has N state coordinates, the combined state has HN coordinates. Within that linear recurrent part, input channel 1 cannot directly write to channel 2's state. A nonlinearity and a position-wise linear mixing layer combine the features afterward.

S5 instead uses one multi-input, multi-output (MIMO) system. A dense B matrix writes all H input features into P state coordinates. A dense C matrix reads those shared coordinates into H output features. P is an independent modeling choice and can be far smaller than HN. The following diagram compares the linear SSMs, before the surrounding layer's nonlinear processing. §2.3, §3.1 and Figure 2

The continuous system is:

Here u,y\in\mathbb R^H, x\in\mathbb R^P, B\in\mathbb R^{P\times H} and C\in\mathbb R^{H\times P}. The general state-space equation permits dense D, but S5 restricts D to a diagonal feedthrough, stored as H real values. That path passes the current input directly to the output; it is distinct from the residual connection around a whole layer.

MIMO moves channel mixing inside the state-space operator, so S5 does not require S4's additional position-wise mixing layer for that purpose. The SSM itself remains linear. Nonlinear activations and the surrounding normalization, dropout and residual machinery make a stack of S5 layers a nonlinear sequence model. The official layer wrapper separates that wrapper from the SSM in ssm.py.

A shared state creates a capacity tradeoff. Fewer coordinates can reduce recurrent work and memory, but all features must now share the same memory bank. S5 is not an exact compression of an arbitrary collection of independently learned S4 systems. To see what is preserved, we need to distinguish a change of coordinates from a change of model.

Diagonal dynamics, dense communication

Suppose a diagonalizable state matrix satisfies A=V\Lambda V^{-1}. Describe the same state in eigen-coordinates, \tilde x=V^{-1}x. Substitution gives:

The recurrent transition now scales each mode independently. The input and output projections remain dense. An input feature can excite many modes, and an output feature can read many modes. Diagonal state dynamics do not imply independent input/output channels. §3.2, equations 4–5

Changing the angle in the figure changes the off-diagonal entries of A in the original coordinates, while the reconstructed derivative stays equal. This identity requires transforming B and C as well. Simply erasing the off-diagonal entries of A while retaining the old projections would change the dynamics.

S5 does not repeatedly eigendecompose a dense learned A during training. It initializes an eigenbasis from a structured matrix, transforms the projections, and then learns diagonal complex eigenvalues and dense complex projections directly. Complex modes conveniently represent decay and oscillation. When a real system has conjugate mode pairs, storing one partner is sufficient: the other partner's output is its complex conjugate. §3.2 and Appendix B.1

This makes the recurrent algebra cheap, but the data are discrete. Before scanning the sequence, the continuous parameters must become an actual per-sample update.

Discretization controls both retention and input gain

For an input held constant over an interval Δ, zero-order hold (ZOH) integrates the continuous equation exactly. For diagonal mode p:

S5 learns a positive interval for each mode, parameterized through its logarithm. A mode's real eigenvalue part controls decay; its imaginary part controls oscillation. The interval determines how much of each occurs in one discrete step. The input gain changes with that same interval. Updating only the transition while leaving the input gain fixed would no longer implement the same held-input integration. §3.2, equation 6

For \operatorname{Re}\lambda_p<0 and positive Δ, the discrete multiplier has magnitude below one. This is a statement about the isolated mode's transition, not a bound on the entire network's outputs or gradients. Dense projections, nonlinearities and layer stacking still matter. The official implementation offers optional negative-real-part clipping through clip_eigs; stability is not an unconditional consequence of a class named S5. It also supports bilinear discretization as an alternative to ZOH. Pinned discretizers and S5SSM

The apparent singularity at a zero eigenvalue is removable: the input gain tends to Δ. Our reference uses a short series near zero so it can check that boundary explicitly. That numerical guard is an authored addition, rather than a claim about the exact bytes of the original JAX implementation.

We now have one affine update per observation. The next step changes how those updates are scheduled, without changing their meaning.

Associative scans preserve the recurrence

Write an update as f_k(x)=a_k\odot x+b_k, where b_k=\bar B_k u_k. Applying an earlier update i and then a later update j yields:

Store a map as the pair (a,b). Composition defines an associative operation:

Associativity allows a tree of compositions. It does not allow arbitrary reordering: in general, exchanging two time steps changes the result. A scan returns every prefix map, rather than only the final reduction. With zero initial state, the b component of each prefix is the state at that time. With a nonzero initial state, apply the prefix's a component to it too. §2.2 and Appendix H

A work-efficient scan over L diagonal P-dimensional updates uses O(LP) arithmetic and logarithmic dependency depth in L. The paper describes parallel time O(P log L) with L processors. These are algorithmic statements under a parallel execution model. The dense B/C projections cost O(LPH), and the scan states occupy O(LP) space during whole-sequence processing. Streaming inference can instead keep just the current state per layer. Constant recurrent state does not imply constant training memory. §2.2 and §3.4

A dense transition would turn composition into matrix multiplication and make this route much more expensive. S5's proposition comparing complexity with S4 assumes P is O(H). Neither that asymptotic comparison nor our small Python loop establishes a GPU latency advantage. Compiler fusion, parallel kernels, shapes, dtype and hardware determine measured runtime.

The pair operation still works when each observation has a different interval, because each step can have its own a and b. A fixed convolution kernel assumes time invariance; affine scan composition does not. That flexibility will matter for irregular observations, but first we need a useful starting spectrum.

Initialization is part of the method

A negative real part can keep a mode from exploding while still giving it a poor memory of a long sequence. S5 inherits its initialization strategy through a carefully limited connection to S4 and the HiPPO framework.

The original HiPPO-LegS matrix is not stably diagonalizable in the way S5 needs. S4 represents it as a normal component plus a low-rank correction. S5 initializes from the normal component, HiPPO-N, which can be diagonalized with a unitary basis. Dropping the low-rank component changes the finite-dimensional system; it is not an exact basis transformation of HiPPO-LegS. The paper motivates this approximation through the S4/S5 relationship and an asymptotic HiPPO result, then evaluates it empirically. §4.2 and Appendix B.1

For zero-based indices, the construction uses p_n=\sqrt{n+1/2} and A_N=A_{\mathrm{LegS}}+pp^T. This has diagonal −½ and skew-symmetric off-diagonal entries. Its initial eigenvalues therefore have real part −½, with imaginary parts that supply different oscillatory modes. The official initializer computes the spectrum and transforms the input and output projections into that basis. Pinned ssm_init.py

The intervals spread these modes over multiple discrete timescales. Appendix B.1 samples log Δ uniformly between log 0.001 and log 0.1 by default. Path-X lowers the minimum to 0.0001. Those are initialization ranges, not constraints that intervals must remain within throughout optimization. The paper also uses repeated smaller HiPPO-N blocks for initialization; B and C remain dense. Appendix B.1

The S4 connection deserves a precise boundary. Under shared S4 state matrices, shared intervals, matching S5 state size P=N, concatenated input vectors and compatible initial conditions, the S5 state can be interpreted as the sum of the independent S4 states. A common output matrix then reads that sum. This does not reproduce a general S4 layer's original output projections, nonlinearity and mixing layer. Proposition 2 concerns related latent dynamics under tied assumptions; the practical model relaxes those assumptions by learning more modes, per-mode intervals and using block initialization. §4.1–4.3

The ablations support those practical choices rather than proving one universal best initialization. Table 5 reports Image accuracy rising from 82.01% for the constrained scalar-interval S5 to 85.04% with per-mode intervals, then 88.00% for the less constrained configuration. The last comparison also changes latent size and initialization, so attributing its whole improvement to just one change would be unjustified. Appendix E, Table 5

Sampling intervals carry physical meaning

When observations arrive at irregular intervals, compute a different discretization for each interval and scan the resulting affine maps. When the sampling rate changes uniformly, rescale Δ by the old sampling rate divided by the new sampling rate. A model trained at 16 kHz sees twice as much elapsed time per observation at 8 kHz, so its interval doubles. This is easy to reverse accidentally if “sampling rate” and “sampling interval” are treated as synonyms. §6.2–6.3

The figure holds a constant input, making exact agreement at common physical times checkable. Real audio has changing content, and decimation can discard information or introduce aliasing. A continuous-time parameterization provides a principled adjustment; it does not promise unchanged accuracy under arbitrary resampling.

The paper evaluates irregular observations on a pendulum regression task, using 50 noisy 24×24 images sampled over a duration of 100 and targets derived from the pendulum angle. That is evidence for the reported task and protocol, rather than evidence that any learned SSM identifies the underlying physical dynamics. §6.3

The executable reference below makes the linear mechanics inspectable before we turn to the larger trained systems.

Paper equations beside executable Python

Read these five rows in order: input projection, discretization, scan, real readout and initialization. Each links the relevant paper section and pinned official code to an independently authored Python region. The recursive scan computes every prefix, including odd-length sequences; it is not a final-state-only reduction.

1. Mix input channels before the recurrence

Original paper

Enlarge figure

Math and tensor shapes

\begin{aligned}\tilde B&=V^{-1}B\\\tilde C&=CV\\b_k&=\bar B u_k\end{aligned}

u: [H], B̄: [P,H] → b: [P]; sequence: [L,H] → [L,P]

Every mode may receive every input feature. This dense multiplication is separate from the cheap elementwise transition. The reference validates the projection width; it does not compute an eigendecomposition.

Pinned official implementation

Python implementation

python

def project(matrix, vector):
    """Multiply a dense matrix by a vector with explicit shape checks."""
    if any(len(row) != len(vector) for row in matrix):
        raise ValueError("Matrix width must equal vector length")
    return [sum(weight * value for weight, value in zip(row, vector))
            for row in matrix]


def elements(a_bar, b_bar, inputs):
    """Build (diagonal transition, dense input projection) at each time."""
    return [(list(a_bar), project(b_bar, u)) for u in inputs]

2. Integrate a held input exactly

Original paper

Enlarge figure

Math and tensor shapes

\begin{aligned}\bar\lambda_p&=e^{\Delta_p\lambda_p}\\\bar B_{ph}&=\Delta_p\phi_1(\Delta_p\lambda_p)\tilde B_{ph}\\\phi_1(z)&=(e^z-1)/z\end{aligned}

λ, Δ: [P]; B̃ and B̄: [P,H]; φ₁(0)=1

The series branch handles zero and nearly zero poles without division by zero or subtracting nearly equal numbers. This is a numerical guard added to the teaching reference; the pinned official implementation uses the direct quotient. It is not an Euler input term ΔBu.

Pinned official implementation

Python implementation

python

def discretize(poles, b_matrix, steps):
    """Return diagonal ZOH A and dense B for per-mode positive steps."""
    if not poles or len(poles) != len(b_matrix) or len(poles) != len(steps):
        raise ValueError("Expected one B row and interval per pole")
    if any(not math.isfinite(dt) or dt < 0 for dt in steps):
        raise ValueError("Intervals must be finite and nonnegative")
    a_bar, b_bar = [], []
    for pole, row, dt in zip(poles, b_matrix, steps):
        z = pole * dt
        a_bar.append(cmath.exp(z))
        if abs(z) < 1e-4:
            term, phi = 1.0, 1.0
            for n in range(1, 9):
                term *= z / (n + 1)
                phi += term
            gain = dt * phi
        else:
            gain = (cmath.exp(z) - 1) / pole
        b_bar.append([gain * value for value in row])
    return a_bar, b_bar

3. Recover all prefixes of affine updates

Original paper

Enlarge figure

Math and tensor shapes

\begin{aligned}(a_i,b_i)\bullet(a_j,b_j)&=(a,b)\\a&=a_j\odot a_i\\b&=a_j\odot b_i+b_j\end{aligned}

elements: L pairs of [P] vectors → L prefix maps; x₀=0 → xₖ=prefix.b

Pair reduction followed by prefix reconstruction uses linear work and logarithmic dependency depth. Python executes this algorithm serially. A parallel runtime such as JAX supplies actual parallel scheduling; associativity is exact algebraically, with floating-point regrouping differences.

Pinned official implementation

Python implementation

python

def compose(left, right):
    """Compose earlier left then later right; multiplication is elementwise."""
    a_left, b_left = left
    a_right, b_right = right
    return ([r * l for l, r in zip(a_left, a_right)],
            [r * l + b for r, l, b in zip(a_right, b_left, b_right)])


def prefix_scan(items):
    """Recursive work-efficient scan; Python executes it serially."""
    if len(items) <= 1:
        return list(items)
    paired = [compose(items[i], items[i + 1])
              for i in range(0, len(items) - 1, 2)]
    prefixes = prefix_scan(paired)
    return [prefixes[i // 2] if i % 2 else
            item if i == 0 else compose(prefixes[i // 2 - 1], item)
            for i, item in enumerate(items)]

4. Reconstruct real outputs and keep the direct path

Original paper

Enlarge figure

Math and tensor shapes

\begin{aligned}r_k&=2\operatorname{Re}(\tilde C_+\tilde x_{k,+})\\y_k&=r_k+D\odot u_k\end{aligned}

half-spectrum x: [P/2], C: [H,P/2], D: [H] → y: [H]

Use factor two only when the omitted modes and projections are conjugate partners. A full-spectrum computation instead takes the real part once. The direct term is not doubled. These are preactivations; nonlinearities, normalization, dropout and residual connections belong to the surrounding layer.

Pinned official implementation

Python implementation

python

def readout(states, inputs, c_matrix, direct, half_spectrum=False):
    """Read real outputs; factor two is only for paired half-spectrum storage."""
    factor = 2 if half_spectrum else 1
    outputs = []
    for state, u in zip(states, inputs):
        hidden = project(c_matrix, state)
        outputs.append([factor * value.real + d * source
                        for value, d, source in zip(hidden, direct, u)])
    return outputs

5. Build the normal starting matrix and interval range

Original paper

Enlarge figure

Math and tensor shapes

\begin{aligned}A_N&=A_{\mathrm{LegS}}+pp^T\\p_n&=\sqrt{n+\tfrac12}\\\log\Delta_p&\sim U(\log\delta_{min},\log\delta_{max})\end{aligned}

normal matrix: [P,P]; default interval bounds: 0.001 to 0.1

The helper constructs the real matrix before diagonalization. Its deterministic log grid illustrates coverage of scales; the official initializer samples log intervals randomly. No eigenbasis, learned weights or training claims are hidden in this compact function. Path-X uses a smaller minimum interval, 0.0001.

Pinned official implementation

Python implementation

python

def hippo_normal(size):
    """Construct HiPPO-N = HiPPO-LegS + p p^T, before diagonalization."""
    if size < 1:
        raise ValueError("State size must be positive")
    return [[-0.5 if n == k else
             (1 if n < k else -1) * math.sqrt((n + 0.5) * (k + 0.5))
             for k in range(size)] for n in range(size)]


def log_spaced_steps(count, low=0.001, high=0.1):
    """Deterministic log grid for teaching, not upstream random initialization."""
    if count < 1 or not 0 < low <= high:
        raise ValueError("Expected positive count and ordered positive bounds")
    return [math.exp(math.log(low) + i / max(count - 1, 1)
                     * math.log(high / low)) for i in range(count)]

Original paper excerpts: Jimmy T. H. Smith, Andrew Warrington, Scott W. Linderman. Paper v3. Cropped/rasterized from the original PDF; no scientific content redrawn.

Download the standard-library Python reference and its requirements note. Run python s5.py. It checks scan/loop agreement, full/half-spectrum readout equivalence, the zero-pole limit and normal-matrix structure. The output consists of deterministic numerical checks and a tiny untrained example. No checkpoint, training data or GPU is required.

The reference intentionally stops at a causal linear SSM. It constructs HiPPO-N before eigendecomposition, but does not implement a general eigensolver, optimizer, residual network or dataset pipeline. Its deterministic interval grid illustrates logarithmic spacing; the official initializer uses random samples. Its recursive Python scan executes serially, so it is not a throughput benchmark. The official JAX code supplies the actual model and experiment configurations.

Published results: strong long-range classification, qualified comparisons

Table 1 of v3 reports these test accuracies in percent. The S5 results are means across three seeds. Sequence lengths range from 1,024 for Image and Pathfinder to 16,384 for Path-X. The average weights these six tasks equally; it is not an average over individual examples or tokens.

Source: v3 Table 1, with uncertainty in Appendix F Table 7. These are cited measurements, not reproduced results. The abstract gives coarser 87.4% and 98.5% figures; the table entries above retain the published precision. The mean of the six displayed S5 scores is approximately 87.462%, consistent with 87.46% after rounding.

S5 has the highest Path-X score among these rows, but does not win every task or the overall average against the quadratic Mega variant. Its comparison with S4 is about a competitive long-range model with simpler recurrent execution, not a universal accuracy or speed dominance claim. The linear Mega-chunk variant is a separate configuration, reported at 85.66% average in the same table.

The released implementation supports bidirectional scans for offline tasks. A backward scan sees future observations and cannot be used unchanged for causal streaming. The small causal reference above therefore demonstrates the operator, while published classification systems must be interpreted with the task-specific configurations in Appendix G. Task-specific width, state size, block initialization, nonlinear layers and tuning prevent the aggregate leaderboard from isolating MIMO as the sole cause of improvement.

Speech Commands provides a different test of the continuous-time formulation. Table 2 reports S5 at 96.52% for 16 kHz and 94.53% for zero-shot 8 kHz, using a 280K-parameter model and means across three seeds. The 8 kHz inputs are constructed by naive decimation and use interval rescaling without fine-tuning. Liquid-S4 reports 96.78% at 16 kHz and 90.00% at 8 kHz. This supports useful rate transfer for the stated audio task; it does not establish general invariance to a sensor or domain shift. §6.2 and Table 2

For performance claims, the scope must stay equally narrow. The scan's linear work and reduced recurrent state are useful design properties. End-to-end latency includes projections, nonlinear layers, memory traffic and runtime overhead. This review supplies no new training run or measured timing; reproducing a speed comparison requires the exact task configuration, batch and sequence shapes, precision, compiler, hardware and timing boundary.

What the simplification buys, and what it costs

S5 gives one recurrent formulation for streaming and whole-sequence processing. It replaces a bank of channel-specific memories with a shared modal memory; puts channel mixing in B and C; initializes diagonal dynamics from a structured spectrum; and uses scan composition for offline parallelism. The same affine algebra tolerates time-varying discretizations.

The costs follow directly from those choices. Fixed-size state compresses history, so representable memory depends on the mode count and spectrum. Dense projections can dominate even when the scan is cheap. Diagonalization does not eliminate sensitivity to initialization, interval choice or numerical precision. Bidirectional benchmark settings do not transfer directly to online generation. And the linear SSM's coefficients are learned parameters, not automatically input-dependent gates; scan compatibility alone should not be confused with the selective mechanism of later Mamba models.

The central lesson is a useful separation: what the state stores, how channels address it, and how time steps are scheduled are different design decisions. S5 makes those decisions explicit enough to inspect, implement and test independently.

Related reading

• Mamba: how input-dependent selection changes the fixed linear dynamics considered here.
• Mamba-2: a further connection between structured state-space models and attention computation.
• VMamba: how sequence recurrences are arranged into spatial scan orders for images.
• Linear attention approximations: another route to sequence processing through finite recurrent summaries.
• Layer normalization: the normalization component around a state-space layer, separate from its linear recurrence.
• S5 paper v3 and pinned official source: derivations, ablations, complete experiment settings and JAX implementation.