Combining Recurrent, Convolutional, and Continuous-time Models with Linear State Space Layers

1 Introduction . A longstanding challenge in machine learning is efficiently modeling sequential data longer than a few thousand time steps . The usual paradigms for designing sequence models involve recurrence ( e.g . RNNs ) , convolutions ( e.g . CNNs ) , or differential equations ( e.g . NDEs ) , which each come with tradeoffs . For example , RNNs are a natural stateful model for sequential data that require only constant computation/storage per time step , but are slow to train and suffer from optimization difficulties ( e.g. , the `` vanishing gradient problem '' [ 39 ] ) , which empirically limits their ability to handle long sequences . CNNs encode local context and enjoy fast , parallelizable training , but are not sequential , resulting in more expensive inference and an inherent limitation on the context length . NDEs are a principled mathematical model that can theoretically address continuous-time problems and long-term dependencies [ 37 ] , but are very inefficient . Ideally , a model family would combine the strengths of these paradigms , providing properties like parallelizable training ( convolutional ) , stateful inference ( recurrence ) and time-scale adaptation 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . ( differential equations ) , while handling very long sequences in a computationally efficient way . Several recent works have turned to this question . These include the CKConv , which models a continuous convolution kernel [ 44 ] ; several ODE-inspired RNNs , such as the UnICORNN [ 47 ] ; the LMU , which speeds up a specific linear recurrence using convolutions [ 12 , 58 ] ; and HiPPO [ 24 ] , a generalization of the LMU that introduces a theoretical framework for continuous-time memorization . However , these model families come at the price of reduced expressivity : intuitively , a family that is both convolutional and recurrent should be more restrictive than either . Our first goal is to construct an expressive model family that combines all 3 paradigms while preserving their strengths . The Linear State-Space Layer ( LSSL ) is a simple sequence model that maps a 1-dimensional function or sequence u ( t ) 7→ y ( t ) through an implicit state x ( t ) by simulating a linear continuous-time state-space representation in discrete-time ẋ ( t ) = Ax ( t ) +Bu ( t ) ( 1 ) y ( t ) = Cx ( t ) +Du ( t ) , ( 2 ) where A controls the evolution of the system and B , C , D are projection parameters . The LSSL can be viewed as an instantiation of each family , inheriting their strengths ( Fig . 1 ) : • LSSLs are recurrent . If a discrete step-size ∆t is specified , the LSSL can be discretized into a linear recurrence using standard techniques , and simulated during inference as a stateful recurrent model with constant memory and computation per time step . • LSSLs are convolutional . The linear time-invariant systems defined by ( 1 ) + ( 2 ) are known to be explicitly representable as a continuous convolution . Moreover , the discrete-time version can be parallelized during training using convolutions [ 12 , 44 ] . • LSSLs are continuous-time . The LSSL itself is a differential equation . As such , it can perform unique applications of continuous-time models , such as simulating continuous processes , handling missing data [ 45 ] , and adapting to different timescales . Surprisingly , we show that LSSLs do not sacrifice expressivity , and in fact generalize convolutions and RNNs . First , classical results from control theory imply that all 1-D convolutional kernels can be approximated by an LSSL [ 59 ] . Additionally , we provide two results relating RNNs and ODEs that may be of broader interest , e.g . showing that some RNN architectural heuristics ( such as gating mechanisms ) are related to the step-size ∆t and can actually be derived from ODE approximations . As corollaries of these results , we show that popular RNN methods are special cases of LSSLs . The generality of LSSLs does come with tradeoffs . In particular , we describe and address two challenges that naive LSSL instantiations face when handling long sequences : ( i ) they inherit the limitations of both RNNs and CNNs at remembering long dependencies , and ( ii ) choosing the state matrix A and timescale ∆t appropriately are critical to their performance , yet learning them is computationally infeasible . We simultaneously address these challenges by specializing LSSLs using a carefully chosen class of structured matrices A , such that ( i ) these matrices generalize prior work on continuous-time memory [ 24 ] and mathematically capture long dependencies with respect to a learnable family of measures , and ( ii ) with new algorithms , LSSLs with these matrices A can be theoretically sped up under certain computation models , even while learning the measure A and timescale ∆t . We empirically validate that LSSLs are widely effective on benchmark datasets and very long time series from healthcare sensor data , images , and speech . • On benchmark datasets , LSSLs obtain SoTA over recent RNN , CNN , and NDE-based methods across sequential image classification tasks ( e.g. , by over 10 % accuracy on sequential CIFAR ) and healthcare regression tasks with length-4000 time series ( by up to 80 % reduction in RMSE ) . • To showcase the potential of LSSLs to unlock applications with extremely long sequences , we introduce a new sequential CelebA classification task with length-38000 sequences . A small LSSL comes within 2.16 accuracy points of a specialized ResNet-18 vision architecture that has 10x more parameters and is trained directly on images . • Finally , we test LSSLs on a difficult dataset of high-resolution speech clips , where usual speech pipelines pre-process the signals to reduce the length by 100x . When training on the raw length16000 signals , the LSSL not only ( i ) outperforms previous methods by over 20 accuracy points in 1/5 the training time , but ( ii ) outperforms all baselines that use the pre-processed length-160 sequences , overcoming the limitations of hand-crafted feature engineering . Summary of Contributions • We introduce Linear State-Space Layers ( LSSLs ) , a simple sequence-to-sequence transformation that shares the modeling advantages of recurrent , convolutional , and continuous-time methods . Conversely , we show that RNNs and CNNs can be seen as special cases of LSSLs ( Section 3 ) . • We prove that a structured subclass of LSSLs can learn representations that solve continuous-time memorization , allowing it to adapt its measure and timescale ( Section 4.1 ) . We also provide new algorithms for these LSSLs , showing that they can be sped up computationally under an arithmetic complexity model Section 4.2 . • Empirically , we show that LSSLs stacked into a deep neural network are widely effective on time series data , even ( or especially ) on extremely long sequences ( Section 5 ) . 2 Technical Background . We summarize the preliminaries on differential equations that are necessary for this work . We first introduce two standard approximation schemes for differential equations that we will use to convert continuous-time models to discrete-time , and will be used in our results on understanding RNNs . We give further context on the step size or timescale ∆t , which is a particularly important parameter involved in this approximation process . Finally , we provide a summary of the HiPPO framework for continuous-time memorization [ 24 ] , which will give us a mathematical tool for constructing LSSLs that can address long-term dependencies . Approximations of differential equations . Any differential equation ẋ ( t ) = f ( t , x ( t ) ) has an equivalent integral equation x ( t ) = x ( t0 ) + ∫ t t0 f ( s , x ( s ) ) ds . This can be numerically solved by storing some approximation for x , and keeping it fixed inside f ( t , x ) while iterating the equation . For example , Picard iteration is often used to prove the existence of solutions to ODEs by iterating the equation xi+1 ( t ) : = xi ( t0 ) + ∫ t t0 f ( s , xi ( s ) ) ds . In other words , it finds a sequence of functions x0 ( t ) , x1 ( t ) , . . . that approximate the solution x ( t ) of the integral equation . Discretization . On the other hand , for a desired sequence of discrete times ti , approximations to x ( t0 ) , x ( t1 ) , . . . can be found by iterating the equation x ( ti+1 ) = x ( ti ) + ∫ ti+1 ti f ( s , x ( s ) ) ds . Different ways of approximating the RHS integral lead to different discretization schemes . We single out a discretization method called the generalized bilinear transform ( GBT ) which is specialized to linear ODEs of the form ( 1 ) . Given a step size ∆t , the GBT update is x ( t+ ∆t ) = ( I − α∆t ·A ) −1 ( I + ( 1− α ) ∆t ·A ) x ( t ) + ∆t ( I − α∆t ·A ) −1B · u ( t ) . ( 3 ) Three important cases are : α = 0 becomes the classic Euler method which is simply the first-order approximation x ( t + ∆t ) = x ( t ) + ∆t · x′ ( t ) ; α = 1 is called the backward Euler method ; and α = 12 is called the bilinear method , which preserves the stability of the system [ 61 ] . In Section 3.2 we will show that the backward Euler method and Picard iteration are actually related to RNNs . On the other hand , the bilinear discretization will be our main method for computing accurate discrete-time approximations of our continuous-time models . In particular , define A and B to be the matrices appearing in ( 3 ) for α = 12 . Then the discrete-time state-space model is xt = Axt−1 +But ( 4 ) yt = Cxt +Dut . ( 5 ) ∆t as a timescale . In most models , the length of dependencies they can capture is roughly proportional to 1∆t . Thus we also refer to the step size ∆t as a timescale . This is an intrinsic part of converting a continuous-time ODE into a discrete-time recurrence , and most ODE-based RNN models have it as an important and non-trainable hyperparameter [ 24 , 47 , 58 ] . On the other hand , in Section 3.2 we show that the gating mechanism of classical RNNs is a version of learning ∆t . Moreover when viewed as a CNN , the timescale ∆t can be viewed as controlling the width of the convolution kernel ( Section 3.2 ) . Ideally , all ODE-based sequence models would be able to automatically learn the proper timescales . Continuous-time memory . Consider an input function u ( t ) , a fixed probability measure ω ( t ) , and a sequence of N basis functions such as polynomials . At every time t , the history of u before time t can be projected onto this basis , which yields a vector of coefficients x ( t ) ∈ RN that represents an optimal approximation of the history of u with respect to the provided measure ω . The map taking the function u ( t ) ∈ R to coefficients x ( t ) ∈ RN is called the High-Order Polynomial Projection Operator ( HiPPO ) with respect to the measure ω . In special cases such as the uniform measure ω = I { [ 0 , 1 ] } and the exponentially-decaying measure ω ( t ) = exp ( −t ) , Gu et al . [ 24 ] showed that x ( t ) satisfies a differential equation ẋ ( t ) = A ( t ) x ( t ) +B ( t ) u ( t ) ( i.e. , ( 1 ) ) and derived closed forms for the matrix A . Their framework provides a principled way to design memory models handling long dependencies ; however , they prove only these few special cases .

Summary: This paper proposes a sequence-to-sequence model whose main building block is a linear layer called LSSL based on a standard dynamical system in control theory, including an internal ODE. LSSL is shown to implement both a convolutional and recurrent network, and, conversely, to be able, when stacked, to approximate any convolution and recurrent network with gating mechanisms, making the proposed model expressive. By endowing LSSL with the preexisting memory mechanism HiPPO, the resulting model MLSSL, and its counterpart SLLSSL that learns this operator, become computationally tractable and avoid gradient vanishing issues. The models are empirically validated on a number of regression and classification tasks showing their computational and performance advantages.