PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Albert Gu, Isys Johnson, Karan Goel, Khaled Saab, Tri Dao, Atri Rudra, Christopher Ré

arXiv:2110.13985v1cs.LGcs.AI

TL;DR

Long-sequence models trade off recurrence, convolution, and continuous-time modeling capabilities. The paper introduces Linear State-Space Layers, which combine these perspectives and use structured state matrices for long-range memory; stacked LSSLs achieve state-of-the-art results across image, healthcare, and speech benchmarks.

  • Problem

    Existing recurrent, convolutional, and neural differential-equation models have complementary strengths and computational tradeoffs for efficiently modeling sequences longer than a few thousand steps.

  • Method

    The paper introduces LSSLs, which simulate linear continuous-time state-space systems while supporting recurrent, convolutional, and continuous-time views, and uses structured matrices for continuous-time memorization.

  • Results

    LSSLs obtain state-of-the-art results across sequential image classification, healthcare regression, and speech tasks, including over 10% accuracy on sequential CIFAR and up to 80% RMSE reduction on length-4000 healthcare series.

  • Takeaways & Limitations

    LSSLs provide a single sequence-modeling framework that preserves the stated strengths of recurrence, convolution, and continuous-time methods while handling very long sequences effectively.

  • Takeaways & Limitations

    Theoretical fast algorithms for LSSLs are established under exact arithmetic rather than floating-point computation, leaving practical numerical stability unresolved.

Abstract

from arXiv · show

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.

1 Introduction

LSSLs are designed to combine the strengths of recurrent, convolutional, and continuous-time sequence models while addressing their tradeoffs on long sequences. The paper develops structured LSSLs for long-range memory and reports strong results across image, healthcare, and speech tasks.

  • Model motivation and design: LSSLs combine parallelizable training, stateful inference, and continuous-time capabilities in one sequence model.They simulate a linear continuous-time state-space representation and can be viewed as recurrent, convolutional, or continuous-time models.
  • Expressivity and connections: LSSLs generalize convolutional kernels and encompass popular RNN methods as special cases.Control-theoretic results connect LSSLs to arbitrary 1-D convolutional kernels, while discretization and ODE analyses relate them to RNN heuristics such as gating.
  • Long-range memory: Structured state matrices give LSSLs continuous-time memorization and make learning the state matrix and timescale theoretically computationally feasible.These structures address long-dependency memory while new algorithms provide speedups under certain computation models.
  • Empirical evaluation: Over 10% accuracy on sequential CIFAR and up to 80% RMSE reduction on length-4000 healthcare series establish strong benchmark performance.The paper reports state-of-the-art results over recent RNN-, CNN-, and NDE-based methods on sequential image classification and healthcare regression.
  • Empirical evaluation: A small LSSL comes within 2.16 accuracy points of a specialized ResNet-18 on length-38000 sequential CelebA sequences.The ResNet-18 uses 10x more parameters and is trained directly on images.
  • Empirical evaluation: On raw length-16000 speech signals, LSSLs outperform previous methods by over 20 accuracy points and beat baselines using length-160 hand-crafted features.The speech result is reported with one-fifth the training time and contrasts raw-sequence processing with standard 100x preprocessing.

2 Technical Background

The paper introduces numerical approximation and discretization tools for differential equations, emphasizing the step size Δt as a timescale and using HiPPO for continuous-time memorization.

  • Approximations of differential equations: Differential equations can be approximated by iterating an integral equation for the state trajectory.Picard iteration keeps an approximation fixed inside the differential function while producing successive approximations.
  • Discretization: Discretization approximates states at desired time points, with different right-hand-side integral approximations yielding different schemes.The generalized bilinear transform specializes this process to linear ordinary differential equations.
  • Discretization: The generalized bilinear transform updates the state using A, B, the step size Δt, and parameter α through an implicit matrix expression.This update is the paper’s main method for computing accurate discrete-time approximations of continuous-time models.
  • Discretization: The cases α = 0, α = 1, and α = 1/2 correspond respectively to Euler, backward Euler, and the bilinear method.The bilinear method is identified as stability-preserving and is used as the main discretization method later.
  • Timescales and continuous-time memory: The step size Δt acts as a timescale: shorter capturable dependencies generally correspond to larger Δt, while 1/Δt approximates dependency length.The paper relates Δt to ODE discretization, RNN gating, and convolutional kernel width.
  • Timescales and continuous-time memory: Continuous-time memory projects an input history onto basis functions, producing coefficients that optimally approximate the history under a measure ω.HiPPO supplies a mathematical framework for constructing long-dependency memory models.

