Structured State Space Models

This is my seminar paper on structured state space models for the “Selected Topics in Machine Learning Research” seminar organized by the DAML group at TUM. Download the original seminar paper (PDF).

Introduction

Large-scale sequence modeling becomes more and more relevant. The most prominent example is language modeling, which is largely dependent on scaling up the context windows models can take into account. Existing architectures are inherently flawed with respect to that problem. Transformers scale quadratically in the sequence length, resulting in slow inference and high resource requirements. RNNs suffer from poor performance, mainly caused by their inability to retain long-term context in the hidden state. Structured state space models are a promising alternative because they offer two alternative formulations: A recurrent formulation that allows for efficient inference, and a parallel formulation based on convolutions that allows for parallel training and circumvents gradient problems (Gu et al., 2022). Additionally, structured SSMs are a theoretically stronger model family for continuous input signals. Transformers and classical RNNs are designed for discrete sequences, having no notion of continuity. In contrast, structured SSMs are based on mathematical foundations that explicitly model the input sequence as a continuous signal. That allows them to handle varying sampling intervals, even within single sequences, without any modifications.

My main contribution is a gentle introduction to structured state space models and their advantages and disadvantages. Additionally, I provide an overview over applications and corresponding state-of-the-art architectures, and give my prediction for further developments in the field.

Background

Online Function Approximation

Main Idea In machine learning, the main idea of online function approximation is to provide models with mathematically principled, compressed representations of sequence histories. The first architecture that followed this approach was Legendre Memory Units (LMUs) (Voelker et al., 2019), but in the following I will introduce a more general solution.

HiPPO Framework The HiPPO framework (Gu et al., 2020) generalizes LMUs to a more abstract framework for online function approximation. It takes in a continuous input signal $u(t)$ and aims at approximating it with respect to a probability measure $\mu(t)$. The measure can be conceptualized as a weight function for the approximation error. Different instantiations of the HiPPO framework refer to different choices of measure. For instance, picking $\mu(t)$ as a window with limited context and uniform value yields LMUs. Since the approximation is to be computed online, meaning only $u_{\leq \tau}$ is known at time $\tau$, the measure’s support $(-\infty, \tau]$ varies. As a consequence, the measure $\mu^{(\tau)}(t)$ is a function of time $t$ and parameterized by the current point in time $\tau$.

The compression happens by projecting $u(t)$ onto a polynomial function basis $\{ g_n \}_{n < N}$. The basis functions are chosen such that they are orthogonal with respect to the measure, meaning ${\forall n, m < N : \langle g_n, g_m \rangle_{\mu^{(t)}} = \delta_{n,m}}$ where the inner product with respect to any measure is defined as ${ \langle f, g \rangle_\mu = \int_0^\infty f(t) g(t) d\mu(t) }$. The compressed representation of the input function history up to $\tau$ is the basis functions’ coefficients, which are computed by projecting the input function onto the respective basis function via inner product: ${ c_n^{(\tau)} = \langle u_{\leq \tau}, g_n \rangle_{\mu^{(\tau)}} }$. Recalling that the goal is to compute the compressed representation $c(t) \in \mathbb{R}^N$ online, the idea now is to differentiate the coefficients with respect to time. It turns out that $\frac{d}{dt}c(t)$ approximately evolves as a system of linear ordinary differential equations of the form

where $A(t) \in \mathbb{R}^{N \times N}$ and $B(t) \in \mathbb{R}^{N \times 1}$ can be derived in closed-form for the respective measure. While the operators $A, B$ depend on time in the general case, (Gu et al., 2020) show that they are time-invariant for all HiPPO instantiations they work with. For the actual computation, the dynamics have to be discretized, which means they take the general form ${c_{k+1} = A_k c_k + B_k u_k}$. This introduces a discretization step size parameter $\Delta t$, which is crucial for performance.

Theoretically, the approximation of $u(t)$ can now be restored from the coefficients as $\hat{u}(t) = \sum_{n < N} c_n(t) g_n(t)$. The plots below visualize the reconstructed signal for two different instantiations of the HiPPO framework (Gu et al., 2020). That being said, in practice the approximation is not actually realized. Instead, the models work directly on the compressed representation of the approximated signal history $c(t)$.

Plots of input signal u(t) (black), HiPPO reconstruction û(t) (red), and measure function μ(t) (green) for the HiPPO instantiation LegS (exponentially decaying measure), and LegT (uniform, windowed measure). Visualization code: state-spaces/s4.

State Space Models

(Gu et al., 2021) introduce the Linear State Space Layer (LSSL), the first (non-structured) SSM that incorporates ideas from the HiPPO framework. It can be written as

