Simplified State Space Layers for Sequence Modeling

By Jimmy Smith, Andrew Warrington, Scott W. Linderman. International Conference on Learning Representations, 2022

1,076 citations recorded by the source corpus.

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.

Subfields: Computer Networks and Communications, Computer Vision and Pattern Recognition, Hardware and Architecture

View Paper: https://www.semanticscholar.org/paper/6d7d141c75af752ffc0d8a6184cca3f9323d6c74

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Computer Networks and Communications

In Computer Vision and Pattern Recognition

In Hardware and Architecture

Citations in the corpus

VideoMamba: State Space Model for Efficient Video Understanding (2024, neutral): “Based on it, various models have been developed ( e.g. , S5 [67], H3 [21], and GSS [57]).”
Resurrecting Recurrent Neural Networks for Long Sequences (2023, methodology): “, 2022a;b) have suggested that the specific initialization used by S4 is often not crucial for performance, and that the discretization rules which achieve best performance may not be the most accurate in theory (Smith et al., 2022).”
LocalMamba: Visual State Space Model with Windowed Selective Scan (2024, neutral): “Despite initial challenges in training, owing to their computational and memory intensity, recent advancements [9–11, 16, 43] have significantly ameliorated these issues, positioning deep SSMs as formidable competitors against CNNs and Transformers.”
MambaHSI: Spatial–Spectral Mamba for Hyperspectral Image Classification (2025, methodology): “Smith et al. [50] designed a S5 layer by introducing MIMO SSM and efficient parallel scan into S4 layer.”
MEGABYTE: Predicting Million-byte Sequences with Multiscale Transformers (2023, background): “…which typically involve forms of token-level recurrence (Katharopoulos et al., 2020; Schlag et al., 2021) or state space models (Gu et al., 2021; Smith et al., 2022; Ma et al., 2022), or (3) sparse approximations of attention (Kitaev et al., 2020; Beltagy et al., 2020; Child et al., 2019; Wu et…”
ChangeMamba: Remote Sensing Change Detection With Spatiotemporal State Space Model (2024, methodology): “Based on S4, Smith et al. [46] proposed a new S5 model by introducing MIMO SSM and efficient parallel scan into the S4 model.”
VMamba: Visual State Space Model (2024, background): “State Space Models (SSMs) are recently proposed models that are introduced into deep learning as state space transforming [18, 17, 43].”
PointMamba: A Simple State Space Model for Point Cloud Analysis (2024, methodology): “The S5 [48] proposes a parallel scan and the MIMO SSM, enabling the state space model to be efficiently utilized and widely implemented.”
Is Mamba Effective for Time Series Forecasting? (2024, neutral): “Transformer [52] exhibits”
Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model (2024, neutral): “[53] proposes a new S5 layer by introducing MIMO SSM and efficient parallel scan into S4 layer.”
Parallelizing Linear Transformers with the Delta Rule over Sequence Length (2024, neutral): “As is the case in vector-valued linear RNNs [ 62, 103], the use of an associative operator enables the use of parallel scan [13] to calculate S 1 , . . . , S L in O (log L ) steps and O ( L ) work (ignoring the terms associated with the associative operation) if the inputs x 1 , . . . , x L are…”
Gated Linear Attention Transformers with Hardware-Efficient Training (2023, background): “Experiments on language modeling demonstrate that gated linear attention Transformers can perform respectably compared to strong baselines.”
MambaIR: A Simple Baseline for Image Restoration with State-Space Model (2024, methodology): “Recently, structured state-space sequence models (S4), especially the improved version Mamba, have emerged as an efficient and effective backbone for constructing deep networks [18,22,24,52,62].”
Benchmarking and Building Long-Context Retrieval Models with LoCo and M2-BERT (2024, methodology): “We also describe existing state-of-the-art approaches for retrieval models and compare them to M2-BERT.”