3 Linear State-Space Layers (LSSL)

LSSLs are sequence-to-sequence layers obtained by discretizing linear continuous-time state-space models. They provide recurrent inference, convolutional training, continuous-time behavior, and theoretical links to RNNs and convolutional filters.

  • Definition and implementation: An LSSL maps sequences by discretizing a linear state-space representation with parameters A, B, C, D, and Δt.For hidden dimension H, the layer independently broadcasts selected parameters across feature channels.
  • Different computational views: In its recurrent view, an LSSL carries prior context in a state and uses fixed computation and storage per timestep.This provides stateful inference over potentially unbounded input sequences.
  • Different computational views: In its convolutional view, the LSSL output is a non-circular convolution that can be computed for the whole sequence using three FFTs.The convolutional representation supports parallelized training.
  • Computational bottleneck: Learning A and Δt is computationally difficult because the recurrent and convolutional views require repeated state operations and Krylov-function recomputation.With fixed parameters, A and the convolution kernel can be cached; learning them removes that efficiency.
  • Connections to convolutions: Rational convolutional filters can be represented by state-space models, so arbitrary filters can be approximated by LSSLs through rational approximation.This establishes the converse connection from convolutional filters to LSSL representations.
  • Connections to RNNs: RNN gating can be interpreted as learning an ODE discretization step size or timescale rather than merely as an optimization heuristic.A one-dimensional gated recurrence is exactly the backward-Euler discretization of a linear ODE under the stated parameterization.
  • Connections to RNNs: Deep stacked LSSLs with position-wise nonlinearities can approximate nonlinear ODEs, moving nonlinear modeling power into depth.The paper states that infinitely deep order-1 stacks can approximate any nonlinear ODE of the specified form.

4 Combining LSSLs with Continuous-time Memorization

The paper specializes LSSLs with structured state matrices that support continuous-time memorization of long dependencies and admits faster theoretical computation. These benefits address the basic layer’s memory and training limitations, but the efficiency theorem assumes exact arithmetic.

  • Limitations of basic LSSLs: Generic LSSLs can inherit recurrence and convolution limitations on long dependencies, and random state matrices are empirically ineffective.Repeated multiplication by A may suffer from vanishing gradients in the recurrent view.
  • Continuous-time memorization: HiPPO provides continuous-time memorization operators for histories under a measure ω, but previously known operators covered only special cases and were non-trainable.The operator hippo(ω) maps a function to a continuous representation of its past.
  • Continuous-time memorization: For an arbitrary measure ω, the optimal memorization operator hippo(ω) has a low recurrence-width structured state matrix A.This generalizes the known continuous-time memorization construction beyond the previously identified special cases.
  • Continuous-time memorization: For measures associated with classical orthogonal polynomials, hippo(ω) is 3-quasiseparable.Quasiseparability supplies additional algorithmic structure beyond the general low-recurrence-width result.
  • Structured LSSLs: The full LSSL uses a state matrix A from the quasiseparable family, while LSSL-f fixes A to a prescribed HiPPO matrix.The structured family is intended to retain the continuous-time memorization interpretation while permitting learning.
  • Efficient algorithms: Restricting A to quasiseparable matrices provides both long-range-memory theory and computational efficiency for recurrent and convolutional views.These matrices support linear-time matrix-vector multiplication and fast Krylov functions.
  • Efficient algorithms: For constant k, the Krylov function of a k-quasiseparable A can be computed in quasi-linear time and space Õ(N + L) with logarithmic depth.The theorem uses an exact arithmetic model and does not account for bit complexity or numerical stability.
  • Scope and caveat: The paper treats the efficiency result as a theoretical proof of concept because Theorem 2 is established over exact arithmetic rather than floating-point numbers.This is the explicit computational caveat for the structured LSSL approach.

5 Empirical Evaluation

LSSLs substantially improve results across long-sequence image, healthcare, speech, and timescale-shift benchmarks while retaining computational and modeling benefits from recurrent, convolutional, and continuous-time views.

  • Image and Time Series Benchmarks: LSSL sets SoTA on sequential CIFAR by more than 10 accuracy points, using at least 5x fewer parameters than the previous SoTA.
  • Image and Time Series Benchmarks: LSSL reduces RMSE by more than two-thirds on all length-4000 BIDMC vital-sign prediction datasets.The tasks predict respiratory rate, heart rate, and oxygen saturation.
  • Modeling and Computational Benefits: LSSLs reach target benchmark performance in a fraction of the previous model’s time and adapt to test-time timescale shifts by changing ∆t at inference.Learning A and ∆t adds less than 1% to the parameter count relative to base models.
  • Modeling and Computational Benefits: Random A matrices perform poorly, while structured and trainable A matrices consistently improve performance despite negligible parameter differences.The ablations support continuous-time memorization theory and learning the memorization measure.

