Simplified State Space Layers for Sequence Modeling

In plain words

Many sequence tasks, such as reading a long document or listening to audio, need a model to connect things that are thousands of steps apart. The earlier S4 layer handled this well, but it did so with many small, separate memory units and a fairly involved calculation that turned them into a convolution.

S5 keeps the same kind of memory, a linear state space model, but uses one shared memory block in which all the inputs and outputs are mixed together. Because this block is simple and diagonal, its step-by-step updates can be computed in parallel with a standard "scan" trick, the same kind of operation used for prefix sums. The result runs in the time domain, works for sequences that are sampled at irregular times, and is about as fast as S4.

The authors also show that the starting values for the memory match those of S4 closely enough to keep its accuracy on long-range benchmarks. The numbers reported are the authors' own.

TL;DR

S5 is a state space layer whose state transition is diagonal. It runs as a parallel linear recurrence over time, with per-mode learnable timescales. It is initialised from a diagonalised HiPPO-N matrix and matches S4's asymptotic cost when the state size is on the order of the feature size §3.4. The paper reports state-of-the-art results on Long Range Arena (LRA) and Speech Commands at 8 kHz, and near-S4 results on the rest. Treat the numbers as author-reported.

Problem

S4 applies a bank of independent single-input, single-output state space models (SSMs), one per feature channel. Efficient training uses a frequency-domain convolution kernel, which needs a non-trivial computation and assumes a time-invariant system with regularly spaced observations §2.3 §3.3. The HiPPO-LegS matrix that S4 uses for initialisation cannot be diagonalised stably, which blocks a direct parallel-scan version §2.3 §4.2. Efficient transformers still perform poorly on very long-range tasks §1.

Method

The S5 layer has one MIMO SSM with latent size P and input and output size H §3.1. The continuous system is diagonalised, A = V Λ V⁻¹, and rewritten with reparameterised input and output matrices B̃ = V⁻¹B and C̃ = CV eq. 4 eq. 5. Each diagonal mode gets a learnable timescale Δ ∈ ℝ^P, and the system is discretised by zero-order hold §3.2 eq. 6. A feedthrough matrix D is restricted to be diagonal §3.2.

Steps for one layer, applied to a sequence u₁:L of shape (L, H):

1. Discretise each diagonal mode with its own timescale §3.2 eq. 6.
2. Run the linear recurrence over time with a parallel scan on the pairs (Λ̄, B̄uₖ). The scan operator is associative §2.2 app. H.
3. Read out the real part of C̃xₖ and add the feedthrough term Duₖ §3.2 eq. 5. The code doubles the real part, which the paper does not state (see Paper vs code).
4. Apply a gated activation, GELU(y) ⊙ σ(W·GELU(y)), with W a dense H×H layer app. G.1.

Equations as stated in the paper:

\overline{\Lambda} = e^{\Lambda \Delta}, \quad \overline{B} = \Lambda^{-1}(\overline{\Lambda} - I)\tilde{B}, \quad \overline{C} = \tilde{C}, \quad \overline{D} = D eq. 6

x_k = \overline{\Lambda} x_{k-1} + \overline{B} u_k eq. 2, specialised to the diagonal case of §3.2

Pseudocode, following the JAX listing in app. A and the official code:

# S5 layer, one sequence  [§3.2] [app. A] [code: s5/ssm.py#L60-L88]
# Inputs:  u (L, H) real; Lambda (P,) complex; B_tilde (P, H) complex
#          C_tilde (H, P) complex; D (H,) real; log_step (P,) real
#          W (H, H) dense; b (H,) dense bias
# Output:  (L, H) real
def s5_layer(u, Lambda, B_tilde, C_tilde, D, log_step, W, b):
    # Per-mode timescale, log-parameterised  [§3.2] [app. B.1.3]
    delta = exp(log_step)                                   # (P,)
    # Zero-order-hold discretisation of each diagonal mode  [§3.2] [eq. 6]
    Lambda_bar = exp(Lambda * delta)                        # (P,)
    B_bar = ((Lambda_bar - 1) / Lambda)[:, None] * B_tilde  # (P, H)
    # Drive term B_bar u_k for every step
    drive = u @ B_bar.T                                     # (L, P)
    # Linear recurrence x_k = Lambda_bar * x_{k-1} + B_bar u_k,
    # computed as a parallel scan  [§2.2] [app. H]
    states = parallel_scan(linear_op, Lambda_bar, drive)    # (L, P)
    # Read out the real part of C_tilde x_k  [§3.2] [eq. 5]
    # The factor 2 is code-only, see Paper vs code
    # [code: s5/ssm.py#L86]
    ys = 2 * real(states @ C_tilde.T)                       # (L, H)
    # Feedthrough with a diagonal D  [§3.2] [eq. 5]
    ys = ys + D * u                                         # (L, H)
    # Gated activation, half_glu1 variant  [app. G.1]
    # The dense gate has a bias in the code, not stated in the paper
    hidden = gelu(ys)                                       # (L, H)
    gate = sigmoid(hidden @ W.T + b)                        # (L, H)
    return hidden * gate                                    # (L, H)

