Efficiently Modeling Long Sequences with Structured State Spaces

1 INTRODUCTION . A central problem in sequence modeling is efficiently handling data that contains long-range dependencies ( LRDs ) . Real-world time-series data often requires reasoning over tens of thousands of time steps , while few sequence models address even thousands of time steps . For instance , results from the long-range arena ( LRA ) benchmark ( Tay et al. , 2021 ) highlight that sequence models today perform poorly on LRD tasks , including one ( Path-X ) where no model performs better than random guessing . Since LRDs are perhaps the foremost challenge for sequence models , all standard model families such as continuous-time models ( CTMs ) , RNNs , CNNs , and Transformers include many specialized variants designed to address them . Modern examples include orthogonal and Lipschitz RNNs ( Arjovsky et al. , 2016 ; Erichson et al. , 2021 ) to combat vanishing gradients , dilated convolutions to increase context size ( Bai et al. , 2018 ; Oord et al. , 2016 ) , and an increasingly vast family of efficient Transformers that reduce the quadratic dependence on sequence length ( Katharopoulos et al. , 2020 ; Choromanski et al. , 2020 ) . Despite being designed for LRDs , these solutions still perform poorly on challenging benchmarks such as LRA ( Tay et al. , 2021 ) or raw audio classification ( Gu et al. , 2021 ) . An alternative approach to LRDs was recently introduced based on the state space model ( SSM ) ( Fig . 1 ) . SSMs are a foundational scientific model used in fields such as control theory , computational neuroscience , and many more , but have not been applicable to deep learning for concrete theoretical reasons . In particular , Gu et al . ( 2021 ) showed that deep SSMs actually struggle even on simple tasks , but can perform exceptionally well when equipped with special state matrices A recently derived to solve a problem of continuous-time memorization ( Voelker et al. , 2019 ; Gu et al. , 2020a ) . Their Linear State Space Layer ( LSSL ) conceptually unifies the strengths of CTM , RNN and CNN models , and provides a proof of concept that deep SSMs can address LRDs in principle . Unfortunately , the LSSL is infeasible to use in practice because of prohibitive computation and memory requirements induced by the state representation . For state dimension N and sequence length L , computing the latent state requires O ( N2L ) operations and O ( NL ) space – compared to a Ω ( L + N ) lower bound for both . Thus for reasonably sized models ( e.g . N = 256 in Gu et al . ( 2021 ) ) , the LSSL uses orders of magnitude more memory than comparably-sized RNNs or CNNs . Although theoretically efficient algorithms for the LSSL were proposed , we show that these are numerically unstable . In particular , the special A matrix is highly non-normal in the linear algebraic sense , which prevents the application of conventional algorithmic techniques . Consequently , although the LSSL showed that SSMs have strong performance , they are currently computationally impractical as a general sequence modeling solution . In this work , we introduce the Structured State Space ( S4 ) sequence model based on the SSM that solves the critical computational bottleneck in previous work . Technically , S4 reparameterizes the structured state matrices A appearing in Voelker et al . ( 2019 ) ; Gu et al . ( 2020a ) by decomposing them as the sum of a low-rank and skew-symmetric term . Additionally , instead of expanding the standard SSM in coefficient space , we compute its truncated generating function in frequency space , which can be simplified into a multipole-like evaluation . Combining these two ideas , we show that the low-rank term can be corrected by the Woodbury identity while the skew-symmetric term can be diagonalized stably , ultimately reducing to a well-studied and theoretically stable Cauchy kernel ( Pan , 2001 ; 2017 ) . This results in Õ ( N + L ) computation and O ( N + L ) memory usage , which is essentially tight for sequence models . Compared to the LSSL , S4 is up to 30× faster with 400× less memory usage , while exceeding the LSSL ’ s performance empirically . Empirically , S4 significantly advances the state-of-the-art for LRD . On the LRA benchmark for efficient sequence models , S4 is as fast as all baselines while outperforming them by 20+ points on average . S4 is the first model to solve the difficult LRA Path-X task ( length-16384 ) , achieving 88 % accuracy compared to 50 % random guessing for all prior work . On speech classification with length-16000 sequences , S4 halves the test error ( 1.7 % ) of specialized Speech CNNs – by contrast , all RNN and Transformer baselines fail to learn ( ≥ 70 % error ) . Towards a general-purpose sequence model . Beyond LRD , a broad goal of machine learning is to develop a single model that can be used across a wide range of problems . Models today are typically specialized to solve problems from a particular domain ( e.g . images , audio , text , time-series ) , and enable a narrow range of capabilities ( e.g . efficient training , fast generation , handling irregularly sampled data ) . This specialization is typically expressed via domain-specific preprocessing , inductive biases , and architectures . Sequence models provide a general framework for solving many of these problems with reduced specialization – e.g . Vision Transformers for image classification with less 2D information ( Dosovitskiy et al. , 2020 ) . However , most models such as Transformers generally still require substantial specialization per task to achieve high performance . Deep SSMs in particular have conceptual strengths that suggest they may be promising as a general sequence modeling solution . These strengths include a principled approach to handling LRDs , as well as the ability to move between continuous-time , convolutional , and recurrent model representations , each with distinct capabilities ( Fig . 1 ) . Our technical contributions enable SSMs to be applied successfully to a varied set of benchmarks with minimal modification : • Large-scale generative modeling . On CIFAR-10 density estimation , S4 is competitive with the best autoregressive models ( 2.85 bits per dim ) . On WikiText-103 language modeling , S4 substantially closes the gap to Transformers ( within 0.8 perplexity ) , setting SoTA for attention-free models . • Fast autoregressive generation . Like RNNs , S4 can use its latent state to perform 60× faster pixel/token generation than standard autoregressive models on CIFAR-10 and WikiText-103 . • Irregularly sampled time-series . Like specialized CTMs , S4 can adapt to changes in time series sampling frequency without retraining , e.g . at 0.5× frequency on speech classification . • Learning with weaker inductive biases . With no architectural changes , S4 surpasses Speech CNNs on speech classification , outperforms the specialized Informer model on 40/50 time-series forecasting problems , and matches a 2-D ResNet on sequential CIFAR with over 90 % accuracy . 2 BACKGROUND : STATE SPACES . Sections 2.1 to 2.4 describe the four properties of SSMs in Fig . 1 : the classic continuous-time representation , addressing LRDs with the HiPPO framework , the discrete-time recurrent representation , and the parallelizable convolution representation . In particular , Section 2.4 introduces the SSM convolution kernel K , which is the focus of our theoretical contributions in Section 3 . 2.1 STATE SPACE MODELS : A CONTINUOUS-TIME LATENT STATE MODEL . The state space model is defined by the simple equation ( 1 ) . It maps a 1-D input signal u ( t ) to an N -D latent state x ( t ) before projecting to a 1-D output signal y ( t ) . x′ ( t ) = Ax ( t ) + Bu ( t ) y ( t ) = Cx ( t ) + Du ( t ) ( 1 ) SSMs are broadly used in many scientific disciplines and related to latent state models such as Hidden Markov Models ( HMM ) . Our goal is to simply use the SSM as a black-box representation in a deep sequence model , where A , B , C , D are parameters learned by gradient descent . For the remainder of this paper , we will omit the parameter D for exposition ( or equivalently , assume D = 0 ) because the term Du can be viewed as a skip connection and is easy to compute . 2.2 ADDRESSING LONG-RANGE DEPENDENCIES WITH HIPPO . Prior work found that the basic SSM ( 1 ) actually performs very poorly in practice . Intuitively , one explanation is that linear first-order ODEs solve to an exponential function , and thus may suffer from gradients scaling exponentially in the sequence length ( i.e. , the vanishing/exploding gradients problem ( Pascanu et al. , 2013 ) ) . To address this problem , the LSSL leveraged the HiPPO theory of continuous-time memorization ( Gu et al. , 2020a ) . HiPPO specifies a class of certain matrices A ∈ RN×N that when incorporated into ( 1 ) , allows the state x ( t ) to memorize the history of the input u ( t ) . The most important matrix in this class is defined by equation ( 2 ) , which we will call the HiPPO matrix . For example , the LSSL found that simply modifying an SSM from a random matrix A to equation ( 2 ) improved its performance on the sequential MNIST benchmark from 50 % to 98 % . ( HiPPO Matrix ) Ank = ( 2n+ 1 ) 1/2 ( 2k + 1 ) 1/2 if n > k n+ 1 if n = k 0 if n < k . ( 2 ) 2.3 DISCRETE-TIME SSM : THE RECURRENT REPRESENTATION . To be applied on a discrete input sequence ( u0 , u1 , . . . ) instead of continuous function u ( t ) , ( 1 ) must be discretized by a step size ∆ that represents the resolution of the input . Conceptually , the inputs uk can be viewed as sampling an implicit underlying continuous signal u ( t ) , where uk = u ( k∆ ) . To discretize the continuous-time SSM , we follow prior work in using the bilinear method ( Tustin , 1947 ) , which converts the state matrix A into an approximation A . The discrete SSM is xk = Axk−1 + Buk A = ( I −∆/2 ·A ) −1 ( I + ∆/2 ·A ) yk = Cxk B = ( I −∆/2 ·A ) −1∆B C = C. ( 3 ) Equation ( 3 ) is now a sequence-to-sequence map uk 7→ yk instead of function-to-function . Moreover the state equation is now a recurrence in xk , allowing the discrete SSM to be computed like an RNN . Concretely , xk ∈ RN can be viewed as a hidden state with transition matrix A. Notationally , throughout this paper we use A , B , . . . to denote discretized SSM matrices defined by ( 3 ) . Note that these matrices are a function of both A as well as a step size ∆ ; we suppress this dependence for notational convenience when it is clear . 2.4 TRAINING SSMS : THE CONVOLUTIONAL REPRESENTATION . The recurrent SSM ( 3 ) is not practical for training on modern hardware due to its sequentiality . Instead , there is a well-known connection between linear time-invariant ( LTI ) SSMs such as ( 1 ) and continuous convolutions . Correspondingly , ( 3 ) can actually be written as a discrete convolution . For simplicity let the initial state be x−1 = 0 . Then unrolling ( 3 ) explicitly yields x0 = Bu0 x1 = ABu0 + Bu1 x2 = A 2 Bu0 + ABu1 + Bu2 . . . y0 = CBu0 y1 = CABu0 + CBu1 y2 = CA 2 Bu0 + CABu1 + CBu2 . . . This can be vectorized into a convolution ( 4 ) with an explicit formula for the convolution kernel ( 5 ) . yk = CA k Bu0 + CA k−1 Bu1 + · · ·+ CABuk−1 + CBuk y = K ∗ u . ( 4 ) K ∈ RL : = KL ( A , B , C ) : = ( CA i B ) i∈ [ L ] = ( CB , CAB , . . . , CA L−1 B ) . ( 5 ) In other words , equation ( 4 ) is a single ( non-circular ) convolution and can be computed very efficiently with FFTs , provided that K is known . However , computing K in ( 5 ) is non-trivial and is the focus of our technical contributions in Section 3 . We call K the SSM convolution kernel or filter .

Summary: The paper extends previous work on linear state space models (SSM), where the state transition matrix is fixed to be a highly structured _HiPPO_ matrix, which has provably beneficial properties for memorizing long-term information from continuous-time signals. The main contribution of the paper is regarding the computational aspect of the model. Namely, a novel approach is proposed for computing the convolutional kernel associated to the discretized SSM unrolled over time. This is done by first showing that the state transition matrix can be decomposed as the sum of a normal (i.e. diagonalizable with orthonormal eigenvectors) and a low-rank (rank 1 or 2) matrix, and this representation is then combined with techniques from numerical linear algebra to reduce the problem to computing the Cauchy kernel (i.e. a well-studied problem). The resulting structured state space model (S3) is then placed into a deep neural network setting, and extensive experiments are carried out on various tasks, such as: 1) Long Range Arena, a benchmark collection for scalable transformers, 2) raw speech classification (length-16k audio signals), 3) generative modelling on CIFAR-10 and WikiText-103, 4) sequential image classification on sMNIST, pMNIST and sCIFAR. Overall, the model seems to perform very well on each task either performing close to SoTA or setting a new high score.