6 Discussion

The discussion presents LSSL as a principled unification of recurrent, convolutional, and continuous-time models, while identifying practical limitations in inference evaluation, numerical algorithms, and memory use.

  • Discussion: LSSL theoretically and empirically generalizes the main modern time-series model families and addresses long-term memory with continuous-time memorization theory.
  • Limitations: The experiments did not evaluate applications leveraging LSSL’s theoretically possible efficient recurrent inference.Follow-up work later showed that speeding up some inference applications is possible in practice.
  • Limitations: Theorem 2’s fast algorithm was not implemented and was later found numerically unstable, leaving practical fast algorithms as an open question.Its contribution is therefore characterized as a proof of concept under an arithmetic-complexity model.
  • Limitations: LSSL and LSSL-fixed use O(NL) rather than O(L) space for length-L sequences, requiring multi-GPU training for the largest speech and image experiments.
  • Limitations: Follow-up work revisited and resolved the computation and space-complexity issues with a new Structured State Space parameterization and algorithms.

A Related Work

Related work places LSSL among continuous-time RNNs, HiPPO and LMU memory models, continuous-time CNNs, and neural differential equations, while distinguishing its trainable structured state spaces.

  • HiPPO: HiPPO and LMU provide continuous-time memory dynamics, but prior methods could not train effective A matrices or ∆t and imposed a 1-D state bottleneck.
  • HiPPO: LSSL removes the bottleneck by using H independent 1-D-to-1-D copies, while allowing multiple timescales and increasing computation by a factor of H.
  • HiPPO: LSSL generalizes HiPPO operators to a parameterized structured class, enabling A-matrix training and continuous-time memorization interpretations.
  • HiPPO: Theorem 2 introduced algorithms for structured matrices, but the presented algorithm was later found impractical and replaced by improved subsequent work.
  • Continuous-time CNNs: Unlike finite-kernel CNNs, CKConv parameterizes a continuous convolution kernel that can be sampled at any resolution.
  • Continuous-time RNNs: Continuous-time RNN research spans dynamical-systems analyses, gradient-stability architectures, explicit ODE models, and ODE-RNN variants.
  • Gating mechanisms: LSSL’s sigmoid-gate result gives an exact discretized damped-ODE interpretation rather than merely motivating gates as arbitrary monotonic functions.
  • Computation: When A and ∆t are fixed, cached transition and Krylov computations enable faster processing, but materializing the Krylov matrix requires O(NL) space.

B.4 Deep Neural Network Architecture

Deep LSSL models stack state-space layers with feedforward mixing, residual connections, and normalization, while discretizing linear continuous-time dynamics into recurrent updates.

  • Architecture: Deep LSSL models stack identical LSSL layers in a simple deep neural-network architecture.The overall parameter count is M · H · (H + N), with small and large configurations of approximately 200K and 2M parameters.
  • Channels: The model uses M output channels, where C ∈ R^M×N and D ∈ R^M×1 map one-dimensional inputs to M-dimensional outputs.
  • Feedforward: A position-wise H·M → H feedforward layer mixes hidden features and projects outputs back to dimension 1 when needed.An optional GeLU nonlinearity separates the LSSL from this projection.
  • Feedforward: The feedforward factorization resembles depth-wise separable convolutions, while residual connections and layer normalization support stacking.
  • Discretization: Linear state-space dynamics discretize to updates involving A, B, ∆t, and a convex-combination parameter α.α = 0 and α = 1 give forward and backward Euler methods, while α = 1/2 gives the trapezoid-rule bilinear method.
  • Discretization: The bilinear method with α = 1/2 was numerically better than forward and backward Euler in prior synthetic and end-to-end experiments.

C.2 RNNs are LSSLs: Proof of Results in Section 3.2

