[Paper Review] Combining Recurrent, Convolutional, and Continuous-time Models with Linear State-Space Layers

Albert Gu, Isys Johnson+5 more|arXiv (Cornell University)|Oct 26, 2021

Neural Networks and Applications143 citations

TL;DR

The paper introduces Linear State-Space Layers (LSSLs) that unify recurrent, convolutional, and continuous-time models, providing long-range memory and stateful, parallelizable computation, achieving state-of-the-art results on very long time-series tasks.

ABSTRACT

Recurrent neural networks (RNNs), temporal convolutions, and neural differential equations (NDEs) are popular families of deep learning models for time-series data, each with unique strengths and tradeoffs in modeling power and computational efficiency. We introduce a simple sequence model inspired by control systems that generalizes these approaches while addressing their shortcomings. The Linear State-Space Layer (LSSL) maps a sequence $u \\mapsto y$ by simply simulating a linear continuous-time state-space representation $\\dot{x} = Ax + Bu, y = Cx + Du$. Theoretically, we show that LSSL models are closely related to the three aforementioned families of models and inherit their strengths. For example, they generalize convolutions to continuous-time, explain common RNN heuristics, and share features of NDEs such as time-scale adaptation. We then incorporate and generalize recent theory on continuous-time memorization to introduce a trainable subset of structured matrices $A$ that endow LSSLs with long-range memory. Empirically, stacking LSSL layers into a simple deep neural network obtains state-of-the-art results across time series benchmarks for long dependencies in sequential image classification, real-world healthcare regression tasks, and speech. On a difficult speech classification task with length-16000 sequences, LSSL outperforms prior approaches by 24 accuracy points, and even outperforms baselines that use hand-crafted features on 100x shorter sequences.

Motivation & Objective

• Motivate a unified sequence model that combines strengths of RNNs, CNNs, and neural differential equations for long sequences.
• Define and analyze Linear State-Space Layers (LSSLs) as discretized linear state-space systems that can behave like recurrent and convolutional layers.
• Incorporate continuous-time memorization via HiPPO-based structured A matrices to enable long-range memory.
• Develop efficient training and inference approaches for LSSLs, including structured matrices for tractable Krylov/FFT computations.
• Demonstrate that LSSLs achieve state-of-the-art performance on benchmarks with very long sequences across vision, healthcare, and speech tasks.

Proposed method

• Define LSSL as a sequence map u(t) -> y(t) via a discretized linear state-space model with dot{x}=Ax+Bu and y=Cx+Du.
• Show three perspectives of LSSL: recurrent (timewise state update), convolutional (Krylov-like kernel representation), and continuous-time (differential equation) forms.
• Use bilinear (alpha=1/2) discretization to obtain stable discrete-time updates x_t = Ā x_{t-1} + B̄ u_t, y_t = C x_t + D u_t.
• Leverage HiPPO-based memory operators with structured A (LRW/quasi-separable) to enable principled continuous-time memorization and long-range dependencies.
• Propose computationally efficient algorithms for Krylov function computations with structured A matrices, enabling near-linear time and parallelizable training.
• Stack LSSL layers with normalization and residual connections to form deep architectures.
• Provide ablations showing the importance of learning timescale Δt and structured A, and demonstrate robustness to test-time timescale shifts.

Experimental results

Research questions

• RQ1 Can a single linear state-space formulation emulate and unify the capabilities of RNNs, CNNs, and continuous-time models for time-series data?
• RQ2 How can we design A and Δt to enable long-range memory while keeping training and inference efficient?
• RQ3 Do HiPPO-inspired, structured A matrices enable continuous-time memorization within LSSLs and improve long-dependency modeling?
• RQ4 Do stacked LSSLs achieve state-of-the-art performance on very long sequences across diverse domains (vision, healthcare, speech)?

Key findings

• LSSLs generalize convolutions to the continuous-time domain and provide recurrent-like stateful inference.
• With HiPPO-structured A, LSSLs can memorize long-range dependencies and allow learning the memorization measure and timescale.
• LSSLs achieve state-of-the-art results on sequential image classification benchmarks (e.g., sequential CIFAR) and healthcare regression tasks with long sequences, with substantial RMSE reductions.
• On very long speech sequences (length 16000), LSSLs outperform prior methods by over 20 accuracy points and even beat baselines using hand-crafted features on much shorter inputs.
• A fixed or learned Δt enables timescale adaptation, including robustness to test-time shifts, with learning Δt providing complementary gains.
• LSSLs trained and stacked as deep networks achieve strong performance while using fewer parameters than previous SOTA models.