where $u(t) \in \mathbb{R}$ is the input signal, $x(t) \in \mathbb{R}^{N}$ is the state, $y(t) \in \mathbb{R}^{M}$ is the output, and $A \in \mathbb{R}^{N \times N}, B \in \mathbb{R}^{N}, C \in \mathbb{R}^{M \times N}, D \in \mathbb{R}^{M}$ are learnable parameter matrices. $Du(t)$ acts as a projected residual connection and is irrelevant for any further derivations, which is why I omit it from now on.

A big difference to other sequence models is that the LSSL layer’s continuous dynamics are manually discretized using a timescale parameter $\Delta t$. That makes it a good choice for settings such as dealing with irregularly spaced time series. Note that for sequences that are not based on a continuous signal, $\Delta t$ can also be learned as a model parameter.

With random initialization, the LSSL does not achieve any noteworthy performance. But when provided with the compressed input function history using ideas from the HiPPO framework, it becomes a strong model. Its state dynamics take the same form as the HiPPO coefficient dynamics above. That allows to integrate the HiPPO framework by initializing $A$ and $B$ with the derived HiPPO matrix and bias for the chosen instantiation. It should be noted that the state dynamic does not stay fixed after the initialization. $A$ and $B$ are further optimized during training of the layer.

Finally, I address how to model multichannel sequences with LSSLs. The definition above describes a map $\mathbb{R} \rightarrow \mathbb{R}^{M}$ for each sequence element, but what is needed for sequences with $H$ channels is a map $\mathbb{R}^{H} \rightarrow \mathbb{R}^{H}$. (Gu et al., 2021) achieve this by training one LSSL per input-channel, and combining the parallel LSSL layers with a position-wise MLP $\mathbb{R}^{H \cdot M} \rightarrow \mathbb{R}^{H}$.

Structured State Space Models

Efficient Computation as Convolution

Problem Primitive state space models such as LSSL-based architectures have advantages such as dealing well with irregularly spaced sequences, and allowing for dynamic selection of discretization interval during training and inference. But fundamentally they are still recurrent architectures, which means they suffer from the typical problems such as slow training and gradient issues during training (Gu et al., 2022). That also impacts their performance on long sequences, even though the HiPPO framework integration provides a theoretically strong fundament for good long-range performance.

Convolutional Formulation (Gu et al., 2022) introduce an alternative way to compute SSM outputs for all sequence positions at the same time. This enables fast, parallel training and eliminates gradient problems during training caused by unrolling the recurrent computational graph. Note that this formulation is not mutually exclusive with the recurrent formulation. You can train in parallel using the following convolutional formulation, but still do efficient, recurrent inference.

Plugging the LSSL’s recurrent formulation into itself a few times yields a closed-form solution for the output after several recurrent steps

where for instance $\overline{A}$ is the discretized version of $A$. Therefore, it can be computed as a single long convolution ${ y = \overline{K} \ast u }$, meaning a convolution with a kernel ${ \overline{K} = \big( \overline{CA^iB} \big)_{i \in [L]} }$ of the same length $L$ as the input signal.

The convolution itself can be computed efficiently in the Fourier domain, but the kernel is non-trivial to compute due to the powers of $A$ (Gu et al., 2022). The following approaches aim at finding an efficient way to compute the kernel by enforcing certain structures of $A$, thereby introducing the model family of Structured SSMs.

Normal Plus Low-Rank (S4) Performing the conjugation ${ (A, B, C) \sim (V^{-1}AV, V^{-1}B, CV) }$ does not change the output an SSM computes, it only performs a base change of the state. That means one can conjugate $A$ without changing the SSMs behavior. In principle, that allows to diagonalize the HiPPO matrix via conjugation ${ A = V^{-1}\Lambda V }$, which would make computing the convolution kernel trivial and computationally inexpensive. Unfortunately, this primitive diagonalization is numerically unstable for larger states, because $V$ would have entries exponential in the state size.

Instead, (Gu et al., 2022) exploit that HiPPO matrices can be decomposed as the sum of a normal matrix and a low-rank correction ${ A = N - PQ^T = V \Lambda V^\ast - PQ^T }$ where ${ V \in \mathbb{C}^{N \times N} }$ is unitary, $\Lambda$ is diagonal, and ${ P,Q \in \mathbb{R}^{N \times r} }$ form a low-rank factorization. Furthermore, the normal plus low-rank form can be turned into a diagonal plus low-rank form by the conjugation ${ V \Lambda V^\ast - PQ^T = V(\Lambda - (V^\ast P)(V^\ast Q)^\ast)V^\ast }$.