Operator for the scan, on pairs (a, b) that stand for the map x ↦ a·x + b app. H:

# Binary operator for the parallel scan  [app. H] [code: s5/ssm.py#L46-L57]
# Each element is a pair (a, b) meaning x -> a * x + b, with a diagonal
def linear_op(first, second):
    # Compose the earlier element with the later one
    a_first, b_first = first
    a_second, b_second = second
    # Associative, so the scan can be evaluated as a tree  [app. H]
    return a_second * a_first, a_second * b_first + b_second

Initialisation, following app. B.1:

• Λ starts from the eigenvalues of the diagonalised HiPPO-N matrix, with real part −½ app. B.1.1.
• B̃ = V⁻¹B and C̃ = CV, where B and C are sampled. The paper does not name the distribution app. B.1.2.
• D is sampled from a standard normal app. B.1.2.
• Each Δ element is drawn log-uniformly from [0.001, 0.1], or from [0.0001, 0.1] for Path-X app. B.1.3.
• Optionally, the state matrix is block-diagonal, with J copies of the HiPPO-N block app. B.1.1.

Hyperparameters

Settings shared by every task:

Task-specific hyperparameters from app. G.2. The paper's values come first, and the last column compares them with the run scripts in the code. The effective learning rate in the code is the global LR, lr_factor × ssm_lr_base train.py.

Results

Author-reported throughout. Setup: LRA's six tasks, sequence lengths 1,024 to 16,384, and three seeds for S5 and S4 §6.1 app. F.1. Speech Commands 35-way with 16 kHz audio and three seeds §6.2 app. F.2. Pendulum regression on a 24×24 image sequence with 50 irregular samples §6.3 app. G.4. Hardware is a 16 GB V100 for the runtime comparison app. C.2 and an RTX 2080 Ti for the pendulum runs app. F.3.

Notes on the table:

• The LRA numbers are means over three seeds, with standard deviations in app. F.1. Across all rows in that table, S5 is highest on Retrieval and Path-X. Mega, which has quadratic cost, is higher on ListOps, Text and Pathfinder. S5 is lower than S4-LegS on Image.
• The paper's own ablation (LRA, app. E): with P = N and a scalar Δ, S5 reaches 57.20 on ListOps and 96.59 on Path-X, against 62.15 and 98.58 for the full model. The state-initialisation ablation app. E shows discrete-time HiPPO-N at 61.93 on Text, against 89.31 for continuous-time HiPPO-N.
• Runtime app. C.2: S5 with P = H trains between 0.7 and 1.9 times as fast as S4D, depending on the task. The configured S5 reaches 4.7 times on Path-X. Memory is 0.7 to 1.2 times S4D.
• Pendulum: S5 takes 0.25 s per training epoch, against 22 s for the authors' re-run of CRU app. F.3. The paper also reports two ablations: S5-drop 6.68 and S5-append 4.13 (lower is better) app. F.3.

Paper vs code

The official code is at github.com/lindermanlab/S5, linked from the paper's introduction §1. It is MIT-licensed and was checked at code_commit.

Discrepancies:

