PaperlayerView exact revision on arXiv

Source-linked AI summary

Simplified State Space Layers for Sequence Modeling

Jimmy T. H. Smith, Andrew Warrington, Scott W. Linderman

arXiv:2208.04933v3cs.LG

TL;DR

Long-range sequence modeling requires handling dependencies across thousands of timesteps, while efficient transformers can still perform poorly on very long-range tasks. The paper introduces S5, a single MIMO state space layer using diagonalized dynamics and parallel scans, and connects it mathematically to S4 to guide initialization. S5 matches S4’s computational efficiency and achieves strong long-range results, including 87.4% average LRA performance and 98.5% on Path-X.

  • Problem

    Long-range sequence modeling remains challenging because relevant information may be encoded between observations thousands of timesteps apart, and efficient transformers still perform poorly on very long-range tasks.

  • Method

    S5 replaces S4’s bank of independent SISO SSMs with one MIMO SSM using diagonalized dynamics and parallel scans, while an S4–S5 relationship informs its initialization and parameterization.

  • Results

    S5 retains S4-like computational efficiency and achieves state-of-the-art long-range performance, averaging 87.4% on LRA and reaching 98.5% accuracy on Path-X.

  • Takeaways & Limitations

    S5 provides an efficient recurrent time-domain layer that supports long-range modeling and can handle time-varying or irregularly sampled SSMs.

  • Takeaways & Limitations

    The S4–S5 dynamics equivalence is established only under simplifying assumptions, including shared state matrices and timescales.

Abstract

from arXiv · show

Models using structured state space sequence (S4) layers have achieved state-of-the-art performance on long-range sequence modeling tasks. An S4 layer combines linear state space models (SSMs), the HiPPO framework, and deep learning to achieve high performance. We build on the design of the S4 layer and introduce a new state space layer, the S5 layer. Whereas an S4 layer uses many independent single-input, single-output SSMs, the S5 layer uses one multi-input, multi-output SSM. We establish a connection between S5 and S4, and use this to develop the initialization and parameterization used by the S5 model. The result is a state space layer that can leverage efficient and widely implemented parallel scans, allowing S5 to match the computational efficiency of S4, while also achieving state-of-the-art performance on several long-range sequence modeling tasks. S5 averages 87.4% on the long range arena benchmark, and 98.5% on the most difficult Path-X task.

1 INTRODUCTION

Long-range sequence modeling remains difficult because relevant observations can be thousands of timesteps apart, while efficient transformers still perform poorly on very long-range tasks. S5 addresses this by replacing S4’s independent SSM bank with one MIMO SSM and parallel scans, retaining S4-like efficiency and strong benchmark performance.

  • Efficient long-sequence modeling is difficult because task-relevant information may link observations thousands of timesteps apart.
  • S5 replaces S4’s bank of independent SISO SSMs with one MIMO SSM and uses an efficient parallel scan.This removes S4’s convolutional and frequency-domain computation of the convolution kernel.
  • S5’s resulting layer has the same computational complexity as S4 while operating recurrently in the time domain.
  • S5’s initialization and parameterization are informed by a mathematical relationship between S4 and S5 and by HiPPO-based initialization.A diagonal approximation is used because S4’s specific HiPPO matrix cannot be stably diagonalized for S5.
  • 87.4% average LRA performance and 98.5% accuracy on Path-X demonstrate strong performance across long-range sequence modeling tasks.

2 BACKGROUND

The background introduces continuous-time and discretized linear SSMs, parallel scans, and S4’s structured architecture. S4 uses independent SISO systems with HiPPO initialization and convolution-based offline processing.

  • A continuous-time linear SSM maps inputs to latent states and outputs using state, input, output, and feedthrough matrices.
  • Discretization yields the recurrence x_k = Ax_k−1 + Bu_k and y_k = Cx_k + Du_k.
  • A parallel scan computes all-prefix compositions and can evaluate a diagonal SSM recurrence in O(PL) total operations.
  • S4 maps H input channels through H independent N-dimensional SISO SSMs, followed by nonlinear activation and position-wise linear mixing.
  • S4 uses HiPPO-initialized state matrices and convolution mode with FFTs for efficient offline processing of evenly spaced sequences.
  • S4 contains independent copies of learnable SSM parameters alongside O(H^2) mixing-layer parameters.

3 THE S5 LAYER

S5 replaces S4’s independent SISO bank with a dense MIMO SSM whose smaller latent state supports efficient diagonal parallel scans. Its parameterization and initialization are designed for recurrent time-domain computation.

  • S5 uses a latent state of size P to transform H-dimensional inputs into H-dimensional SSM outputs through one MIMO system.Because features are already mixed, S5 does not require S4’s additional position-wise linear layer.
  • S5 diagonalizes its continuous-time state matrix so the linear recurrence can be computed efficiently with parallel scans.
  • The discretized S5 system uses ZOH parameters derived from the diagonalized dynamics and a timescale parameter.
  • S5 learns diagonalized input, output, feedthrough, state, and timescale parameters, with diagonal feedthrough restricted in practice.
  • HiPPO-like initialization is used because the HiPPO-N matrix can be stably diagonalized for efficient parallel scans.
  • P is smaller than S4’s HN effective latent size, enabling efficient scans for both online generation and offline processing.
  • P = O(H) gives S5 the same order of runtime and memory complexity as S4.

