PaperlayerView exact revision on arXiv

Source-linked AI summary

Efficiently Modeling Long Sequences with Structured State Spaces

Albert Gu, Karan Goel, Christopher Ré

arXiv:2111.00396v3cs.LG

TL;DR

Long-range sequence modeling remains difficult at practical sequence lengths because existing approaches can be computationally or numerically infeasible. S4 reparameterizes state space models for efficient computation, achieving strong results across diverse benchmarks, including 88% accuracy on Path-X.

  • Problem

    Sequence models struggle to efficiently handle long-range dependencies across sequences spanning tens of thousands of time steps.

  • Method

    S4 reparameterizes structured state matrices with low-rank and normal components, enabling stable diagonalization and efficient Cauchy-kernel computation.

  • Results

    S4 outperforms baselines across diverse long-range-dependency benchmarks while remaining similarly fast, including 88% accuracy on length-16384 Path-X.

  • Takeaways & Limitations

    S4 shows potential as an effective general sequence modeling solution across modalities and model capabilities.

  • Takeaways & Limitations

    Prior state space approaches remain impractical because their state representations require prohibitive computation and memory, while naive diagonalization is numerically infeasible.

Abstract

from arXiv · show

A central goal of sequence modeling is designing a single principled model that can address sequence data across a range of modalities and tasks, particularly on long-range dependencies. Although conventional models including RNNs, CNNs, and Transformers have specialized variants for capturing long dependencies, they still struggle to scale to very long sequences of $10000$ or more steps. A promising recent approach proposed modeling sequences by simulating the fundamental state space model (SSM) \( x'(t) = Ax(t) + Bu(t), y(t) = Cx(t) + Du(t) \), and showed that for appropriate choices of the state matrix \( A \), this system could handle long-range dependencies mathematically and empirically. However, this method has prohibitive computation and memory requirements, rendering it infeasible as a general sequence modeling solution. We propose the Structured State Space sequence model (S4) based on a new parameterization for the SSM, and show that it can be computed much more efficiently than prior approaches while preserving their theoretical strengths. Our technique involves conditioning \( A \) with a low-rank correction, allowing it to be diagonalized stably and reducing the SSM to the well-studied computation of a Cauchy kernel. S4 achieves strong empirical results across a diverse range of established benchmarks, including (i) 91\% accuracy on sequential CIFAR-10 with no data augmentation or auxiliary losses, on par with a larger 2-D ResNet, (ii) substantially closing the gap to Transformers on image and language modeling tasks, while performing generation $60\times$ faster (iii) SoTA on every task from the Long Range Arena benchmark, including solving the challenging Path-X task of length 16k that all prior work fails on, while being as efficient as all competitors.

1 Introduction

S4 addresses the long-range dependency problem by making structured state space models computationally practical while retaining their theoretical ability to model long sequences. Across diverse benchmarks, it achieves strong accuracy, efficiency, generalization, and generation speed with minimal task-specific modification.

  • Motivation: Long-range dependency tasks remain difficult: Path-X contains length-16,000 sequences where prior models perform no better than random guessing.Real-world time-series data may require reasoning over tens of thousands of time steps.
  • Limitation: O(N^2L) computation and O(NL) space make the prior LSSL infeasible, versus an Ω(L + N) lower bound for both.For N = 256, LSSL uses orders of magnitude more memory than comparably sized RNNs or CNNs.
  • Method: S4 reparameterizes structured state matrices as low-rank plus normal terms and computes a truncated generating function in frequency space.This resolves the critical computational bottleneck in prior state space modeling.
  • Long-range dependency benchmarks: 88% accuracy on length-16384 Path-X establishes S4 as the first model to solve the task, compared with 50% random guessing for prior work.S4 is as fast as all LRA baselines while outperforming them by 20+ points on average.
  • General-purpose sequence modeling: 60× faster pixel/token generation than standard autoregressive models demonstrates S4’s recurrent-state advantage on CIFAR-10 and WikiText-103.S4 also reaches 2.85 bits per dim on CIFAR-10 density estimation and comes within 0.8 perplexity of Transformers on WikiText-103.
  • General-purpose sequence modeling: Over 90% accuracy on sequential CIFAR, surpassing Speech CNNs and the specialized Informer, shows S4 can learn with weaker inductive biases.S4 also adapts to 0.5× sampling frequency on speech classification without retraining.

2 Background: State Spaces

State-space models map scalar inputs through an N-dimensional latent state to scalar outputs, with HiPPO matrices enabling history memorization. After bilinear discretization, the model supports both recurrent computation and an equivalent convolution whose kernel is central to efficient implementation.

  • 2.1 Continuous-Time SSM: SSMs map a 1-D input signal u(t) to an N-D latent state x(t), then project it to a 1-D output y(t).The model parameters A, B, C, and D can be learned by gradient descent; D is omitted because Du acts as a skip connection.
  • 2.2 HiPPO and Long-Range Dependencies: HiPPO specifies matrices A that allow the state x(t) to memorize the history of the input u(t), addressing long-range dependencies.Replacing a random A with the HiPPO matrix improved sequential MNIST performance, according to the supplied passage.
  • 2.3 Discretization: For discrete inputs, the continuous-time SSM is sampled with step size ∆ and discretized using the bilinear method.The resulting matrices depend on both the continuous state matrix A and step size ∆.
  • 2.3 Discretization: The discrete SSM is a sequence-to-sequence map whose state equation is a recurrence, so x_k functions as an RNN-like hidden state.The discrete transition matrix is A, and the recurrence operates over x_k ∈ R^N.
  • 2.4 Convolution Representation: The recurrent SSM is equivalent to a discrete convolution y = K ∗ u, enabling efficient FFT computation once the SSM convolution kernel K is known.Computing K is non-trivial and is the focus of the paper’s technical contributions.

3 Method: Structured State Spaces (S4)

S4 makes state-space sequence modeling efficient by representing HiPPO-based systems as diagonal-plus-low-rank matrices and reducing their computation to stable Cauchy-kernel operations. This parameterization enables linear-time recurrence updates and near-linear convolution computation while preserving the SSM representations.

  • Motivation: Naive SSM recurrence computation requires O(N 2L) operations and O(NL) space because it repeatedly multiplies by A.This repeated matrix multiplication is the central computational bottleneck addressed by S4.
  • Motivation: Direct diagonalization is numerically infeasible because the HiPPO matrix has exponentially large diagonalization entries in state size N.The resulting change-of-basis matrices would be too ill-conditioned to compute reliably.
  • Efficient Computation: S4 computes convolution spectra with truncated generating functions, Woodbury corrections, inverse FFTs, and stable Cauchy-kernel algorithms.The diagonal-plus-low-rank structure turns matrix inverses into diagonal-case computations and reduces the diagonal case to Cauchy multiplication.
  • NPLR Parameterization: S4 uses a Normal Plus Low-Rank decomposition, and all HiPPO matrices admit this representation with rank r = 1 or r = 2.In particular, the HiPPO matrix used in equation (2) has rank r = 1.
  • Complexity Guarantees: O(N) operations compute one recurrent step, while 4 Cauchy multiplies compute the convolution filter in eO(N + L) operations and O(N + L) space.These results establish efficient recurrent and convolutional representations for S4.

4 Experiments

S4 is evaluated as an efficient general sequence model across long-range reasoning, speech, image and language modeling, forecasting, and sampling-rate changes. It substantially advances long-range benchmarks while retaining strong efficiency and its HiPPO initialization is shown to be central to performance.

  • Efficiency benchmarks: S4 is orders of magnitude more speed- and memory-efficient than LSSL and competitive with efficient Transformer variants in parameter-matched benchmarks.The comparisons include Linear Transformer and Performer.
  • Long Range Arena: 80.48% average accuracy lets S4 outperform all baselines across LRA’s six 1K-16K-step tasks and solve Path-X at length 16,384.Path-X requires reasoning over sequences formed from 128 × 128 inputs.
  • Raw Speech Classification: 98.3% accuracy on length-16000 raw speech classification exceeds all baselines using the reported preprocessing.S4 classifies raw Speech Commands signals rather than relying on MFCC features.
  • General sequence modeling: S4 remains competitive on CIFAR density estimation and WikiText-103 language modeling, while generation is around 60× faster than a vanilla Transformer on both tasks.The CIFAR sequences contain 3072 RGB subpixels, and WikiText-103 predicts tokens from past context.
  • Sampling resolution change: 96.3% accuracy at 0.5× sampling frequency is achieved without retraining by changing S4’s internal step size.S4’s continuous-time formulation adapts to data sampled at different rates.
  • Ablations: 84.27% test accuracy with 100K parameters supports HiPPO initialization as the primary source of S4’s effectiveness, rather than NPLR parameterization alone.Training the SSM improved all methods, while validation accuracies across initializations differed by over 15%.

5 Conclusion

S4 introduces a new state-space-model parameterization that efficiently models long-range dependencies through continuous-time, recurrent, and convolutional views. Results across diverse benchmarks suggest its potential as a general sequence-modeling solution.

  • 5 Conclusion: S4 introduces a new parameterization of the state space model for efficiently modeling long-range dependencies.The model is presented as a principled approach to long-range sequence modeling.
  • 5 Conclusion: S4 supports continuous-time, recurrent, and convolutional views of sequence modeling.These views are part of the model’s unified state-space formulation.
  • 5 Conclusion: Results across established benchmarks covering diverse data modalities and model capabilities suggest S4 could be an effective general sequence-modeling solution.The conclusion bases this assessment on evaluations across multiple established benchmarks.

A Discussion

The discussion situates S4 within prior biologically inspired state-space-model work, highlights broadly applicable parameterization and algorithmic contributions, and notes implementation choices and remaining limitations. In particular, S4 currently uses a GPU-parallelizable naive Cauchy-kernel algorithm, while language modeling still shows a gap to Transformers.

  • Related Work: S4’s parameterization and algorithms apply to a broader family of state-space models, and the associated techniques may be independently useful.The discussion explicitly extends applicability beyond the models investigated in prior work.
  • Implementation: S4’s training core is the Cauchy kernel, for which multiple algorithms offer different computational complexities and levels of sophistication.The kernel is discussed in Sections 3.2 and 3.3 and Appendix C.3.
  • Implementation: The current S4 implementation uses the naive O(NL) Cauchy-kernel algorithm because it is easily parallelized on GPUs and supported by accessible libraries.The implementation leverages the pykeops library.
  • Limitations and Future Directions: S4 is effective across varied data but may not suit every dataset, with Table 8 still showing a gap to Transformers for language modeling.The authors suggest combining S4 with other sequence models to complement their strengths.

B Numerical Instability of LSSL · B.1 HiPPO Diagonalization

The appendix establishes that the HiPPO diagonalization contains exponentially large entries, causing numerical instability and undermining a theoretically fast implementation. It then constructs the diagonalization by identifying eigenvectors and eigenvalues through an index-based induction.

  • B Numerical Instability of LSSL: The explicit HiPPO diagonalization confirms numerical instability because its entries grow exponentially.The same exponential terms also appear in the proposed theoretically fast algorithm, preventing implementation.
  • B.1 HiPPO Diagonalization: The HiPPO matrix is transformed, up to sign and diagonal conjugation, into a structured matrix whose entries follow an index-dependent formula.The displayed matrix and formula distinguish the cases n > k, n = k, and n < k.
  • B.1 HiPPO Diagonalization: The proposed diagonalizing matrix is specified explicitly, with its columns serving as eigenvectors of the HiPPO matrix.The proof’s goal is to show that these columns diagonalize A.
  • B.1 HiPPO Diagonalization: For each j, the j-th column v(j) is defined componentwise and is claimed to be an eigenvector with eigenvalue j + 1.The proof verifies the eigenvector equation for every index k ∈ [N].
  • B.1 HiPPO Diagonalization: When k < j, triangular sparsity makes both sides of the eigenvector equation equal to 0.This follows because either the relevant matrix entry A_ki or the vector component v(j)_i vanishes.
  • B.1 HiPPO Diagonalization: When k ≥ j, the proof expands the eigenvector equation and proceeds by induction on k.The base case k = j is immediate, while the inductive step follows by applying the hypothesis and simplifying factorial expressions.

B.2 Fast but Unstable LSSL Algorithm … C.2 Computing the S4 Recurrent View

The fast LSSL algorithm achieves near-linear operation and space complexity but is numerically unstable because it computes exponentially large intermediate terms. The S4 appendices establish efficient HiPPO NPLR/DPLR representations and an O(N) recurrent implementation via Woodbury-based discretization.

  • B.2 Fast but Unstable LSSL Algorithm: O(N log2 N + L log L) operations and O(N + L) space make the fast LSSL algorithm efficient, but it is numerically unstable.The instability is established despite the algorithm’s stated computational and memory bounds.
  • B.2 Fast but Unstable LSSL Algorithm: The instability arises from computing coefficients of p(x), the characteristic polynomial of A, and p(x)−1 (mod xL).The inverse polynomial can have coefficients larger than those of (1 − x)^N and increasing with L.
  • B.2 Fast but Unstable LSSL Algorithm: A = I exposes the instability: the fast LSSL algorithm requires computing terms exponentially large in N.This case is relevant because discretization can bring the transition matrix closer to the identity.
  • C.1 NPLR Representations of HiPPO Matrices: Theorem 1 shows that the HiPPO-LagT, HiPPO-LegT, and HiPPO-LegS continuous-time memory matrices admit normal plus low-rank representations.The constructions use low-rank corrections to obtain skew-symmetric or closely related diagonalizable forms.
  • C.2 Computing the S4 Recurrent View: The S4 recurrent view assumes A = Λ − P Q* is diagonal plus low-rank and derives a closed form for its discretized transition matrix.The backward Euler inverse is simplified using Woodbury’s Identity, while the forward term follows Euler discretization.
  • C.2 Computing the S4 Recurrent View: O(N) matrix-vector multiplication for both DPLR matrices A0 and A1 yields the efficient S4 recurrence and proves Theorem 2.The matrices are accessed only through matrix-vector multiplications.

C.3 Computing the Convolutional View

S4 computes the SSM convolution efficiently by replacing direct filter construction with generating-function evaluations, Woodbury correction, and Cauchy-kernel multiplication. The resulting algorithm recovers the filter stably in O(L log L) operations while using near-linear space and Cauchy multiplication cost.

  • Reduction 1: SSM Generating Function: S4 introduces an SSM generating function that replaces direct convolution-filter computation with evaluations while preserving information needed to recover the filter.The generating function converts the filter from the time domain to the frequency domain.
  • Reduction 1: SSM Generating Function: O(L log L) operations recover the truncated SSM convolution filter from generating-function evaluations at roots of unity using an inverse DFT and FFT.The evaluations at the roots of unity form a DFT of the filter coefficients.
  • Reduction 2: Woodbury Correction: The Woodbury identity converts generating-function evaluation for a diagonal-plus-low-rank A into operations involving inverses of diagonal matrices and a low-dimensional correction.This avoids inefficiently computing powers of the diagonal-plus-low-rank matrix.
  • Reduction 3: Cauchy Kernel: The diagonal-case generating function is exactly a Cauchy kernel, enabling fast and stable numerical algorithms for its evaluation.Cauchy matrices support fast arithmetic algorithms and numerical methods based on the Fast Multipole Method.
  • Reduction 3: Cauchy Kernel: O(L + N) space and C(L, N) operations suffice to evaluate the reduced expression, where C(L, N) = ˜O(L + N) is the Cauchy matrix-vector multiplication cost.Algorithm 1 is dominated by a constant number of black-box Cauchy-kernel calls.

D Experiment Details and Full Results

This section provides the full experimental procedures, extended results, and citations supporting the evaluation. Its appendices cover benchmarking, long-range dependency experiments, and general sequence modeling experiments.

  • Appendix D.1: Appendix D.1 reports the benchmarking results from Section 4.1.It contains the corresponding full procedures and extended results.
  • Appendix D.2: Appendix D.2 covers the long-range dependency experiments on Long Range Arena and Speech Commands from Section 4.2.These experiments are identified as the LRD evaluation.
  • Appendix D.3: Appendix D.3 covers general sequence modeling experiments on generation, image classification, and forecasting from Section 4.3.It provides the corresponding experimental details and extended results.

D.1 Benchmarking

The benchmarking setup evaluates S4 against LSSL and efficient Transformers under controlled architectural and parameter-matching conditions. Results from Tables 2 and 3 were tested on a single A100 GPU, while method rankings may depend on hyperparameter choices.

  • Tables 2 and 3 were evaluated on a single A100 GPU.
  • Benchmarks against LSSL: For LSSL and S4, single layers used H hidden features, with S4’s state size set to parameter-match LSSL’s and one forward+backward pass benchmarked.LSSL used state size N = H, while S4 used state size N = 4 due to parameterization differences.
  • Benchmarks against Efficient Transformers: Efficient Transformer comparisons used four layers, hidden dimension 256, four heads, projection dimension 128, batch size 32, and roughly 600k parameters.S4 kept depth and hidden dimension constant with parameter tying, yielding state size N = 256.
  • The relative ordering of methods can vary with the exact hyperparameters.

D.2 Long-Range Dependencies … D.4 Visualizations

The appendices detail reproducibility settings for long-range dependency and general sequence-modeling experiments, including task-specific architectures, optimization protocols, generation benchmarks, forecasting comparisons, and filter visualizations. S4 uses fixed or carefully tuned configurations across these evaluations and achieves the best results on 40 of 50 forecasting settings.

  • D.2 Long-Range Dependencies: The Long Range Arena experiments use AdamW with validation-plateau learning-rate decay, reduced HiPPO-parameter learning rates, fixed state size N = 64, and mean pooling for classification.HiPPO-related parameters use a maximum starting learning rate of 0.001, and all models run on a single GPU, with A100 hardware used notably for Path-X.
  • D.3 General Sequence Modeling: The appendix covers CIFAR-10 density estimation, WikiText-103 language modeling, autoregressive generation, sequential image classification, and time-series forecasting.These experiments correspond to the general sequence-modeling evaluations in Section 4.3.
  • D.3.1 CIFAR Density Estimation: CIFAR-10 density estimation uses alternating S4 and feed-forward layers in a UNet-style backbone, with softmax loss and input embeddings.The model uses LAMB with learning rate 0.005; the large model adds dropout 0.1 before linear layers.
  • D.3.2 WikiText-103 Language Modeling: WikiText-103 uses a 16-block S4 Transformer-style backbone with 1024-dimensional features, two S4 layers per block, adaptive embeddings, and sliding non-overlapping evaluation windows.Training uses AdamW, a single cosine cycle up to 800000 steps, 8 A100 GPUs, context size 8192, and dropout 0.25.
  • D.3.3 Autoregressive Generation Speed: Generation speed is benchmarked by throughput on a single A100 GPU, with batch sizes maximized for memory and CIFAR-10 generation limited to 16Gb.The reported baselines use cached implementations where available, while Sparse Transformer could not be run because of unresolved CUDA-kernel issues.
  • D.3.4 Pixel-Level Sequential Image Classification: Sequential image classification trains S4 models with AdamW for up to 200 epochs and compares default LayerNorm against ResNet-18’s BatchNorm, including augmentation and normalization ablations.The augmentation consists of horizontal flips and random crops with symmetric padding.
  • D.3.5 Time Series Forecasting compared to Informer: 40 out of 50 forecasting settings favor S4, which concatenates past context with masked forecast positions and uses the final F outputs as predictions.This simple sequence-to-sequence procedure contrasts with Informer’s specialized encoder-decoder architecture and customized attention module.
  • D.4 Visualizations: The appendix visualizes S4’s learned convolutional filter K̄ for the Pathfinder and CIFAR-10 tasks.These visualizations are provided in Appendix D.4.

D.5 Reproduction

The reproduction section documents updated results and evaluation conventions for Long Range Arena, Speech Commands, and WikiText-103. These updates mainly reflect revised baselines, dataset scope, and training configurations while preserving the S4 model.

  • Speech Commands: Speech Commands evaluation now distinguishes the 10-class subset SC10 from the original 35-class dataset.For an S4 baseline, the paper directs readers to Table 11 from [19], which uses the full dataset and stronger baselines.
  • WikiText-103: 20.95 test perplexity was achieved on WikiText-103 after retraining with batch size 1 and context size 8192.The retrained model also achieved validation perplexity 19.69, compared with the original test perplexity of 21.28 and validation perplexity of 20.88.
  • WikiText-103: The updated WikiText-103 model is essentially identical to the original, so the original table results were updated.A model checkpoint is available in the public repository.
Loading 2111.00396v3…