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: This paper presents a novel parameterization for the established state-space models (SSM), which tackles the scalability problem in linear state-space layers (LSSL) in modeling long-range dependencies for very long sequences. The proposed technique reparameterizes the structured state matrix in LSSL that allows the state to memorize the past, a key component following the continuous-time memorization theory. A complete theoretical description of this reparameterization and empirical evaluations on a diverse set of benchmarks are presented. The proposed model, S3, achieves astonishing results both in terms of performance and computational efficiency.