On the Parameterization and Initialization of Diagonal State Space Models
Abstract
State space models (SSM) have recently been shown to be very effective as a deep learning layer as a promising alternative to sequence models such as RNNs, CNNs, or Transformers.
The first version to show this potential was the S4 model, which is particularly effective on tasks involving long-range dependencies by using a prescribed state matrix called the HiPPO matrix.
While this has an interpretable mathematical mechanism for modeling long dependencies, it introduces a custom representation and algorithm that can be difficult to implement.
On the other hand, a recent variant of S4 called DSS showed that restricting the state matrix to be fully diagonal can still preserve the performance of the original model when using a specific initialization based on approximating S4’s matrix.
This work seeks to systematically understand how to parameterize and initialize such diagonal state space models.
While it follows from classical results that almost all SSMs have an equivalent diagonal form, we show that the initialization is critical for performance.
We explain why DSS works mathematically, by showing that the diagonal restriction of S4’s matrix surprisingly recovers the same kernel in the limit of infinite state dimension.
We also systematically describe various design choices in parameterizing and computing diagonal SSMs, and perform a controlled empirical study ablating the effects of these choices.
Our final model S4D is a simple diagonal version of S4 whose kernel computation requires just 2 lines of code and performs comparably to S4 in almost all settings, with state-of-the-art results for image, audio, and medical time-series domains, and averaging 85% on the Long Range Arena benchmark.
Introduction
A core class of models in modern deep learning are sequence models, which are parameterized mappings operating on arbitrary sequences of inputs.
Recent approaches based on state space models (SSMs) have outperformed traditional deep sequence models such as recurrent neural networks (RNNs), convolutional neural networks (CNNs), and Transformers, in both computational efficiency and modeling ability.
In particular, the S4 model displayed strong results on a range of sequence modeling tasks, especially on long sequences [9].
Its ability to model long-range dependencies arises from being defined with a particular state matrix called the “HiPPO matrix” [6], which allows S4 to be viewed as a convolutional model that decomposes an input onto an orthogonal system of smooth basis functions[10].
However, beyond its theoretical interpretation, actually computing S4 as a deep learning model requires a sophisticated algorithm with many linear algebraic techniques that are difficult to understand and implement.
These techniques were necessitated by parameterizing its state matrix as a diagonal plus low-rank (DPLR) matrix, which is necessary to capture HiPPO matrices.
A natural question is whether simplifications of this parameterization and algorithm are possible.
In particular, removing the low-rank term would result in a diagonal state space model (diagonal SSM) that is dramatically simpler to implement and understand.
Although it is known that almost all SSMs have an equivalent diagonal form—and therefore (complex) diagonal SSMs are fully expressive algebraically—they may not represent all SSMs numerically, and finding a good initialization is critical.
Gu et al.
[9] showed that it is difficult to find a performant diagonal SSM, and that many alternative parameterizations of the state matrix – including by random diagonal matrices – are much less effective empirically, which motivated the necessity of the more complicated HiPPO matrix.
However, recently Gupta [11] made the empirical observation that a variant of S4 using a particular diagonal matrix is nearly as effective as the original S4 method.
This matrix is based on the original HiPPO matrix and is defined by simply chopping off the low-rank term in the DPLR representation.
The discovery of performant diagonal state matrices opens up new possibilities for simplifying deep state space models, and consolidating models such as S4 and DSS to understand and improve them.
First, the strongest version of DSS computes the SSM with a complex-valued softmax that complicates the algorithm, and is actually less efficient than S4. Additionally, DSS and S4 differ in several auxiliary aspects of parameterizing SSMs that can conflate performance effects, making it more difficult to isolate the core effects of diagonal versus DPLR state matrices. Most importantly, DSS relies on initializing the state matrix to a particular approximation of S4’s HiPPO matrix.
While
S4’s matrix has a mathematical interpretation for addressing long-range dependencies, the efficacy of the diagonal approximation to it remains theoretically unexplained.
In this work, we seek to systematically understand how to train diagonal SSMs.
We introduce the S4D method, a diagonal SSM which combines the best of S4’s computation and parameterization and DSS’s initialization , resulting in a method that is extremely simple, theoretically princpled, and empirically effective. • First, we describe S4D, a simple method outlined by S4 for computing diagonal instead of DPLR matrices, which is based on Vandermonde matrix multiplication and is even simpler and more efficient than the DSS.
Outside of the core state matrix, we categorize different representations of the other components of SSMs, introducing flexible design choices that capture both S4 and DSS and allow different SSM parameterizations to be systematically compared (Section 3).
• We provide a new mathematical analysis of DSS’s initialization, showing that the diagonal approximation of the original HiPPO matrix surprisingly produces the same dynamics as S4 when the state size goes to infinity.
We propose even simpler variants of diagonal SSMs using different initializations of the state matrix (Section 4).
• We perform a controlled study of these various design choices across many domains, tasks, and sequence lengths, and additionally compare diagonal (S4D) versus DPLR (S4) variants.
Our best S4D methods are competitive with S4 on almost all settings, with near state-of-the-art results on
image, audio, and medical time series benchmarks, and achieving 85% on the Long Range Arena benchmark (Section 5).
Background
Continuous State Spaces Models S4 investigated state space models (1) that are parameterized maps on signals u ( t ) (cid:55)→ y ( t ) . These SSMs are linear time-invariant systems that can be represented either as a linear ODE (equation (1)) or convolution (equation (2)).
[Equation]
Here the parameters are the state matrix A ∈ C N × N and
other matrices B ∈ C N × 1 , C ∈ C 1 × N .
In the case of diagonal SSMs, A is diagonal and we will overload notation so that A n , B n , C n denotes the entries of the parameters.
An intuitive way to view the convolution kernel (2) is to interpret it as a linear combination (controlled by C ) of basis kernels K n ( t ) (controlled by A , B )
[Equation]
We denote this basis as K ( t ) = K A , B ( t ) =
e t A B if necessary to disambiguate; note that it is a vector of N functions.
In the case of diagonal SSMs, each function K n ( t ) is just e t A n B n .
S4: Structured State Spaces
As a deep learning model, SSMs have many elegant properties with concrete empirical and computational benefits [8].
For example, the convolutional form (2) can be converted into a temporal recurrence that is substantially faster for autoregressive applications [5].
However, making SSMs effective required overcoming two key challenges: choosing appropriate values for the matrices, and computing the kernel (2) efficiently.
First, Gu et al.
[8] showed that naive instantiations of the SSM do not perform well, and instead relied on a particular (real-valued) matrix A called the HiPPO-LegS matrix (4).
These matrices were derived so that the basis kernels K n ( t ) have closed-form formulas L n ( e − t ) , where L n ( t ) are normalized Legendre polynomials. Consequently, the SSM has a mathematical interpretation of decomposing the input signal u ( t ) onto a set of infinitely-long basis functions that are orthogonal respect to an exponentially-decaying measure, giving it long-range modeling abilities [10].
Second, S4 introduced a particular parameterization that decomposed this A matrix into the sum of a normal and rank-1 matrix (5), which can be unitarily conjugated into a (complex) diagonal plus rankmatrix.
Leveraging this structured form, they then introduced a sophisticated algorithm for efficiently computing the convolution kernel (2) for state matrices that are diagonal plus low-rank (DPLR) .
[Equation]
HiPPO-LegS matrix used in S4 )
[Equation] [Equation]
DSS: Diagonal State Spaces S4 was originally motivated by searching for a diagonal state matrix , which would be even more structured and result in very simple computation of the SSM.
However, the HiPPO-LegS matrix cannot be stably transformed into diagonal form [9, Lemma 3.2], and they were unable to find any diagonal matrices that performed well, resulting in the DPLR formulation.
Gupta [11] made the surprising empirical observation that simply removing the low-rank portion of the DPLR form of the HiPPO-LegS matrix results in a diagonal matrix that performs comparably to the original S4 method.
More precisely, their initialization is the diagonal matrix A ( D ) , or the diagonalization of A ( N ) in (5).
They termed A ( N ) the skew-HiPPO matrix, which we will also call the normal-HiPPO matrix. To be more specific and disambiguate these variants, we may also call A ( N ) the HiPPO-LegS-N or HiPPO-N matrix and A ( D ) the HiPPO-LegS-D or HiPPO-D matrix. In addition to this initialization, they proposed a method for computing a diagonal SSM kernel.
Beyond these two core differences, several other aspects of their parameterization differ from S4’s. In Sections 3 and 4, we systematically study the components of DSS: we categorize different ways to parameterize and compute the diagonal state space, and explain the theoretical interpretion of this particular diagonal A matrix.
Parameterizing Diagonal State Spaces
We describe various choices for the computation and parameterization of diagonal state spaces.
Our categorization of these choices leads to simple variants of the core method.
Both DSS and our proposed S4D can be described using a combination of these factors (Section 3.4).
3.1 Discretization
The true continuous-time SSM can be represented as a continuous convolution y (
[Equation] [Equation]
In discrete time, we view an input sequence u 0 , u 1 , . . .
as uniformly-spaced samples from an underlying function u ( t ) and must approximate this integral.
Standard methods for doing so that preserve the convolutional structure of the model exist.
The first step is to discretize the parameters.
Two simple choices that have been used in prior work include
[Equation]
With these methods, the discrete-time SSM output is just
[Equation]
where
[Equation]
These integration rules have both been used in prior works (e.g. LMU and DSS use ZOH [26, 11] while S4 and its predecessors use bilinear [6, 8, 9]).
In Section 5, we show that there is little empirical difference between them.
However, we note that there is a curious phenomenon where the bilinear transform actually perfectly smooths out the kernel used in DSS to match the S4 kernel (Section 4 Fig. 2d).
We additionally note that numerical integration is a rich and well-studied topic and more stable methods of approximating the convolutional integral may exist.
For example, it is well-known that simple rules like the Trapezoid rule [18] can dramatically reduce numerical integration error when the function has bounded second derivative.
3.2 Convolution Kernel
The main computational difficulty of the original S4 model is computing the convolution kernel K .
This is extremely slow for general state matrices A , and S4 introduced a complicated algorithm for
DPLR state matrices.
When A is diagonal, the computation is nearly trivial. By (6),
K ℓ
[Equation]
where V L ( A ) n,ℓ = A ℓn (7) where ◦ is Hadamard product, · is matrix multiplication, and V is known as a Vandermonde matrix .
Time and Space Complexity
The naive way to compute (7) is by materializing the Vandermonde matrix V L ( A ) and performing a matrix multiplication, which requires O ( NL ) time and space.
However, Vandermonde matrices are well-studied and theoretically the multiplication can be computed in (cid:101)
O
( N + L ) operations and O ( N + L ) space.
In fact, Vandermonde matrices are closely related to Cauchy matrices, which are the computational core of S4’s DPLR algorithm, and have identical complexity [17]. Proposition 1. The time and space complexity of computing the kernel of diagonal SSMs is equal to that of computing DPLR SSMs. We note that on modern parallelizable hardware such as GPUs, a simple fast algorithm is to compute (7) with naive summation (using O ( NL ) operations), but without materializing the Vandermonde matrix (using O ( N + L ) space).
Just as with S4, this may require implementing a custom kernel in some modern deep learning frameworks such as PyTorch to achieve the space savings.
3.3 Parameterization
Parameterization of A .
Note that the kernel K ( t ) =
C e t A B blows up to ∞ as t → ∞ if A has any eigenvalues with positive real part.
Goel et al.
[5] found that this is a serious constraint that affects the stability of the model, especially when using the SSM autoregressively.
They propose to force the real part of A to be negative, also known as the left-half plane condition in classical controls, by parameterizing the real part inside an exponential function A = − exp ( A Re ) + i · A Im .
We note that instead of exp , any activation function can be used as long as its range is bounded on one side, such as ReLU, softplus, etc.
The original DSS does not constrain the real part of A , which is sufficient for simple tasks involving fixed-length sequences, but could become unstable in other settings.
Parameterization of B , C .
Another choice in the parameterization is how to represent B and
C . Note that the computation of the final discrete convolution kernel K depends only on the elementwise product B ◦ C (equation (7)).
Therefore DSS chose to parameterize this product directly, which they call W , instead of B and C individually. However, we observe that this is equivalent to keeping independent B and C , and simply freezing
B = 1 while training C .
Therefore, just as S4 has separate parameters A , B , and C and uses a fixed initialization for A and B , S4D also proposes separate A , B , and C and uses fixed initializations for
A (discussed in Section 4) and B (set to 1 ).
Then the difference between S4D and DSS is simply that DSS does not train B .
In our ablations, we show that training B gives a minor but consistent improvement in performance.
3.4 S4D: the Diagonal Version of S
A key component of our exposition is disentangling the various choices possible in representing and computing state space models.
With this categorization, different choices can be mixed and matched to define variants of the core method.
Table 1 compares S4, DSS, and S4D, which have a core structure and kernel computation, but have various choices of other aspects of the parameterization.
[Table]
[Table]
Initialization of Diagonal State Matrices
The critical question remains: which diagonal state matrices A are actually effective?
We comment on the limitations of diagonal SSMs, and then provide three instantiations of S4D that perform well empirically.
Expressivity and Limitations of Diagonal SSMs.
We first present a simplified view on the expressivity of diagonal SSMs mentioned by [11].
First, it is well-known that almost all matrices diagonalize over the complex plane.
Therefore it is critical to use complex-valued matrices in order to use diagonal SSMs.
Proposition 2.
The set D ⊂ C N × N of diagonalizable matrices is dense in C N × N , and has full measure (i.e. its complement has measure 0 ).
It is also well known that the state space ( A , B , C ) is exactly equivalent to (i.e. expresses the same map u (cid:55)→ y )
the state space ( V − 1 AV , V − 1 B , CV ) , known in the SSM literature as a state space transformation.
Therefore Proposition 2 says that (almost) all SSMs are equivalent to a diagonal SSM .
However, we emphasize that Proposition 2 is about expressivity which does not guarantee strong performance of a trained model after optimization.
For example, Gu et al.
[9] and Gupta [11] show that parameterizing A as a dense real matrix or diagonal complex matrix, which are both fully expressive classes, performs poorly if randomly initialized.
Second, Proposition 2 does not take into account numerical representations of data, which was the original reason S4 required a low-rank correction term instead of a pure diagonalization [9, Lemma 3.2].
In Section 5.2, we also show that two different initializations with the same spectrum (i.e., are equivalent to the same diagonal A ) can have very different performance.
S4D-LegS. The HiPPO-LegS matrix has DPLR representation A ( D ) − P P ⊤ , and Gupta [11] showed that simply approximating it with A ( D ) works quite well (5).
Our first result is providing a clean mathematical interpretation of this method.
Theorem 3 shows a surprising fact that does not hold in general for DPLR matrices (Appendix A.1), and arises out of the special structure of this particular matrix.
Theorem 3.
Let A =
A ( N ) − P P ⊤ and B be the HiPPO-LegS matrices, and K A , B ( t ) be its basis.
As the state size N → ∞ , the SSM basis K A ( N ) , B / 2 ( t ) limits to K A , B ( t ) (Fig. 2).
Note that A ( N ) is then unitarily equivalent to A ( D ) , which preserves the stability and timescale [10] of the system.
We define S4D-LegS to be the S4D method for this choice of diagonal A =
A ( D ) .
Theoremexplains the empirical results in [11] whereby this system performed quite close to S4, but was usually slightly worse.
This is because DSS is a variant of S4D-LegS, which by Theorem 3 is a noisy approximation to S4-LegS. Fig.
2 illustrates this result, and also shows a curious phenomenon involving different discretization rules that is open for future work.
S4D-Inv.
To further simplify S4D-LegS, we analyze the structure of A ( D ) = diag ⟨ A ⟩ in more detail.
The real part is easy to understand, which follows from the analysis in [9]: ℜ ( A )
= − 12 1 .
Let the imaginary part be sorted, i.e. ℑ ( A ) n is the n -th largest (positive) imaginary component.
We empirically deduced the following conjecture for the asymptotics of the imaginary part.
Conjecture 4.
As N → ∞ , ℑ ( A )
→ 1 π N
+ c where c ≈ 0 .
5236 is a constant.
For a fixed N , the other eigenvalues satisfy an inverse scaling in n : ℑ ( A ) n = Θ( n − 1 ) .
Fig.
3 empirically supports this conjecture.
Based on Conjecture 4, we propose the initialization
S4D-Inv to use the following inverse-law diagonal matrix which closely approximates S4D-LegS.
[Equation]
S4D-Lin.
While S4D-Inv can be seen as an approximation to the original S4-LegS, we propose an even simpler scaling law for the imaginary parts that can be seen as an approximation of S4-FouT ([10]), where the imaginary parts are simply the Fourier series frequencies (i.e. matches the diagonal part of the DPLR form of S4-FouT).
Fig.
1 ( Right ) illustrates the S4D-Lin basis e t A B , which are simply damped Fourier basis functions.
Experiments
Our experimental study shows that S4D has strong performance in a wide variety of domains and tasks, including the well-studied Long Range Arena (LRA) benchmark where the best S4D variant is competitive with S4 on all tasks and significantly outperforms all non-SSM baselines.
We begin with controlled ablations of the various representations of diagonal state space models.
Sections 5.1 and 5.2 ablate the proposed methods for parameterizing, computing, and initializing diagonal SSMs from Sections 3 and 4.
Section 5.3 show full results of larger models on standard benchmarks,
Methodology and Datasets.
In order to study the effects of different S4 and S4D variants in a controlled setting, we propose the following protocol.
We focus on three datasets covering a varied [Figure] range of data modalities (image pixels, biosignal time series, audio waveforms), sequence lengths (1K, 4K, 16K) and tasks (classification and regression with bidirectional and causal models).
• Sequential CIFAR (sCIFAR).
CIFAR-10 images are flattened into a sequence of length 1024 , and a bidirectional sequence model is used to perform 10-way classification.
• BIDMC Vital Signs.
EKG and PPG signals of length 4000 are used to predict respiratory rate (RR), heart rate (HR), and blood oxygen saturation (SpO2).
We focus on SpO2 in this study. Speech Commands (SC).
2 A 1-second raw audio waveform comprising 16000 samples is used for 35-way spoken word classification.
We use an autoregressive (AR) model to vary the setting; this causal setting more closely imitates autoregressive speech generation, where SSMs have shown recent promise [5].
We fix a simple architecture and training protocol that works generically.
The architecture haslayers and hidden dimension H = 128 , resulting in ∼ 100 K parameters.
All results are averaged over multiple seeds (full protocol and results including std. reported in Appendix B).
5.1 Parameterization, Computation, Discretization
Given the same diagonal SSM matrices A , B , there are many variants of how to parameterize the matrices and compute the SSM kernel described in Section 3. We ablate the different choices described in Table 1.
Results are in Table 2, and show that:
(i)
Computing the model with a softmax instead of Vandermonde product does not make much difference (ii) Training B is consistently slightly better (iii) Different discretizations (Section 3.1) do not make a noticeable difference (iv) Unrestricting the real part of A (Section 3.3) may be slightly better
These ablations show that for a fixed initialization ( A , B ) , different aspects of parameterizing SSMs make little difference overall.
This justifies the parameterization and algorithm S4D uses (Section 3.4), which preserves the choices of the original S4 model and is simpler than DSS. For the remaining
[Table]
of the experiments in Section 5.2 and Section 5.3, we fix the S4D parameterization and algorithm described in Section 3.
Note that this computes exactly the same kernel as the original S4 algorithm when the low-rank portion is set to 0 , allowing controlled comparisons of the critical state matrix A for the remainder of this section.
5.2 S4D Initialization Ablations
The original S4 model proposed a specific formula for the A matrix, and the first diagonal version [11] used a specific matrix based on it.
Our new proposed variants S4D-Inv and S4D-Lin also define precise formulas for the initialization of the A matrix (8).
This raises the question of whether the initialization of the A still needs to be so precise, despite the large simplifications from the original version.
We perform several natural ablations on these initializations, showing that even simple variations of the precise formula can degrade performance.
Imaginary part scaling factor.
The scaling rules for the imaginary parts of S4D-Inv and S4D-Lin are simple polynomial laws, but how is the constant factor chosen and how important is it?
These constants are based on approximations to HiPPO methods (e.g. Conjecture 4).
Note that the range of imaginary components for S4D-Inv and S4D-Lin are quite different (Fig. 3); the largest imaginary part is N 2 π for S4D-Inv and πN for S4D-Lin.
We consider scaling all imaginary parts by a constant factor of 0 .
01 or 100 .
to investigate whether the constant matters.
Note that this preserves the overall shape of the basis functions (Fig. 1, dashed lines) and simply changes the frequencies, and it is not obvious that this should degrade performance.
However, both changes substantially reduce the performance of S4D in all settings.
Randomly initialized imaginary part.
Next, we consider choosing the imaginary parts randomly.
For S4D-Inv, we keep the real parts equal to − 12 and set each imaginary component to A n = −
+ i Nπ
( N 2 u +1 − 1) for u ∼ N ·U [0 , 1] .
Note that when u is equally spaced in [0 , 1] instead of uniformly random, this exactly recovers S4D-Inv (8), so this is a sensible random approximation to it.
Similarly, we consider a variant of S4D-Lin that replaces the n in (9) with N · U [0 , 1] .
Table 3a ( Random Imag ) shows that this small change causes minor degradation in performance.
We additionally note that the randomly initialized imaginary ablation can be interpreted as follows.
Fig.
3 shows the asymptotics of the imaginary parts of SSM matrices, where the imaginary parts of the eigenvalues correspond to y-values corresponding to uniformly spaced nodes on the x-axis.
This ablation then replaces the uniform spacing on the x-axis with uniformly random x values.
Randomly initialized real part.
We considering initializing the real part of each eigenvalue as −U [0 , 1] instead of fixing them to − 12 .
Table 3a(Left, Random Real ) shows that this also causes minor but consistent degradation in performance on the ablation datasets.
Finally, we also consider randomizing both real and imaginary parts, which degrades performance even further.
Ablation: Other S4D matrices.
Other simple variants of initializations show that it is not just the range of the eigenvalues but the actual distribution that is important (Fig. 3). Both S4D-Inv2 and S4D-Quad have real part − 1 2 and imaginary part satisfying the same maximum value as Conjecture 4.
The S4D-Inv2 initialization uses the same formula as S4D-Inv, but replaces a 2 n + 1 in the denominator with n + 1 .
The S4D-Quad initialization uses a polynomial law with power 2 instead of − 1 (S4D-Inv) or 1 (S4D-Lin).
We include two additional methods here that are not based on the proposed S4D-Inv or S4D-Lin methods.
First, S4D-Rand uses a randomly initialized diagonal A , and validates that it performs [Table] [Table] poorly, in line with earlier findings [9, 11].
Second, S4D-Real uses a particular real initialization with A n = − ( n + 1) .
This is the exact same spectrum as the original S4(-LegS) method, which validates that it is not just the diagonalization that matters, highlighting the limitations of Proposition 2.
5.3 Full Comparisons of S4D and S4 Methods
Trainable A , B matrices.
Table 3b shows the performance of all S4D and S4 variants [10] on the ablations datasets.
We observe several interesting phenomena: (i)
Freezing the matrices performs comparably to training them on sCIFAR and BIDMC, but is substantially worse on SC.
We hypothesize that this results from ∆ being poorly initialized for SC, so that at initialization models do not have context over the entire sequence, and training A and B helps adjust for this.
As further evidence, the finite window methods S4-LegT and S4-FouT (defined in [10]) have the most limited context and suffer the most when A is frozen.
(ii)
The full DPLR versions are often slightly better than the diagonal version throughout the entire training curve.
We report the validation accuracy after 1 epoch of training on sCIFAR and SC to illustrate this phenomenon.
Note that this is not a consequence of having more parameters (Appendix B).
Large models on ablation datasets.
Finally, we relax the strict requirements on model size and regularization for the ablation datasets, and show the performance of S4 and S4D variants on the test sets with a larger model (architecture and training details in Appendix B) when the model size and regularization is simply increased (Table 4).
We note that results for each dataset are better than the original S4 model, which was already state-of-the-art on these datasets [8, 9].
We use the same hyperparameter setting for the state-of-the-art S4 model in [10] on the Long Range Arena benchmark for testing long dependencies in sequence models.
S4D variants are highly competitive on all datasets except Path-X, and outperform the S4 variants on several of them.
On
Path-X using this hyperparameter setting with bidirectional models, only S4D-Inv, our simpler approximation to the original S4-LegS model, achieves above random chance, and has an average of 85% on the full LRA suite, more than 30 points better than the original Transformer [24]. [Table] [Table]
Conclusion
State space models based on S4 are a promising family of models for modeling many types of sequential data, with particular strengths for continuous signals and long-range interactions.
These models are a large departure from conventional sequence models such as RNNs, CNNs
[Table]
and Transformers, with many new ideas and moving parts.
This work provides a more in-depth exposition for all aspects of working with S4-style models
[Table]
from their core structures and kernel computation algorithms
[Table]
to miscellaneous choices in their parameterizations, to new theory and methods for their initialization.
We systematically analyzed and ablated each of these components
[Table]
and provide recommendations for building a state space model that is as simple as possible
[Table]
while as theoretically principled and empirically effective as S4.
We believe that S4D can be a strong generic sequence model for a variety of domains
[Table]
that opens new directions for state space models theoretically
[Table]
and is much more practical to understand and implement for practitioners.
Acknowledgments
We gratefully acknowledge the support of DARPA under Nos.
FA86501827865 (SDH) and
FA86501827882 (ASED); NIH under No. U54EB020405 (Mobilize), NSF under Nos.
CCF (Beyond Sparsity), CCF1563078 (Volume to Velocity), and 1937301 (RTML); ONR
under No.
N000141712266 (Unifying Weak Supervision); the Moore Foundation, NXP, Xilinx, LETI-CEA,
Intel, IBM, Microsoft, NEC, Toshiba, TSMC, ARM, Hitachi, BASF, Accenture, Ericsson, Qualcomm,
Analog Devices, the Okawa Foundation, American Family Insurance, Google Cloud, Swiss Re,
Brown Institute for Media Innovation, Department of Defense (DoD) through the National Defense
Science and Engineering Graduate Fellowship (NDSEG) Program, Fannie and John Hertz Foundation, National Science Foundation Graduate Research Fellowship Program, Texas Instruments, and members of the Stanford DAWN project: Teradata, Facebook, Google, Ant Financial, NEC,
VMWare, and Infosys.
The U.S. Government is authorized to reproduce and distribute reprints for
Governmental purposes notwithstanding any copyright notation thereon.
Any opinions, findings, and conclusions or recommendations expressed in this material are those of the authors and do not necessarily reflect the views, policies, or endorsements, either expressed or implied, of DARPA, NIH,
ONR, or the U.S. Government.
A Method Details
Listing 1 Full Numpy example of the parameterization and computation of a 1D S4D-Lin model def parameters(N, dt_min=1e-3, dt_max=1e-1): # Initialization log_dt
=
np.random.rand() * (np.log(dt_max)-np.log(dt_min))
+ np.log(dt_min)
A = -0.5 + 1j * np.pi * np.arange(N//2)
# S4D-Lin initialization
B = np.ones(N//2) + 0j C = np.random.randn(N//2)
+ 1j * np.random.randn(N)
# Variance preserving initialization return log_dt, np.log(-A.real), A.imag, B, C def kernel(L, log_dt, log_A_real, A_imag, B, C):
# Discretization (e.g. bilinear transform) dt, A = np.exp(log_dt), -np.exp(log_A_real) + 1j * A_imag dA, dB =
(1+dt*A/2) / (1-dt*A/2), dt*B / (1-dt*A/2) # Computation (Vandermonde matrix multiplication - can be optimized)
# Return twice the real part - same as adding conjugate pairs return 2 * ((B*C) @
(dA[:, None] ** np.arange(L))).real def forward(u, parameters): L = u.shape[-1]
K = kernel(L, *parameters)
# Convolve y =
u
* K using FFT
K_f, u_f
= np.fft.fft(K, n=2*L), np.fft.fft(u, n=2*L) return np.fft.ifft(K_f*u_f, n=2*L)[..., :L]
A.1 Proofs
We prove Theorem 3, and then show why this it is a surprising result that is not true in general to low-rank perturbations of SSMs.
We start with the interpretation of the S4-LegS matrix shown in [10], which corresponds to Fig.(Left).
Theorem 5.
Let A , B , P be the matrices defined in equation (4) .
The SSM kernels K n ( t ) =
e ⊤ n e t A B have the closed form formula
[Equation]
where L n are the Legendre polynomials shifted and scaled to be orthonormal on the interval [0 , 1] .
Lemma A.1. The functions L n ( e − t ) are a complete orthonormal basis with respect to the measure ω ( t ) = e − t . Proof.
The polynomials are defined to be orthonormal on [0 , 1] , i.e.
[Equation]
By the change of variables t =
e − s with dtds = − e − s , [Equation] which shows the orthonormality.
Completeness follows from the fact that polynomials are complete. Proof of Theorem 3. We start with the standard interpretation of SSMs as convolutional systems. The SSM x ′ ( t ) =
A x ( t ) + B u ( t ) is equivalent to the convolution [Equation] for the SSM kernels (equation (3)). Defining u ( t ) ( s ) =
u ( t − s ) , we can write this as [Equation] where ω ( s ) = e − s and ⟨ p ( s ) , q ( s ) ⟩ ω
=
{+(cid:82) ∞ 
+}
p ( s ) q ( s ) ω ( s ) ds is the inner product in the Hilbert space of L2 functions with respect to measure ω . By Theorem 5, the K n are a complete orthonormal basis in this Hilbert space. There x n ( t ) represents a decomposition of the function u ( t ) with respect to this basis, and can be recovered as a linear combination of these projections [Equation]
Pointwise over the inner times s , [Equation] This implies that [Equation]
Intuitively, due to the function reconstruction interpretation of HiPPO [10], we can approximate u ( t ) using knowledge in the current state x ( t ) . There in the limit N → ∞ , the original SSM is equivalent to
[Equation]
General low-rank perturbations. Finally, we remark that this phenomenon where removing the low-rank correction to a DPLR matrix approximates the original dynamics, is unique to this HiPPO-LegS matrix.
We note that if instead of P P ⊤ , a random rank-1 correction is added to the
HiPPO-LegS matrix in Theorem 3, the resulting SSM kernels look completely different and in fact diverge rapidly as the magnitude of P increases (Fig. 4).
Similarly, Fig. 5a shows a new S4 variant called S4-FouT that is also DPLR [10], but removing the low-rank component dramatically changes the SSM kernels.
Ablation datasets training protocol. The architecture has 4 layers and hidden dimension H = 128 , resulting in around 100K trainable parameters.
The A and B parameters were tied across the H SSM copies; therefore the S4 models have only H × { num. layers } more parameters than S4D models, arising from the P tensor in the DPLR representation A = Λ − P P ⊤ . This choice was made
[Figure]
[Figure]
because it generally does not affect performance much, while reducing parameter count and ensuring that S4 vs. S4D models have very similar numbers of parameters.
All results are averaged over 2 or 3 seeds. All models use learning rate 0 .
01 weight decay, and no other regularization or data augmentation. For the classification tasks (sCIFAR and SC). we use a cosine scheduler with 1 epoch warmup and decaying to 0 . For the regression task (BIDMC), we use a multistep scheduler following [21, 8].
Reported results are all best validation accuracy, except for the large models in Table 4. Full results for parameterization ablations. Table 6 and Table 7 contain the raw results for Tableincluding standard deviations.
Full results for large models on ablations datasets. Tables 8 to 10 show full results comparing our proposed methods against the best models from the literature; citations indicate numbers from prior work.
[Table] No
Vandermonde 84.78 (0.16) 89.62 (0.03) 0.1355 (0.0039) Yes Softmax 85.37 (0.43) 90.06 (0.11) 0.1170 (0.0039) Yes
Vandermonde 85.37 (0.43) 90.34 (0.18) 0.1274 (0.0020)
[Table] [Table] [Table]
Note that earlier works on the Speech Commands dataset typically use pre-processing such as MFCC features [Table] or a 10-class subset of the full 35-class dataset [14, 19, 9]. As we are not aware of a collection of strong baselines for raw waveform classification using the full dataset
[Table] we trained several baselines from scratch for Table 10. The InceptionNet, ResNet-
[Table] and XResNet-50 models are 1D adaptations from Nonaka and Seita [16] of popular CNN architectures for vision.
The ConvNet architecture is a generic convolutional neural network that we tuned for strong performance, comprising: • Four stages, each composed of three identical residual blocks.
• The first stage has model dimension (i [Table]
e. channels, in CNN nomenclature) H = 64 . Each stage doubles the dimension of the previous stage (with a position-wise linear layer) and ends in an average pooling layer of width 4 . Thus, the first stage operates on inputs of length 16384 , dimension 64 (the input is zero-padded from 16000 to 16384) and the last on length 256 , dimension 512 . • Each residual block has a (pre-norm) BatchNorm layer followed by a convolution layer and GeLU activation.
• Convolution layers have a kernel size of 25 . Long Range Arena. Our Long Range Arena experiments follow the same setup as the original S paper with some differences in model architecture and hyperparameters.
The main global differences are as follows:
Bidirectional
The original S4 layer is unidirectional or causal, which is an unnecessary constraint for the classification tasks appearing in LRA.
Goel et al. [5] propose a bidirectional version of S4 that simply concatenates two S4 convolution kernels back-to-back. We use this for all tasks.
GLU feedforward S4 consists of H independent 1-dimensional SSMs, each of which are processed by an independent S4 SSM mapping ( A , B , C , D ) .
These outputs are then mixed with a
[Table]
[Table]
position-wise linear layer, i.e. W y for a learned matrix W ∈ R H × H . Instead of this linear mapping, we use a GLU activation ( W 1 y ) ◦ σ ( W 2 y ) for W 1 ,
W 2 ∈ R H
× H [3].
These have been empirically found to improve linear layers of DNNs in general [22].
Cosine scheduler Instead of the plateau scheduler used in [9], we use a cosine annealing learning rate scheduler for all tasks. Regularization Almost all tasks used no dropout and 0 . 05 weight decay.
Architecture Almost all tasks used an architecture with 6 layers, H = 256 hidden units, BatchNorm, pre-norm placement of the normalization layer. Exceptions to the above rules are described below. Full hyperparameters are in Table 11.
sCIFAR / LRA Image. This dataset is grayscale sequential CIFAR-10, and the settings for this task were taken from S4’s hyperparameters on the normal sequential CIFAR-10 task. In particular, this used LayerNorm [1] instead of BatchNorm [13], a larger number of hidden features H , post-norm instead of pre-norm, and minor dropout. We note that the choice of normalization and increased H do not make a significant difference on final performance, still attaining classification accuracy in the high 80’s.
Dropout does seem to make a difference.
BIDMC.
We used a larger state size of N = 256 , since we hypothesized that picking up higher frequency features on this dataset would help.
We also used a step scheduler that decayed the LR by 0 .
5 every 100 epochs, following prior work [21, 8].
ListOps.
We hypothesized that this task benefits from deeper models, because of the explicit hierarchical nature of the task, so the architecture used here had 8 layers and H = 128 hidden features.
However, results are very close with much smaller models.
We also found that post-norm generalized better than pre-norm, but results are again close (less than 1% difference).
PathX. As described in [10], the initialization range for PathX is decreased from (∆ min , ∆ max ) =
to (∆ min , ∆ max ) =