• Learning rate for Pathfinder, Image and Path-X. The table gives LR 0.005, 0.005 and 0.002 app. G.2. The scripts set lr_factor to 5, 4.5 and 3 times the SSM LR, which gives 0.0045, 0.0045 and 0.0018 run_lra_pathfinder.sh run_lra_cifar.sh run_lra_pathx.sh train.py. The paper's sCIFAR row (0.0045) matches the code, and the paper says sCIFAR uses the LRA-image setting, so the Image row looks like the typo.
• "No exceptions to the defaults" for Text. The paper says this app. G.3.2. The IMDB script sets dt_global, half_glu2 and warmup_end=0 run_lra_imdb.sh.
• GLU variant. The paper's gate is GELU(y) ⊙ σ(W·GELU(y)) app. G.1, which is half_glu1 layers.py. ListOps and IMDB scripts use half_glu2 run_lra_listops.sh run_lra_imdb.sh. That variant gates the un-activated y with GELU only in the gate: y ⊙ σ(W·GELU(y)) layers.py. The paper does not mention it.
• Path-X output projection. The paper says C̃ = CV app. B.1.2. The Path-X script uses complex_normal, which samples C̃ directly with no V run_lra_pathx.sh ssm.py. The paper does not state this.
• Readout factor 2. With conjugate symmetry, the code returns 2 · Re(C̃x) ssm.py. The paper says only that half the eigenvalues are used §3.2. The factor 2 is needed to match the full-state output. Anyone implementing from the paper alone will be off by a factor of two.
• Meaning of P. The table's P is the full count, ssm_size_base (for Path-X, 256). The code halves it to P/2 complex states train.py. §3.2 states the halving, but the table does not.
• 8 kHz rescaling direction. The paper says to rescale Δ by "the ratio between the new and old sampling rates" §6.2. Read literally, 8 kHz over 16 kHz gives 0.5. The code multiplies by 2.0 train.py, which is old over new. The paper's wording is ambiguous, and the code's direction is the one that doubles the time per sample.
• Warmup. The paper's hyperparameters do not mention warmup. Code defaults to one epoch of linear warmup before cosine annealing run_train.py train_helpers.py.
• Clipped eigenvalues. The LRA-image run sets clip_eigs=True run_lra_cifar.sh, which clips the real part of Λ to at most −1e-4 ssm.py. The paper does not mention this.
• Δ in the global group. With dt_global (IMDB and AAN scripts), the timescale Δ moves from the SSM learning-rate group to the global group, with weight decay train_helpers.py. The paper does not describe this option.
• Bidirectional scans. The code runs a forward and a backward scan and concatenates them, so C̃ has 2P columns ssm.py ssm.py. Section 3 describes only the causal model.
• Dropout. The dropout mask is broadcast over axis 0, the time axis of the per-sequence layout, so one mask covers every time step layers.py. The paper does not say this.
• Dense gate bias. The gate layer has a bias layers.py. The paper's W has none.
• Eigenvalue initialisation. The paper says Λ comes from the eigendecomposition of A^Normal app. B.1.1. The code sets the real part to the mean diagonal of S, which is −½, and takes the imaginary parts from eigh on −i·S ssm_init.py. It keeps the first P/2 eigenvalues in ascending order train.py. This picks one member of each conjugate pair, and the paper does not say which. eigh assumes a Hermitian matrix and −i·S is not Hermitian, so the code relies on how eigh reads the input. I did not verify the eigenvalues numerically.
• Abstract rounding. The abstract gives 98.5% on Path-X and the table gives 98.58 tab. 1. Cite the table value.
• Missing code. No pendulum dataset or script is in the repository. There are MNIST and permuted-MNIST loaders dataloading.py dataloading.py, but no run scripts for sMNIST, psMNIST or the colour sCIFAR.

Pitfalls

• Conjugate symmetry has three moving parts: the halved state count, the factor-2 readout, and which half of the eigenvalues to keep. Get all three right, or the output is silently off.
• Parallel scan memory grows as O(P·L) for the stored states, and the paper's runtime claims assume the JAX implementation on a V100 app. C.1 app. C.2.
• Block-diagonal initialisation. With J > 1, the initial Λ repeats the HiPPO-N block. B̃ and C̃ stay dense, and nothing keeps the state matrix block-diagonal during training app. B.1.1.
• Irregular sampling works by passing a different Δ for each step §3.3 §6.3. The pendulum code is not in the repository, so the variable-interval results cannot be reproduced from the code here.
• Speech 8 kHz is a naive decimation of the 16 kHz test set, with Δ rescaled §6.2 app. G.4. It is a zero-shot test of a model trained at 16 kHz, not a retrained model.
• Sampling distributions for B and C are not stated. The code uses lecun-normal for B and truncated-standard-normal for C by default, with complex-normal for Path-X ssm.py run_train.py.
• Optimiser groups. Norm parameters and the SSM parameters are in one group with the SSM learning rate under the default settings train_helpers.py. Changing opt_config changes which of B, C and D are decayed, and the paper states only a general rule app. G.2.1.

Lineage

• Builds on: S4 (Gu et al., 2021), for the layer structure, the convolution mode and the HiPPO initialisation. S5 keeps the S4 architecture around the layer and changes the SSM §1 §3.
• Builds on: the diagonal S4D (Gu et al., 2022) and DSS (Gupta et al., 2022) variants. S5 applies their diagonalised HiPPO-N idea to the MIMO setting §4.2.
• Compared against: S4-LegS and S4D-LegS tab. 1 tab. 2; Liquid-S4 (Hasani et al., 2022); Mega and Mega-chunk (Ma et al., 2022); CRU for pendulum regression app. F.3.
• Superseded by: not known from content already in this wiki.