Using the diagonal plus low-rank decomposition, the convolution kernel $\overline{K}$ can be efficiently evaluated using a few tricks. I briefly sketch (Gu et al., 2022)’s approach: 1) Instead of computing $\overline{K}$ directly, they compute its spectrum by evaluating its truncated generating function at the roots of unity. From this representation, $\overline{K}$ can be efficiently recovered using the inverse Fourier transform. 2) Evaluating the generating function does not involve a power of $\overline{A}$ anymore, but an inverse. There is an identity called Woodbury matrix identity, which reduces the problem of computing the inverse of any matrix plus low-rank correction $(M + LR^\ast)^{-1}$ to computing the inverse of the matrix $M^{-1}$. That can be used to reduce computing $\overline{A^{-1}}$ to the final, truly diagonal case. 3) Computing the diagonal case is equivalent to the computation of so-called Cauchy kernels, which is a common problem with stable, near-linear time algorithms.

The resulting structured SSM is called S4. As already mentioned, it can be evaluated with two different algorithms. S4 Recurrence is the recurrent formulation and one step can be evaluated in $\mathcal O(N)$ operations. S4 Convolution can be used to evaluate the convolution kernel in $\mathcal O(N + L)$ operations and space, which can then be applied in $\mathcal O (L \cdot \mathrm{log} \, L)$. Note that this is a significant advantage over transformers, which are quadratic in the sequence length $L$.

Diagonal Further work such as S4D (Gu, Goel, Gupta, et al., 2022), DSS (Gupta et al., 2022), and GSS (Mehta et al., 2023) shows that dropping the low-rank correction term does not heavily influence performance and allows for simpler and more efficient computation. The resulting diagonal structured SSMs form the basis of most advanced structured SSM architectures, including the ones discussed in further sections.

Addressing Input-Selectivity

Problem While structured SSMs have desirable properties with respect to computation and efficiency, their performance on discrete sequence modeling tasks is not on par with transformers yet. Their core problem is often referred to as missing (input-)selectivity, which means the parameters $A, B, C, D, \Delta$, and therefore the computed operator, are independent of the input. Consequently, SSMs fail even on simple synthetic tasks such Induction Head (retrieving the token after special token), Associative Recall (retrieving a specific value from key-value input) (Fu et al., 2023), and Selective Copying (copying a sequence to unevenly spaced locations marked by special tokens) (Arjovsky et al., 2016). In the following, I present three architectures that address this issue. H3 (Fu et al., 2023) and Hyena (Poli et al., 2023) try to achieve some input-selectivity by designing an architecture around ”classical” structured SSMs. Mamba (Gu & Dao, 2023) goes one step further and explicitly turns some parameters into functions of the input.

H3 (Fu et al., 2023) introduce the H3 architecture, which is specifically designed to solve Induction Head and Associative Recall and achieves a large jump in performance for both of them. Both of these tasks require the ability to perform an operation on a token based on the previous token’s value. H3 starts off by computing 3 projections $Q, K, V$ of the current input token $u$, which are loosely related to the projections of the attention mechanism. Then, the layer output is calculated as

which can be conceptualized as two steps: 1) $K$ is fed through a so-called Shift-SSM, an SSM with $A$ fixed to shift the state one position ”down” ($x_0^{(t)}$ becomes $0$, $x_1^{(t)}$ becomes $x_0^{(t)}$ and so on). Even though $B$ is learned in practice, from assuming it to be fixed to the first basis vector $e_1$ and recalling the LSSL dynamics above, one can see that this keeps a local history of the scalar SSM inputs in the state vector. The Shift-SSM’s output interacts with $V$ via Hadamard product, thereby relating the previous token’s $K$ to the current token’s $V$. 2) The output of this operation is fed through a ”normal” diagonal SSM, and the final layer output is computed as Hadamard product with $Q$.

Additionally, the authors introduce a hardware-optimized CUDA kernel called FlashConv, that speeds up the convolution computation by fusing operations in the SRAM and only then writing back to HBM, similar to the well-known FlashAttention (Dao et al., 2022).

Hyena By introducing Hyena, (Poli et al., 2023) generalize the concept of the H3 layer, thereby drastically improving the performance on recall and reasoning tasks such as the ones mentioned above. To demonstrate the analogy to H3, I first provide a different view on SSMs and H3: SSMs effectively compute an implicitly parameterized long convolution, which means the convolution kernel’s elements are functions of the actual model parameters. H3’s convolutions and Hadamard products can be rewritten as matrix multiplications ${ \mathrm{H3}(q, k, v) = D_q S_\psi D_k S_\varphi \cdot v }$ where ${ D_x = \mathrm{diag}(x) }$ and $S_\psi, S_\varphi$ are Töplitz matrices for the convolution kernels realized by the diagonal SSM and Shift-SSM respectively. The Hyena layer generalizes that to ${ y = D_x^{N_H} S_h^{N_H} \dots D_x^1 S_h^1 \cdot v }$. That means H3 is an instantiation of Hyena with ${N_H}=2$ and a constrained first kernel. Of course, in practice all these matrices are not materialized, but computed by convolutions in the Fourier domain and Hadamard products.