4 RELATIONSHIP BETWEEN S4 AND S5

The paper relates S5 and S4 under simplifying assumptions: their latent state dynamics can be equivalent even though their output projections differ. This connection motivates transferring HiPPO-based initialization to S5, while the assumptions constrain the equivalence.

  • The S4–S5 dynamics analysis assumes equal input/output dimensions, shared S4 state matrices, shared timescales, and a matching S5 state matrix.
  • Under these conditions, Proposition 2 relates an S5 MIMO system to S4 systems with the same state matrix, input vectors, and discretization timescales.
  • S5 outputs are not identical to block-diagonal S4 outputs; they correspond to a modified equivalent output matrix while the underlying state dynamics match.
  • The relationship motivates HiPPO initialization for S5 because HiPPO initialization was important to S4 performance.
  • Exact HiPPO-LegS initialization is unsuitable because its matrix is not stably diagonalizable, motivating diagonalized HiPPO-N instead.
  • The equivalence applies only to constrained versions of S4 and S5, since tied state matrices and shared timescales restrict the comparison.

5 RELATED WORK

Prior work connects linear recurrent models and parallel scans, providing context for S5’s use of parallel-scan computation.

  • Linear surrogate RNNs approximate nonlinear RNNs with stacks of linear RNNs connected by nonlinear layers.
  • QRNNs and SRUs belong to a class of linear surrogate RNNs that can leverage parallel scans.
  • Prior work also studied parallel scans as an efficient computation strategy for recurrent models.

6 EXPERIMENTS

Experiments compare S5 with S4 and other baselines across long-range, speech, and irregularly sampled sequence tasks. S5 matches S4’s performance and efficiency while achieving strong benchmark results and handling variable observation rates.

  • Experimental setup: S5 matches S4’s performance and efficiency when used as a drop-in replacement with comparable computational complexity and parameter counts.The experimental architecture uses a linear encoder, stacked S5 layers, and a linear decoder.
  • Long Range Arena: 98.5% accuracy is achieved on Path-X, the longest LRA task, exceeding all compared models.
  • Raw speech classification: S5 outperforms baselines and previous S4 methods on the 35-way Speech Commands classification task, while performing similarly to Liquid-S4.The task uses high-fidelity recordings of 35 spoken words.
  • Raw speech classification: S5 can apply a model trained on 16kHz speech to decimated 8kHz data without fine-tuning by globally rescaling the continuous-time timescale parameter.The rescaling uses the ratio between the new and old sampling rates.
  • Variable observation interval: S5 achieves lower mean error and markedly faster application than CRU on irregularly sampled pendulum regression.The task uses 50 noisy 24 × 24 images sampled irregularly from a continuous trajectory, with sine and cosine angle targets.
  • Pixel-level 1-D image classification: S5 broadly matches S4 and outperforms several state-of-the-art RNN-based methods on sequential MNIST, permuted sequential MNIST, and sequential CIFAR.

7 CONCLUSION

The conclusion presents S5 as a simpler recurrent alternative to S4 that retains computational efficiency and supports time-varying or irregularly sampled state space models. It also identifies broader modeling opportunities enabled by the MIMO design.

  • S5 replaces S4’s frequency-domain convolutional approach with a purely recurrent time-domain approach using parallel scans.
  • S5 retains S4’s computational efficiency while enabling efficient processing of time-varying state space models.The conclusion connects this capability to sequences sampled at variable rates.
  • The general MIMO state space design may connect deep sequence models with probabilistic state space modeling and parallel filtering or smoothing.

A JAX IMPLEMENTATION OF S5 LAYER

The S5 implementation uses a diagonalized SSM and parallel scan to replace S4’s convolution-based computation while preserving comparable computational complexity. Its initialization combines HiPPO-inspired structure with dense input and output parameters.

  • Input, output, and feed-through matrices: S5 initializes its complex input and output parameters by transforming sampled matrices with the eigenvectors of the initial state matrix.The transformations are B_tilde = V^-1B and C_tilde = CV.
  • State-matrix initialization: S5’s default state matrix is initialized from the normal HiPPO approximation after block-diagonal construction and eigendecomposition.The number of HiPPO-N blocks is controlled by the J hyperparameter.
  • Input, output, and feed-through matrices: S5 samples its feedthrough matrix D elementwise from a standard normal distribution.The feedthrough matrix is restricted to be diagonal in the practical parameterization.
  • Implementation: S5 applies a parallel scan to a diagonalized linear SSM, avoiding S4’s convolutional and frequency-domain computation.The scan computes SSM outputs before a nonlinear activation function produces the layer outputs.
  • Timescale initialization: S5 initializes learnable timescales by sampling log ∆ uniformly between log δmin and log δmax, using a wider lower range for Path-X.The default range is δmin = 0.001 to δmax = 0.1; Path-X uses δmin = 0.0001 and δmax = 0.1.
  • Computational efficiency: S4 and S5 have the same order of runtime and memory complexity when S5’s latent size P = O(H).This holds for both parallelized offline processing and recurrent online generation.