This section shows that common RNN architectural mechanisms can be understood as approximations of shared continuous-time dynamics, with LSSLs providing a unifying formulation.

  • Popular RNN families approximate the same continuous-time dynamics through discretization and Picard iteration.
  • Nonlinear RNNs discretize the linear component using backward Euler, producing the gating mechanism.
  • A special LSSL case approximates the nonlinear dynamics by applying Picard iteration to the nonlinear term.
  • Deep linear RNNs combine depth-wise Picard iteration with time-wise discretization through gates.
  • The section assumes continuous dynamics with a function that is continuous in time and Lipschitz in its state argument.

C.2.1 Intuition / Proof Sketches

The proof sketches connect RNN updates and LSSL layers to continuous-time dynamics by combining discretization with Picard iteration, then establish the computational structure used by LSSLs.

  • Intuition / Proof Sketches: Backward Euler discretization turns continuous dynamics into gated recurrences, while Picard iteration supplies the corresponding nonlinear approximation.
  • Intuition / Proof Sketches: Alternating this LSSL with point-wise nonlinear functions yields a deep LSSL whose outputs converge to the target dynamics.
  • Intuition / Proof Sketches: A special continuous-time LSSL with N = 1, A = −1, B = 1, C = 1, and D = 0 approximates the nonlinear dynamics.
  • Intuition / Proof Sketches: Combining discretization and Picard iteration shows that deep discrete-time linear RNNs approximate the same continuous-time dynamics.
  • Intuition / Proof Sketches: HiPPO operators for classical orthogonal-polynomial measures lead to linear ODEs and structured matrices enabling efficient matrix-vector multiplication.

D.3 Proof of Corollary 4.1: HiPPO for Classical Orthogonal Polynomials

This section proves that HiPPO matrices associated with classical orthogonal polynomials have quasiseparable structure, supporting efficient computation for LSSL state matrices.

  • HiPPO matrices for Jacobi and Laguerre polynomial families are shown to be quasiseparable.
  • Quasiseparable structure bounds the ranks of strictly upper- and lower-diagonal submatrices.
  • The Jacobi construction combines low-rank and diagonal factors to establish a 3-quasiseparable matrix.
  • The Laguerre construction establishes 1-quasiseparability from rank-one structure above and below the diagonal.

E.1 Proof of Theorem 2

This section develops a divide-and-conquer algorithm for computing LSSL Krylov functions with structured state matrices, yielding near-linear complexity in state and sequence lengths.

  • The Krylov function is the coefficient vector of C(I − Ax)^−1B modulo x^L.
  • Quasiseparable matrices enable recursive computation by treating I − Ax as a low-rank perturbation of smaller structured matrices.
  • The recursive reduction computes the required rational-function blocks in O((q^3 + kq^2 + k^2q)N log(N)) operations.
  • For semiseparable matrices with constant q, the computation costs O(N log^2(N)) operations.
  • Theorem 2 gives total complexity N log^2(N) + L log(L) operations, with recursion depth log^2(N) + log(L).
  • The general quasiseparable algorithm remains difficult to implement, motivating a simpler structured subclass based on diagonal and tridiagonal matrices.

F.1 Additional Experiments

Additional experiments show that LSSL remains competitive on irregularly sampled data, benefits from learning both A and ∆t, and adapts its timescales during training.

  • LSSL is competitive with the best prior methods on irregularly sampled CharacterTrajectories data.Some competing methods were specialized for this setting.
  • Learning either A or ∆t provides a consistent performance increase on SC-Raw and smaller sCIFAR models.
  • ∆t values spread across a larger range of timescales during SpeechCommands-Raw training.Figure 2 visualizes the 32 largest and smallest ∆t values at training start and end.

F.2 Methodology

The methodology specifies common optimization settings, dataset-specific training and model choices, and baseline configurations for evaluating LSSL across sequence tasks.

  • All models use Adam with learning-rate decay after validation plateaus, with batch sizes of 50 for MNIST/CIFAR and 32 for other datasets.Batch sizes were reduced when necessary to fit memory.
  • Hyperparameters are selected using validation accuracy or RMSE, with dataset-specific searches over learning rates, dropout, model sizes, and schedulers.The best hyperparameters are reported in Table 11.
  • Section 5 results omit standard deviations because most baseline results came from prior papers without error bars.The authors state that LSSL performance was stable and at least one standard deviation above baseline results for every reported result.
  • Speech Commands experiments compare raw 16000-timestep signals with 161-timestep MFCC features and evaluate LSSL against UniCORNN and NRDE baselines.Baseline runs used reported or separately tuned hyperparameters depending on the representation and sequence length.
  • Convergence comparisons use logs from corresponding baseline state-of-the-art models and wall-clock estimates measured on the same hardware.
Loading 2110.13985v1…