Strictly speaking, the Hyena framework is not just an SSM architecture, but a more general framework of interleaved implicitly parameterized long convolutions and data-controlled gating by Hadamard products. (Poli et al., 2023) compute their kernels as ${ h_t = \mathrm{Window}(t) \cdot \mathrm{MLP}\big(\mathrm{PositionalEncoding}(t)\big) }$ where $\mathrm{Window}$ is a weighting function for the sequence positions, and most commonly chosen as an exponential decay. It is still appropriate to view Hyena as an architectural framework for SSMs. Since SSMs also realize a convolution with an implicit kernel, the Hyena convolutions with the kernels specified above can simply be replaced with SSM layers.

Mamba H3 and Hyena try to achieve some input selectivity by designing an architecture around classical SSMs. With Mamba, (Gu & Dao, 2023) go one step further and actually compute the SSM parameters $\Delta, B, C$ as learnable functions of the input. This results in a massive boost in performance, with Mamba being the first SSM-based model to solve the Selective Copying task and being the first to match the performance of transformers on language modeling.

This performance increase is no free lunch. Since the parameters are now not only functions of the time but also of the input, the convolutional formulation is not possible anymore. To mitigate this, (Gu & Dao, 2023) introduce a hardware-optimized scan algorithm that fuses the discretization and a limited number of recurrence steps to a single CUDA kernel, preventing most IO operations between SRAM and HBM.

Applications & Performance

Long-Range Sequence Modeling Since SSMs scale subquadratically with the sequence length, they allow for deep models, even when large context windows are required. Based on the architectures presented in this seminar paper, there is follow-up work that specifically tackles long range tasks. (Hasani et al., 2023) combine SSMs with ideas from liquid time-constant networks (Hasani et al., 2021). (Gu et al., 2023) derive an alternative formulation of the HiPPO framework and introduce novel instantiations. (Smith et al., 2023) fuse multiple S4 layers to a multi-input multi-output S5 layer. The following table provides an overview over Long Range Arena (LRA) (Tay et al., 2021) benchmark scores. I also include Big Bird (Zaheer et al., 2020), a strong transformer baseline explicitly designed to handle long-range dependencies. The comparison shows that SSM-based architectures perform significantly better than transformer-based architectures.

All Long Range Arena results in this comparison are from (Hasani et al., 2023).

Language Modeling The table below shows results of SSMs and a transformer baseline on the WikiText-103 (Merity et al., 2017) language modeling task. For these smaller models, SSM-based models designed for input-selectivity start to exceed transformer-baselines. Experiments with large-scale models are not publicly available yet.

Other Applications Potential applications of SSMs are manifold. (Qiao et al., 2024) introduce a Mamba-based multimodal LLM. (Zhu et al., 2024) present a drop-in replacement for vision transformers. (Liu et al., 2024; Ruan & Xiang, 2024) perform 2D image segmentation, (Xing et al., 2024) 3D volume segmentation, and (Yang et al., 2024) 4D volume-video segmentation. (Fei et al., 2024) train a diffusion model based on SSMs. (Liang et al., 2024) introduce a 3D point cloud processing backbone. SSMs are also applied to other domains, such as speech (Goel et al., 2022; Zhang et al., 2024), time series (Patro & Agneeswaran, 2024; M. Zhang et al., 2023), and graphs (Behrouz & Hashemi, 2024; Wang et al., 2024).

Conclusion

Structured SSMs allow for fast parallel training and efficient recurrent inference, which makes them a promising model family for sequence modeling. The fact that most state-of-the-art models rely on Mamba as a backbone indicates that input-selectivity is crucial for a lot of tasks, though. Consequently, I suspect ”pure” linear time invariant SSMs will only be used for niche applications such as irregularly spaced sequences in the future. That does not make SSMs useless for tasks that require input-selectivity, such as language modeling. The trend in language modeling seems to go further into the direction of scaling up models as much as possible, which is why efficiency becomes more and more relevant. Replacing transformer layers by SSM layers allows researchers to make fine-grained trade-offs between expressivity and efficiency for individual layers, potentially resulting in higher performance of the full models.

References