C.2 EMPIRICAL RUNTIME COMPARISON

The empirical runtime study compares S4, S4D, and multiple S5 configurations across LRA tasks and sequence lengths using matched implementation conditions. Overall, S5’s runtime and memory usage are broadly similar to S4 and S4D.

  • Experimental setup: S5, S4, and S4D are benchmarked across three LRA tasks with varied sequence lengths using matched hardware and batch sizes.Comparisons use a 16GB NVIDIA V100 GPU, with bidirectionality and conjugate symmetry enabled as reported.
  • S5 configurations: S5 is evaluated with latent sizes P = H, P = N, and task-specific best architectural dimensions.The first two configurations test latent-size effects, while the third supplies the S5 models used for the LRA results.
  • Results: Runtime and memory usage of S5 and S4/S4D are broadly similar, supporting the computational-complexity analysis.The comparison is presented as the broad empirical takeaway of the runtime study.

D RELATIONSHIP BETWEEN S4 AND S5

The S4–S5 relationship shows that S5 can reproduce the latent dynamics of multiple S4 systems while using a different output projection. This connection motivates HiPPO-based initialization and extensions that relax the shared-parameter assumptions.

  • Diagonalizable initialization: The analysis extends the HiPPO-LegS-to-HiPPO-N approximation from scalar-input systems to vector-valued inputs, motivating diagonalizable HiPPO-N initialization for S5.This approximation enables efficient parallel scans.
  • Relaxing assumptions: S5 can implement multiple independent S4 systems through judicious state-matrix initialization, while vector-valued timescale parameters can improve performance.The main reported takeaway is that continuous-time parameterization with HiPPO initialization consistently performs well across tasks, including Path-X.
  • S4–S5 equivalence: S5 latent states equal a linear combination of the latent states computed by H SISO S4 systems under shared state-matrix and timescale assumptions.The equivalence follows from the linearity of the recurrences.
  • Different output projections: The equivalence concerns underlying latent dynamics, not equality of S5 outputs with the outputs of the block-diagonal S4 system.The systems differ through their output matrices.
  • Different output projections: S5 and S4 use different effective output matrices: S4 uses independent diagonal projections, whereas S5 ties dense output matrices across the S4 systems.Both projection matrices have the same number of parameters.

E.3 S4D INITIALIZATION ABLATIONS

The section reports ablations replacing HiPPO-N initialization with alternative diagonal matrices and places these results within the LRA and Speech Commands evaluations.

  • Initialization ablations: S4D-Inv and S4D-Lin are evaluated as alternative diagonal initializations by replacing HiPPO-N while keeping other factors unchanged.The same replacement procedure is applied to S4 and S5 ablations.
  • Evaluation: Table 7 reports test accuracy across the LRA benchmark tasks, with means across three seeds for S4, S4D, and S5.The table marks models that did not exceed random guessing.
  • Evaluation: Table 8 reports Speech Commands test accuracy for 35-way keyword spotting, including zero-shot testing on naively decimated 8kHz examples.Means across three random seeds are reported.

F.3 PENDULUM EXTENDED RESULTS

The extended pendulum results evaluate S5 variants and document the experimental architecture, hyperparameters, and timing conventions used for comparison.

  • Pendulum ablations: S5-drop removes dependence on the inter-sample interval, whereas S5-append supplies the integration timestep to the input encoder.The authors expect S5-drop to perform poorly and S5-append to perform as well as S5.
  • Pendulum evaluation: Table 9 reports test MSE ×10^-3 and runtimes for pendulum regression, using twenty random seeds for S5 methods.Baseline results come from prior work, while CRU is also rerun on the authors’ hardware.
  • Architecture: The sequence model uses a linear encoder, multiple S5 layers, mean pooling, a linear decoder, and Softmax for classification tasks.S5 is used as a drop-in replacement for S4 in the referenced architecture.
  • Experimental setup: Experiments control layer counts, feature sizes, and comparable parameter counts relative to prior reported configurations.Normalization, optimization, and task-specific exceptions are documented in the experimental settings.

G.4 DATASET DETAILS

The section describes the evaluated long-sequence datasets and explains how linear recurrences can be computed with associative parallel scans.

  • Datasets: Pendulum regression uses 100 simulated timesteps and 50 irregularly sampled frames to estimate the sine and cosine of the pendulum angle.The image pixels are corrupted by temporally correlated noise.
  • Parallel scans: For x_k = A x_{k−1} + B u_k, the scan precomputes sequence elements and combines them with an associative binary operator.The operator uses matrix multiplication, matrix-vector multiplication, and elementwise addition.
  • Parallel scans: The worked length-four example shows that the scan’s second tuple entry contains the desired latent state at each step.Parallel computation reorganizes the same recurrence through intermediate pairwise combinations.
Loading 2208.04933v3…