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

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" ), 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 , 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 (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 ; several ODE-inspired RNNs, such as the UnICORNN ; the LMU, which speeds up a specific linear recurrence using convolutions ; and HiPPO , 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)\mapsto y(t)$ through an implicit state $x(t)$ by simulating a linear continuous-time state-space representation in discrete-time

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 $\Delta 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 .

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 , 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 . 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 $\Delta 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 $\Delta 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 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 $\Delta 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 length-16000 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.

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).

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 $\Delta 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 , which will give us a mathematical tool for constructing LSSLs that can address long-term dependencies.

Any differential equation $\dot{x}(t)=f(t,x(t))$ has an equivalent integral equation $x(t)=x(t_{0})+\int_{t_{0}}^{t}f(s,x(s))\mathop{}\!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 $x_{i+1}(t):=x_{i}(t_{0})+\int_{t_{0}}^{t}f(s,x_{i}(s))\mathop{}\!ds$ . In other words, it finds a sequence of functions $x_{0}(t),x_{1}(t),\dots$ that approximate the solution $x(t)$ of the integral equation.

On the other hand, for a desired sequence of discrete times $t_{i}$, approximations to $x(t_{0}),x(t_{1}),\dots$ can be found by iterating the equation $x(t_{i+1})=x(t_{i})+\int_{t_{i}}^{t_{i+1}}f(s,x(s))\mathop{}\!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 $\Delta t$, the GBT update is

Three important cases are: $\alpha=0$ becomes the classic Euler method which is simply the first-order approximation $x(t+\Delta t)=x(t)+\Delta t\cdot x^{\prime}(t)$; $\alpha=1$ is called the backward Euler method; and $\alpha=\frac{1}{2}$ is called the bilinear method, which preserves the stability of the system .

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 $\overline{A}$ and $\overline{B}$ to be the matrices appearing in (3) for $\alpha=\frac{1}{2}$. Then the discrete-time state-space model is

In most models, the length of dependencies they can capture is roughly proportional to $\frac{1}{\Delta t}$. Thus we also refer to the step size $\Delta 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 . On the other hand, in Section 3.2 we show that the gating mechanism of classical RNNs is a version of learning $\Delta t$. Moreover when viewed as a CNN, the timescale $\Delta 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.

Linear State-Space Layers (LSSL)

We define our main abstraction, a model family that generalizes recurrence and convolutions. Section 3.1 first formally defines the LSSL, then discusses how to compute it with multiple views. Conversely, Section 3.2 shows that LSSLs are related to mechanisms of the most popular RNNs.

Given a fixed state space representation $A,B,C,D$, an LSSL is the sequence-to-sequence mapping defined by discretizing the linear state-space model (1) and (2).

Computationally, the discrete-time LSSL can be viewed in multiple ways (Fig. 1).

For simplicity let the initial state be $x_{-1}=0$. Then (4)+(5) explicitly yields

Then $y$ is simply the (non-circular) convolution $y=\mathcal{K}_{L}(\overline{A},\overline{B},C)\ast u+Du$, where

We make a note that the bottleneck of (i) the recurrence view is matrix-vector multiplication (MVM) by the discretized state matrix $\overline{A}$ when simulating (4), and (ii) the convolutional view is computing the Krylov function $\mathcal{K}_{L}$ (7). Throughout this section we assumed the LSSL parameters were fixed, which means that $\overline{A}$ and $\mathcal{K}_{L}(A,B,C)$ can be cached for efficiency. However, learning the parameters $\overline{A}$ and $\Delta t$ would involve repeatedly re-computing these, which is infeasible in practice. We revisit and solve this problem in Section 4.2.

2 Expressivity of LSSLs

For a model to be both recurrent and convolutional, one might expect it to be limited in other ways. Indeed, while also observe that certain recurrences can be replaced with a convolution, they note that it is not obvious if convolutions can be replaced by recurrences. Moreover, while the LSSL is a linear recurrence, popular RNN models are nonlinear sequence models with activation functions between each time step. We now show that LSSLs surprisingly do not have limited expressivity.

A well-known fact about state-space systems (1)+(2) is that the output $y$ is related to the input $u$ by a convolution $y(t)=\int h(\tau)u(t-\tau)d\tau$ with the impulse response $h$ of the system. Conversely, a convolutional filter $h$ that is a rational function of degree $N$ can be represented by a state-space model of size $N$ . Thus, an arbitrary convolutional filter $h$ can be approximated by a rational function (e.g., by Padé approximants) and represented by an LSSL.

In the particular case of LSSLs with HiPPO matrices (Sections 2 and 4.1), there is another intuitive interpretation of how LSSL relate to convolutions. Consider the special case when $A$ corresponds to a uniform measure (in the literature known as the LMU or HiPPO-LegT matrix). Then for a fixed $dt$, equation (1) is simply memorizing the input within sliding windows of $\frac{1}{\Delta t}$ elements, and equation (2) extracts features from this window. Thus the LSSL can be interpreted as automatically learning convolution filters with a learnable kernel width.

We show two results about RNNs that may be of broader interest. Our first result says that the ubiquitous gating mechanism of RNNs, commonly perceived as a heuristic to smooth optimization , is actually the analog of a step size or timescale $\Delta t$.

A (1-D) gated recurrence $x_{t}=(1-\sigma(z))x_{t-1}+\sigma(z)u_{t}$, where $\sigma$ is the sigmoid function and $z$ is an arbitrary expression, can be viewed as the GBT($\alpha=1$) (i.e., backwards-Euler) discretization of a 1-D linear ODE $\dot{x}(t)=-x(t)+u(t)$.

Applying a discretization requires a positive step size $\Delta t$. The simplest way to parameterize a positive function is via the exponential function $\Delta t=\exp(z)$ applied to any expression $z$. Substituting this into (3) with $A=-1,B=1,\alpha=1$ exactly produces the gated recurrence. ∎

While Lemma 3.1 involves approximating continuous systems using discretization, the second result is about approximating them using Picard iteration (Section 2). Roughly speaking, each layer of a deep linear RNN can be viewed as successive Picard iterates $x_{0}(t),x_{1}(t),...$ approximating a function $x(t)$ defined by a non-linear ODE. This shows that we do not lose modeling power by using linear instead of non-linear recurrences, and that the nonlinearity can instead be “moved” to the depth direction of deep neural networks to improve speed without sacrificing expressivity.

(Infinitely) deep stacked LSSL layers of order $N=1$ with position-wise non-linear functions can approximate any non-linear ODE $\dot{x}(t)=-x+f(t,x(t))$.

3 Deep LSSLs

The full LSSL architecture in a deep neural network is defined similarly to standard sequence models such as deep ResNets and Transformers, involving stacking LSSL layers connected with normalization layers and residual connections. Full architecture details are described in Appendix B, including the initialization of $A$ and $\Delta t$, computational details, and other architectural details.

Combining LSSLs with Continuous-time Memorization

In Section 3 we introduced the LSSL model and showed that it shares the strengths of convolutions and recurrences while also generalizing them. We now discuss and address its main limitations, in particular handling long dependencies (Section 4.1) and efficient computation (Section 4.2).

The generality of LSSLs means they can inherit the issues of recurrences and convolutions at addressing long dependencies (Section 1). For example, viewed as a recurrence, repeated multiplication by $\overline{A}$ could suffer from the vanishing gradients problem . We confirm empirically that LSSLs with random state matrices $A$ are actually not effective (Section 5.4) as a generic sequence model.

However, one advantage of these mathematical continuous-time models is that they are theoretically analyzable, and specific $A$ matrices can be derived to address this issue. In particular, the HiPPO framework (Section 2) describes how to memorize a function in continuous time with respect to a measure $\omega$ . This operator mapping a function to a continuous representation of its past is denoted $\operatorname{\mathsf{hippo}}(\omega)$, and was shown to have the form of equation (1) in three special cases. However, these matrices are non-trainable in the sense that no other $A$ matrices were known to be $\operatorname{\mathsf{hippo}}$ operators.

To address this, we theoretically resolve the open question from , showing that $\operatorname{\mathsf{hippo}}(\omega)$ for any measure $\omega$ To be precise, the measures that correspond to orthogonal polynomials . results in (1) with a structured matrix $A$.

For an arbitrary measure $\omega$, the optimal memorization operator $\operatorname*{hippo}(\omega)$ has the form $\dot{x}(t)=Ax(t)+Bu(t)$ (1) for a low recurrence-width (LRW) state matrix $A$.

For measures covering the classical orthogonal polynomials (OPs) (in particular, corresponding to Jacobi and Laguerre polynomials), there is even more structure.

For $\omega$ corresponding to the classical OPs, $\operatorname{\mathsf{hippo}}(\omega)$ is $3$-quasiseparable.

Although beyond the scope of this section, we mention that LRW matrices are a type of structured matrix that have linear MVM . In Appendix D we define this class and prove Theorem 1. Quasi-separable matrices are a related class of structured matrices with additional algorithmic properties. We define these matrices in Definition 4 and prove Corollary 4.1 in Section D.3.

Theorem 1 tells us that a LSSL that uses a state matrix $A$ within a particular class of structured matrices would carry the theoretical interpretation of continuous-time memorization. Ideally, we would be able to automatically learn the best $A$ within this class; however, this runs into computational challenges which we address next (Section 4.2). For now, we define the LSSL-fixed or LSSL-f to be one where the $A$ matrix is fixed to one of the HiPPO matrices prescribed by .

2 Theoretically Efficient Algorithms for the LSSL

Although $A$ and $\Delta t$ are the most critical parameters of an LSSL which govern the state-space (c.f. Section 4.1) and timescale (Sections 2 and 3.2), they are not feasible to train in a naive LSSL. In particular, Section 3.1 noted that it would require efficient matrix-vector multiplication (MVM) and Krylov function (7) for $\overline{A}$ to compute the recurrent and convolutional views, respectively. However, the former seems to involve a matrix inversion (3), while the latter seems to require powering $\overline{A}$ up $L$ times.

In this section, we show that the same restriction of $A$ to the class of quasiseparable (Corollary 4.1), which gives an LSSL the ability to theoretically remember long dependencies, simultaneously grants it computational efficiency.

First of all, it is known that quasiseparable matrices have efficient (linear-time) MVM . We show that they also have fast Krylov functions, allowing efficient training with convolutions.

We remark that Theorem 2 is non-obvious. To illustrate, it is easy to see that unrolling (7) for a general matrix $A$ takes time $LN^{2}$. Even if $A$ is extremely structured with linear computation, it requires $LN$ operations and linear depth. The depth can be reduced with the squaring technique (batch multiply by $A,A^{2},A^{4},\dots$), but this then requires $LN$ intermediate storage. In fact, the algorithm for Theorem 2 is quite sophisticated (Appendix E) and involves a divide-and-conquer recursion over matrices of polynomials, using the observation that (7) is related to the power series $C(I-Ax)^{-1}B$ .

Unless specified otherwise, the full LSSL refers to an LSSL with $A$ satisfying Corollary 4.1. In conclusion, learning within this structured matrix family simultaneously endows LSSLs with long-range memory through Theorem 1 and is theoretically computationally feasible through Theorem 2. We note the caveat that Theorem 2 is over exact arithmetic and not floating point numbers, and thus is treated more as a proof of concept that LSSLs can be computationally efficient in theory. We comment more on the limitations of the LSSL in Section 6.

Empirical Evaluation

We test LSSLs empirically on a range of time series datasets with sequences from length 160 up to 38000 (Sections 5.1 and 5.2), where they substantially improve over prior work. We additionally validate the computational and modeling benefits of LSSLs from generalizing all three main model families (Section 5.3), and analyze the benefits of incorporating principled memory representations that can be learned (Section 5.4).

Our tasks have extensive prior work and we evaluate against previously reported best results. We highlight our primary baselines, three very recent works explicitly designed for long sequences: CKConv (a continuous-time CNN) , UnICORNN (an ODE-inspired RNN) , and Neural Controlled/Rough Differential Equations (NCDE/NRDE) (a sophisticated NDE) . These are the only models we are aware of that have experimented with sequences of length >10k.

1 Image and Time Series Benchmarks

We test on the sequential MNIST, permuted MNIST, and sequential CIFAR tasks (Fig. 3), popular benchmarks which were originally designed to test the ability of recurrent models to capture long-term dependencies of length up to 1k . LSSL sets SoTA on sCIFAR by more than 10 points. We note that all results were achieved with at least 5x fewer parameters than the previous SoTA (Appendix F).

We additionally use the BDIMC healthcare datasets (Fig. 3), a suite of widely studied time series regression problems of length 4000 on estimating vital signs. LSSL reduces RMSE by more than two-thirds on all datasets.

2 Speech and Image Classification for Very Long Time Series

Raw speech is challenging for ML models due to high-frequency sampling resulting in very long sequences. Traditional systems involve complex pipelines that require feeding mixed-and-matched hand-crafted features into DNNs . Table 2 reports results for the Speech Commands (SC) dataset for classification of 1-second audio clips. Few methods have made progress on the raw speech signal, instead requiring pre-processing with standard mel-frequency cepstrum coefficients (MFCC). By contrast, LSSL sets SoTA on this dataset while training on the raw signal. We note that MFCC extracts sliding window frequency coefficients and thus is related to the coefficients $x(t)$ defined by LSSL-f (Section 2, Section 4.1, , Appendix D). Consequently, LSSL may be interpreted as automatically learning MFCC-type features in a trainable basis.

To stress-test the LSSL’s ability to handle extremely long sequences, we create a challenging new sequential-CelebA task, where we classify $178\times 218$ images = 38000-length sequences for 4 facial attributes: Attractive (Att.), Mouth Slightly Open (MSO), Smiling (Smil.), Wearing Lipstick (WL) . We chose the 4 most class-balanced attributes to avoid well-known problems with class imbalance. LSSL-f comes close to matching the performance of a specialized ResNet-18 image classification architecture that has $10\times$ the parameters (Table 1). We emphasize we are the first to demonstrate that this is possible to do with a generic sequence model.

3 Advantages of Recurrent, Convolutional, and Continuous-time Models

We validate that the generality of LSSLs endows it with the strengths of all three families.

Convergence Speed. As a recurrent and NDE model that incorporates new theory for continuous-time memory (Section 4.1), the LSSL has strong inductive bias for sequential data, and converges rapidly to SoTA results on our benchmarks. With its convolutional view, training can be parallelized and it is also computationally efficient in practice. Table 3 compares the time it takes the LSSL-f to achieve SoTA, in either sample (measured by epochs) or computational (measured by wall clock) complexity. In all cases, LSSLs reached the target in a fraction of the time of the previous model.

Timescale Adaptation. Table 2 also reports the results of continuous-time models that are able to handle unique settings such as missing data in time series, or test-time shift in timescale (we note that this is a realistic problem, e.g., when deployed healthcare models are tested on EEG signals that are sampled at a different rate ). We note that many of these baselines were custom designed for such settings, which is of independent interest. On the other hand, LSSLs perform timescale adaptation by simply changing its $\Delta t$ values at inference time, while still outperforming the performance of prior methods with no shift. Additional results on the CharacterTrajectories dataset from prior work are in Appendix F, where LSSL is competitive with the best baselines.

4 LSSL Ablations: Learning the Memory Dynamics and Timescale

We demonstrate that the $\Delta t$ and $A$ parameters, which LSSLs are able to automatically learn in contrast to prior work, are indeed critical to the performance of these continuous-time models. We note that learning $\Delta t$ adds only $O(H)$ parameters and learning $A$ adds $O(N)$ parameters, adding less than 1% parameter count compared to the base models with $O(HN)$ parameters.

Memory dynamics $A$. We validate that vanilla LSSLs suffer from the modeling issues described in Section 4. We tested that LSSLs with random $A$ matrices (normalized appropriately) perform very poorly (e.g., 62% on pMNIST). Further, we note the consistent increase in performance from LSSL-f to LSSL despite the negligible parameter difference. These ablations show that (i) incorporating the theory of Theorem 1 is actually necessary for LSSLs, and (ii) further training the structured $A$ is additionally helpful, which can be interpreted as learning the measure for memorization (Section 4.1).

Timescale $\Delta t$. Section 3.2 showed that LSSL’s ability to learn $\Delta t$ is its direct generalization of the critical gating mechanism of popular RNNs, which previous ODE-based RNN models cannot learn. We note that on sCIFAR, LSSL-f with poorly-specified $\Delta t$ gets only $49.3\%$ accuracy. Additional results in Appendix F show that learning $\Delta t$ alone provides an orthogonal boost to learning $A$, and visualizes the noticeable change in $\Delta t$ over the course of training.

Discussion

In this work we introduced a simple and principled model (LSSL) inspired by a fundamental representation of physical systems. We showed theoretically and empirically that it generalizes and inherits the strengths of the main families of modern time series models, that its main limitations of long-term memory can be resolved with new theory on continuous-time memorization, and that it is empirically effective on difficult tasks with very long sequences.

Related work. The LSSL is related to several rich lines of work on recurrent, convolutional, and continuous-time models, as well as sequence models addressing long dependencies. Appendix A provides an extended related work connecting these topics.

Tuning. Our models are very simple, consisting of identical L(inear)SSL layers with simple position-wise non-linear modules between layers (Appendix B). Our models were able to train at much higher learning rates than baselines and were not sensitive to hyperparameters, of which we did light tuning primarily on learning rate and dropout. In contrast to previous baselines , we did not use hyperparameters for improving stability and regularization such as weight decay, gradient clipping, weight norm, input dropout, etc. While the most competitive recent works introduce at least one hyperparameter of critical importance (e.g. depth and step size , $\alpha$ and $\Delta t$ , $\omega_{0}$ ) that are difficult to tune, the LSSL-fixed has only $\Delta t$, which the full LSSL can even learn automatically (at the expense of speed).

Limitations. Sections 1, 3 and 1 mention that a potential benefit of having the recurrent representation of LSSLs may endow it with efficient inference. While this is theoretically possible, this work did not experiment on any applications that leverage this. Follow-up work showed that it is indeed possible in practice to speed up some applications at inference time.

Theorem 2’s algorithm is sophisticated (Appendix D) and was not implemented in the first version of this work. A follow-up to this paper found that it is not numerically stable and thus not usable on hardware. Thus the algorithmic contributions in Theorem 2 serve the purpose of a proof-of-concept that fast algorithms for the LSSL do exist in other computation models (i.e., arithmetic operations instead of floating point operations), and leave an open question as to whether fast, numerically stable, and practical algorithms for the LSSL exist.

As described in Appendix B, by freezing the $A$ matrix and $\Delta t$ timescale, the LSSL-fixed is able to be computed much faster than the full LSSL, and is comparable to prior models in practice (Table 3). However, beyond computational complexity, there is also a consideration of space efficiency. Both the LSSL and LSSL-fixed suffer from a large amount of space overhead (described in Appendix B) – using $O(NL)$ instead of $O(L)$ space when working on a 1D sequence of length $L$ – that essentially stems from using the latent state representation of dimension $N$. Consequently, the LSSL can be space inefficient and we used multi-GPU training for our largest experiments (speech and high resolution images, Tables 2 and 1).

These fundamental issues with computation and space complexity were revisited and resolved in follow-up work to this paper, where a new state space model (the Structured State Space) provided a new parameterization and algorithms for state spaces.

Conclusion and future work. Modern deep learning models struggle in applications with very long temporal data such as speech, videos, and medical time-series. We hope that our conceptual and technical contributions can lead to new capabilities with simple, principled, and less engineered models. We note that our pixel-level image classification experiments, which use no heuristics (batch norm, auxiliary losses) or extra information (data augmentation), perform similar to early convnet models with vastly more parameters, and is in the spirit of recent attempts at unifying data modalities with a generic sequence model . Our speech results demonstrate the possibility of learning better features than hand-crafted processing pipelines used widely in speech applications. We are excited about potential downstream applications, such as training other downstream models on top of pre-trained state space features.

We thank Arjun Desai, Ananya Kumar, Laurel Orr, Sabri Eyuboglu, Dan Fu, Mayee Chen, Sarah Hooper, Simran Arora, and Trenton Chang for helpful feedback on earlier drafts. We thank David Romero and James Morrill for discussions and additional results for baselines used in our experiments. This work was done with the support of Google Cloud credits under HAI proposals 540994170283 and 578192719349. AR and IJ are supported under NSF grant CCF-1763481. KS is supported by the Wu Tsai Neuroscience Interdisciplinary Graduate Fellowship. We gratefully acknowledge the support of NIH under No. U54EB020405 (Mobilize), NSF under Nos. CCF1763315 (Beyond Sparsity), CCF1563078 (Volume to Velocity), and 1937301 (RTML); ONR under No. N000141712266 (Unifying Weak Supervision); ONR N00014-20-1-2480: Understanding and Applying Non-Euclidean Geometry in Machine Learning; N000142012275 (NEPTUNE); the Moore Foundation, NXP, Xilinx, LETI-CEA, Intel, IBM, Microsoft, NEC, Toshiba, TSMC, ARM, Hitachi, BASF, Accenture, Ericsson, Qualcomm, Analog Devices, the Okawa Foundation, American Family Insurance, Google Cloud, Salesforce, Total, the HAI-AWS Cloud Credits for Research program, the Stanford Data Science Initiative (SDSI), and members of the Stanford DAWN project: Facebook, Google, and VMWare. The Mobilize Center is a Biomedical Technology Resource Center, funded by the NIH National Institute of Biomedical Imaging and Bioengineering through Grant P41EB027060. The U.S. Government is authorized to reproduce and distribute reprints for Governmental purposes notwithstanding any copyright notation thereon. Any opinions, findings, and conclusions or recommendations expressed in this material are those of the authors and do not necessarily reflect the views, policies, or endorsements, either expressed or implied, of NIH, ONR, or the U.S. Government.

References

Appendix A Related Work

We provide an extended related work comparing the LSSL to previous recurrent, convolutional, and continuous-time models.

The LSSL is most closely related to the HiPPO framework for continuous-time memory and its predecessor, the Legendre Memory Unit (LMU) . The HiPPO-RNN and the LMU define dynamics of the form of equation (1), and incorporate it into an RNN architecture. A successor to the LMU, the LMU-FFT keeps the original linear dynamics, allowing the LMU to be computed with a cached convolution kernel.

These methods all suffer from two main limitations. First, the state matrix $A$ and discretization timescale $\Delta t$ cannot be trained due to both limitations in theoretical understanding of which $A$ matrices are effective, as well as computational limitations. Second, (1) is a 1-D to $N$-D map, requiring states to be projected back down to 1-D. This creates an overall 1-D bottleneck in the state, limiting the expressivity of the model.

Compared to these, the LSSL does not use a conventional RNN architecture, instead keeping the linear recurrence (4) and downprojecting it with the second part of the state space representation (5). To avoid the 1-D feature bottlneck, it simply computes $H$ copies of this 1-D to 1-D independently, creating an overall $H$-dimensional sequence-to-sequence model. However, this exacerbates the computational issue, since the work is increased by a factor of $H$.

This work resolves the expressivity issue with new theory. Compared to HiPPO and the LMU, LSSL allows training the $A$ matrix by showing generalized theoretical results for the HiPPO framework, showing that there is a parameterized class of structured state spaces that are HiPPO operators.

The LSSL makes progress towards the second issue with new algorithms for these structured matrices (Theorem 2). However, as noted in Sections 4.2 and 6, the algorithm presented in Theorem 2 was later found to be not practical, and an improved representation and algorithm was found in subsequent work.

The connection from RNNs to continuous-time models have been known since their inception, and recent years have seen an explosion of CT-RNN (continuous-time RNN) models based on dynamical systems or ODEs. We briefly mention a few classic and modern works along these lines, categorizing them into a few main topics.

First are theoretical works that analyze the expressivity of RNNs from a continuous-time perspective. The connection between RNNs and dynamical systems has been studied since the 90s , fleshing out the correspondence between different dynamical systems and RNN architectures . Modern treatments have focused on analyzing the stability and dynamics of RNNs.

Second, a large class of modern RNNs have been designed that aim to combat vanishing gradients from a dynamical systems analysis. These include include the AntisymmetricRNN , iRNN , and LipschitzRNN , which address the exploding/vanishing gradient problem by reparatermizing the architecture or recurrent matrix based on insights from an underlying dynamical system.

Third is a class of models that are based on an explicit underlying ODE introduced to satisfy various properties. This category includes the UnICORNN and its predecessor coRNN which discretize a second-order ODE inspired by oscillatory systems. Other models include the Liquid Time-Constant Networks (LTC) and successor CfC , which use underlying dynamical systems with varying time-constants with stable behavior and provable rates of expressivity measured by trajectory length. The LTC is based on earlier dynamic causal models (DCM) , which are a particular ODE related to state spaces with an extra bilinear term. Finally, the LMU and HiPPO also fall in this category, whose underlying ODEs are mathematically derived for continuous-time memorization.

Fourth, the recent family of neural ODEs , originally introduced as continuous-depth models, have been adapted to continuous-time, spawning a series of “ODE-RNN” models. Examples include the ODE-RNN , GRU-ODE-Bayes , and ODE-LSTM , which extend adjoint-based neural ODEs to the discrete input setting as an alternative to standard RNNs. Neural Controlled Differential Equations (NCDE) and Neural Rough Differential Equations (NRDE) are memory efficient versions that integrate observations more smoothly and can be extended to very long time series.

As a special case of continuous-time RNNs, some works have observed the relation between gating mechanisms and damped dynamical systems . Some examples of continuous-time RNNs based on such damped dynamical systems include the LTC and iRNN . Compared to these, Lemma 3.1 shows a stronger result that sigmoid gates are not just motivated by being an arbitrary monotonic function with range $(0,1)$, but the exact formula appears out of discretizing a damped ODE.

Appendix B Model Details

Section 3.1 noted that some of the computations for using the LSSL are expensive to compute. When the LSSL fixes the $A$ and $\Delta t$ parameters (e.g. when they are not trained, or at inference time), these computational difficulties can be circumvented by caching particular computations. In particular, this case applies to the LSSL-f. Note that in this case, the other state-space matrices $C$ and $D$ comprise the $O(HN)$ trainable parameters of the fixed-transition LSSL.

In particular, we assume that there is a black-box inference algorithm for this system, i.e. matrix-vector multiplication by $\overline{A}$ (an example of implementing this black box for a particular structured class is in Section E.2). We then compute and cache

the transition matrix $\overline{A}$, which is computed by applying the black-box $\overline{A}$ MVM algorithm to the identity matrix $I$.

which is computed in a parallelized manner by the squaring technique for exponentiation, i.e. batch multiply by $\overline{A},(\overline{A})^{2},(\overline{A})^{4},\dots$.

At inference time, the model can be unrolled recurrently with $\overline{A}$. At training time, the convolutional filter $K_{L}(\overline{A},\overline{B},C)$ (equation (7)) is computed with a matrix multiplication $C\cdot K(\overline{A},\overline{B})$ before convolving with the input $u$.

Table 5 provides more detailed complexity of this version of the LSSL with fixed $A,\Delta t$.

Note that as mentioned in Section 6, this cached algorithm is fairly fast, but the main drawback is that materializing the Krylov matrix (8) requires $O(NL)$ instead of $O(L)$ space.

B.2 Initialization of A𝐴A

Note that the LSSL-f is the LSSL with a non-trainable $A$ (and $\Delta t$), so that $A$ is fixed to the above matrix.

B.3 Initialization of Δ​tΔ𝑡\Delta t

One distinction between the LSSL and the most related prior work is that the inclusion of the projection (2) makes the layer a 1-dimensional to 1-dimensional map, instead of 1-D to $N$-D . This enables us to concatenate $H$ copies of this map (at the expense of computation, cf. Sections 4.2 and D). Even when $\Delta t$ is not trained as in the LSSL-f, these $H$ copies allow multiple timescales to be considered by setting $\Delta t$ differently for each copy.

In particular, we initialize $\Delta t$ log-uniformly in a range $\Delta t_{min},\Delta t_{max}$ (i.e., $\Delta t$ is initialized within this range, such that $\log\Delta t$ is uniformly distributed). The maximum and minimum values were generally chosen to be a factor of $100$ apart such that the length of the sequences in the dataset are contained in this range. Specific values for each model and dataset are in Appendix F. We did not search over these as a hyperparameter, but we note that it can be tuned for additional performance improvements in our experiments.

B.4 Deep Neural Network Architecture

The Deep LSSL models used in our experiments simply stack together LSSL layers in a simple deep neural network architecture. We note the following architecture details.

We call $M$ the number of channels in the model.

There are two drawbacks with the current definition of LSSL:

They are defined by running $H$ independent copies of a state-space model, which means the $H$ input features do not interact at all.

If the channel dimension is $M>1$, then the LSSL is a map from dimension $1$ to $M$, which means residuals cannot be applied.

These are both addressed by introducing a position-wise feedforward layer after the LSSL of shape $H\cdot M\to H$. This simultaneously mixes the hidden features, and projects the output back to dimension $1$ if necessary. There is also an optional non-linearity in between the LSSL and this feedforward projection; we fix it to the GeLU activation function in our models.

We note that this factorization of parallel convolutions on the $H$ features followed by a position-wise linear map is very similar to depth-wise separable convolutions .

To stack multiple layers of LSSLs together, we use very standard architectures for deep neural networks. In particular, we use residual connections and a layer normalization (either pre-norm or post-norm) in the style of standard Transformer architectures. Whether to use pre-norm or post-norm was chosen on a per-dataset basis, and depended on whether the model overfit; recent results have shown that pre-norm architectures are more stable , so we used it on harder datasets with less overfitting. We note that we could have additionally inserted MLP modules in between LSSL layers, in the style of Transformers , but did not experiment with this.

The overall parameter count of an LSSL model is $M\cdot H\cdot(H+N)$.

We primarily used two model sizes in our experiments, which were chosen simply to produce round numbers of parameters:

LSSL small ($\approx 200K$ parameters): $6$ layers, $H=128,N=128,M=1$.

LSSL large ($\approx 2M$ parameters): $4$ layers, $H=256,N=256,M=4$.

We did not search over additional sizes, but for some datasets reduced the model size for computational reasons.

Appendix C LSSL Proofs

This section gives refinements of the statements in Section 3, additional results, and proofs of all results.

Section C.1 has a more detailed (and self-contained) summary of basic methods in ODE approximation which will be used in the results and proofs.

Section C.2 give more general statements and proofs of Lemma 3.1 and Lemma 3.2 in Lemma C.1 and Theorem 4, respectively.

We consider the standard setting of a first-order initial value problem (IVP) ordinary differential equation (ODE) for a continuous function $f(t,x)$

This differential form has an equivalent integral form

Sections C.1.1 and C.1.2 overview the Picard theorem and first-order numerical integration methods, which apply to any IVP (9). Section C.1.3 then shows how to specialize it to linear systems as in equation (1).

At a high level, the basic approximation methods considered here use the integral form (10) and approximate the integral in the right-hand side by simple techniques.

The Picard-Lindelöf Theorem gives sufficient conditions for the existence and uniqueness of solutions to an IVP. As part of the proof, it provides an iteration scheme to compute this solution.

In the IVP (9), if there is an interval around $t_{0}$ such that $f$ is Lipschitz in its second argument, then there is an open interval $I\ni t_{0}$ such that there exists a unique solution $x(t)$ to the IVP in $I$. Furthermore, the sequence of Picard iterates $x^{(0)},x^{(1)},\dots$ defined by

C.1.2 Numerical Integration Methods

Many methods for numerical integration of ODEs exist, which calculate discrete-time approximations of the solution. We discuss a few of the simplest methods, which are first-order methods with local error $O(h^{2})$ .

These methods start by discretizing (10) into the form

Here we assume a sequence of discrete times $t_{0},t_{1},t_{2},\dots$ is fixed. For convenience, let $x_{k}$ denote $x(t_{k})$ and let $\Delta t_{k}:=t_{k}-t_{k-1}$. The goal is now to approximate the integral in the RHS of (11).

The Euler method approximates (11) by holding the left endpoint constant throughout the integral (i.e., the “rectangle rule” with left endpoint), $f(s,x(s))\approx f(t_{k-1},x(t_{k-1}))$. The discrete-time update becomes

The backward Euler method approximates (11) by holding the right endpoint constant throughout the integral (i.e., the “rectangle rule” with right endpoint), $f(s,x(s))\approx f(t_{k},x(t_{k}))$. The discrete-time update becomes

C.1.3 Discretization of State-Space Models

In the case of a linear system, the IVP is specialized to the case

Note that here $u$ is treated as a fixed external input, which is constant from the point of view of this ODE in $x$. Let $u_{k}$ denote the average value in each discrete time interval,

The integral equation (11) can be specialized to this case, and more generally a convex combination of the left and right endpoints can be taken to approximate the integral, weighing them by $1-\alpha$ and $\alpha$ respectively. Note that the case $\alpha=0,1$ are specializations of the forward and backward Euler method, and the case $\alpha=\frac{1}{2}$ is the classic “trapezoid rule” for numerical integration.

This derives the generalized bilinear transform (GBT) . The bilinear method is the case $\alpha=\frac{1}{2}$ of special significance, and was numerically found to be better than the forward and backward Euler methods $\alpha=0,1$ both in synthetic function approximation settings and in end-to-end experiments [24, Figure 4].

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

We provide more detailed statements of Lemmas 3.1 and 3.2 from Section 3.2. In summary, LSSLs and popular families of RNN methods all approximate the same continuous-time dynamics

by viewing them with a combination of two techniques.

We note that these results are about two of the most commonly used architecture modifications for RNNs. First, the gating mechanism is ubiquitous in RNNs, and usually thought of as a heuristic for smoothing optimization . Second, many of the effective large-scale RNNs use linear (gated) recurrences and deeper models, which is usually thought of as a heuristic for computational efficiency . Our results suggest that neither of these are heuristics after all, and arise from standard ways to approximate ODEs.

Non-linear RNNs discretize the dynamics (14) by applying backwards Euler discretization to the linear term, which arises in the gating mechanism of RNNs (Section C.2.2, Lemma C.1).

A special case of LSSLs approximates the dynamics (14) (in continuous-time) by applying Picard iteration to the non-linear term (Section C.2.3, Theorem 4).

Deep linear RNNs approximate the dynamics (14) with both Picard iteration in the depth direction to linearize the non-linear term, and discretization (gates) in the time direction to discretize the equation (Section C.2.4, Corollary C.3).

In the remainder of this section, we assume that there is an underlying function $x(t)$ that satisfies (14) on some interval for any initial condition, and that $f$ is continuous and Lipschitz in its second argument. Our goal is to show that several families of models approximate this in various ways.

We sketch the idea of how LSSLs capture popular RNNs. More precisely, we will show how approximating the dynamics (14) in various ways lead to types of RNNs and LSSLs.

The first step is to look at the simpler dynamics

where there is some input $u(t)$ that is independent of $x$. (In other words, in (14), the function $f(t,x)$ does not depend on the second argument.)

By directing applying the GBT discretization with $\alpha=1$, this leads to a gated recurrence (Lemma 3.1).

The second step is that by applying the backwards Euler discretization more directly to (14), this leads to a gated RNN where the input can depend on the state (Lemma C.1).

Alternatively, we can apply Picard iteration on (14), which says that the iteration

Differentiating, this equation is the ODE

To move to discrete-time, this continuous-time layer can be discretized with gates as in Lemma 3.1, leading to deep linear RNNs such as the QRNN, or with the bilinear discretization, leading to the discrete-time LSSL. We note again that in the discrete-time LSSL, $\overline{A}$ and $\overline{B}$ play the role of the gates $\sigma,1-\sigma$.

C.2.2 Capturing gates through discretization

where $\overline{f}(k,x)$ is an arbitrary function that is Lipschitz in its second argument (e.g., it may depend on an external input $u_{k}$).

Then equation (15) is a discretization of the dynamics (14) with step sizes $\Delta t_{k}=\exp(z_{k})$, i.e. $x_{k}\approx x(t_{k})$ where $t_{k}=\sum_{i=1}^{k}\Delta t_{i}$.

Apply the backwards Euler discretization (13) to equation (14) to get

Note that $\frac{\Delta t_{k}}{1+\Delta t_{k}}=\frac{e^{z_{k}}}{1+e^{z_{k}}}=\frac{1}{1+e^{-z_{k}}}$ and $\frac{1}{1+\Delta t_{k}}=1-\frac{\Delta t_{k}}{1+\Delta t_{k}}$, thus

Here we are denoting $\overline{f}(k,x)=f(t_{k},x)$ to be a discrete-time version of $f$ evaluatable at the given timesteps $t_{k}$. ∎

Note that a potential external input function $u(t)$ or sequence $u_{k}$ is captured through the abstraction $f(t,x)$. For example, a basic RNN could define $\overline{f}(k,x)=f(t_{k},x)=\text{tanh}(Wx+Uu_{k})$.

C.2.3 Capturing non-linearities through Picard iteration

The main result of this section is Theorem 4 showing that LSSLs can approximate the same dynamics as the RNNs in the previous section. This follows from a technical lemma.

Let $f(t,x)$ be any function that satisfies the conditions of the Picard-Lindelöf Theorem (Theorem 3).

A (continuous-time) deep LSSL with order $N=1$ and $A=-1,B=1,C=1,D=0$ approximates the non-linear dynamics (14).

Applying the definition of an LSSL (equations (1)+(2)) with these parameters results in a layer mapping $u(t)\mapsto y(t)$ where $y$ is defined implicitly through the ODE

This can be seen since the choice of $C,D$ implies $y(t)=x(t)$ and the choice of $A,B$ gives the above equation.

Consider the deep LSSL defined by alternating this LSSL with position-wise (in time) non-linear functions

(and $z_{0}=z(t_{0})=x(t_{0})=x_{0}$). Note that

Since $f$ satisfies the conditions of the Picard Theorem (i.e., is continuous in the first argument and Lipschitz in the second), so does the function $g$ where $g(t,x):=e^{-At}f(t,e^{At}x)$ for some interval around the initial time.

C.2.4 Capturing Deep, Linear, Gated RNNs

We finally note that several types of RNNs exist which were originally motivated by approximating linearizing gated RNNs for speed. Although these were treated as a heuristic for efficiency reasons, they are explained by combining our two main technical results.

Lemma C.1 shows that a single-layer, discrete-time, non-linear RNN approximates the dynamics (14) through discretization, which arises in the gating mechanism.

Theorem 4 shows that a deep, continuous-time, linear RNN approximates (14) through Picard iteration, where the non-linearity is moved to the depth direction.

Combining these two results leads to Corollary C.3, which says that a deep, discrete-time, linear RNN can also approximate the same dynamics (14).

This is a discretization of the dynamics (14) with step sizes $\Delta t_{k}=\exp(z_{k})$, i.e. $x_{k}\approx x(t_{k})$ where $t_{k}=\sum_{i=1}^{k}\Delta t_{i}$.

By Lemma C.1, the first equation is a discretization of the continuous-time equation

uses the continuous-time version $f$ of $\overline{f}$. But by Lemma C.2, this is an approximation of the dynamics (14) using Picard iteration. ∎

Notable examples of this type of model include the Quasi-RNN or QRNN and the Simple Recurrent Unit (SRU) , which are among the most effective models in practice. We remark that these are the closest models to the LSSL and suggest that their efficacy is a consequence of the results of this section, which shows that they are not heuristics.

We note that there are many more RNN variants that use a combination of these gating and linearization techniques that were not mentioned in this section, and can be explained similarly.

Appendix D LSSL Proofs and Algorithms

This section proves the results in Section 4.1, and is organized as follows:

Section D.1 gives a self-contained synopsis of the HiPPO framework .

Section D.2 proves Theorem 1, which shows that the $\operatorname{\mathsf{hippo}}$ operators for any measure lead to a simple linear ODE of the form of equation (1).

Section D.3 proves Corollary 4.1, including a formal definition of quasiseparable matrices (i.e., how LSSL matrices are defined) in Definition 4.

This section is technically involved and we adopt notation to simplify reasoning about the shapes of objects. In particular, we use bold capitals (e.g. $\mathbf{A}$) to denote matrices and bold lowercase (e.g. $\mathbf{b}$) to denote vectors. For example, equation (1) becomes $\dot{x}=\mathbf{A}x+\mathbf{b}u$. These conventions are adopted throughout Appendices D and E.

D.1 Preliminaries: HiPPO Framework and Recurrence Width

This section summarizes technical preliminaries taken directly from prior work. We include this section so that this work is self-contained and uses consistent notation, which may deviate from prior work. For example, we use modified notation from Gu et al. in order to follow conventions in control theory (e.g., we denote input by $u$ and state by $x$ as in (1)).

Section D.1.1 formally defines the HiPPO operator mathematically as in [24, Section 2.2], and Section D.1.2 overviews the steps to derive the HiPPO operator as in [24, Appendix C]. Section D.1.3 defines the class of Low Recurrence Width (LRW) matrices, which is the class of matrices that our generalization of the HiPPO results (Theorem 1) uses.

$\mathsf{proj}_{t}$ takes a function $u$ restricted up to time $t$, $u_{\leq t}:=u(x)|_{x\leq t}$, and maps it to a polynomial $g^{(t)}\in\mathcal{G}$, that minimizes the approximation error $\|u_{\leq t}-g^{(t)}\|_{L_{2}(\mu^{(t)})}$.

D.1.2 HiPPO Framework for Deriving the HiPPO Operator

The main ingredients of HiPPO consists of an approximation measure and an orthogonal polynomial basis. We recall how they are defined in (we note that compared to Gu et al. , our notation has changed from input $f(t)$ coefficients (state) $c(t)$ to input $u(t)$ and coefficients (state) $x(t)$, following conventions in controls).

At every $t$, the approximation quality is defined with respect to a measure $\mu^{(t)}$ supported on $(-\infty,t]$. We assume that the measures $\mu^{(t)}$ have densities $\omega(t,Y):=\frac{\mathop{}\!d\mu^{(t)}}{\mathop{}\!dY}$. Note that this implies that integrating with respect to $\mathop{}\!d\mu^{(t)}$ is the same as integrating with respect to $\omega(t,Y)\mathop{}\!dY$.

In the general framework, HiPPO does not require an orthogonal polynomial basis as the selected basis. The choice of basis is generalized by tilting with $\chi$.

We express the coefficients $x_{n}(t)$ calculated by the HiPPO framework as:

To use this to derive $\dot{x}_{n}(t)$, let $h(t,Y)=u(Y)p_{n}(t,Y)\omega(t,Y)$. We see that

This allows $\dot{x}_{n}(t)$ to be written as

Although Gu et al. describe the framework in the full generality above and use $\chi$ as another degree of freedom, in their concrete derivations they always fix $\chi=\omega$. Our general results also use this setting. For the remainder of this section, we assume the “full tilting” case $\chi=\omega$. In particular, this means that in Eq. 18, we essentially substitute $\omega$ above with $1$ and divide each term by the inverse square root of our normalization constant, $\zeta$, to get the coefficient dynamics that we will use in our arguments:

Now, if we can show that each of the integrated terms in (18) are linear combinations of $x_{n}(t)$, this would be the same as saying that $\dot{x}_{n}(t)=\mathbf{A}(t)x(t)+\mathbf{b}(t)u(t)$ for some $\mathbf{A}(t)$. Therefore, the incremental update operation would be bounded by the runtime of the matrix-vector operation $\mathbf{A}(t)x(t)$.

D.1.3 Recurrence Width

An $N\times N$ matrix $\mathbf{A}$ has recurrence width $t$ if the polynomials $a_{i}(X)=\sum_{j=0}^{N-1}\mathbf{A}[i,j]X^{j}$ satisfy $\deg(a_{i})\leq i$ for $i<t$, and

D.2 Proof of Theorem 1

This section proves Theorem 1, which is restated formally in Corollary D.4. Section D.2.1 proves some results relating orthogonal polynomials to recurrence width (Section D.1.3). Section D.2.2 proves Corollary D.4. Sections D.2.3 and D.2.4 provides examples showing how Corollary D.4 can be specialized to exactly recover the HiPPO-LegT, HiPPO-LagT, HiPPO-LegS methods .

Next we introduce the following lemma, which will be useful in our arguments:

For any $n$, there exists ordered sets of coefficients $\alpha_{n}=\{\alpha_{n,i}\},\beta_{n}=\{\beta_{n,i}\}$,

$p^{\prime}_{n}(Z)=\sum_{i=0}^{n-1}\alpha_{n,i}p_{i}(Z)$

$Zp^{\prime}_{n}(Z)=\sum_{i=0}^{n-1}\beta_{n,i}p_{i}(Z)$

Follows from the fact that $p_{i}(z)$ for $0\leq i<N$ forms a basis and the observation of the degrees of the polynomials on the LHS. ∎

Let $\mathbf{S}$ be the “right shift" matrix, i.e. for any matrix $\mathbf{M}$, $\mathbf{M}\mathbf{S}$ has the columns of $\mathbf{M}$ shifted to right by one. Note that $\mathbf{S}^{T}$ corresponds to the “left shift" matrix.

Recall that $\mathbf{P}$ has the coefficients of the OP polynomials $p_{0}(Z),\dots,p_{N-1}(Z)$ as its rows. Then note that

The claim on the structure of $\mathbf{D}_{2}$ then follows from the above expression. The claim on runtime of computing $\mathbf{D}_{2}\mathbf{z}$ follows from essentially the same argument as for $\mathbf{D}_{1}\mathbf{z}$.

Finally, we make the following observation:

Let $\mathbf{A}^{\prime}$ and $\mathbf{B}^{\prime}$ be defined such that $\mathbf{A}^{\prime}[n,i]=\alpha_{n,i}$ and $\mathbf{B}^{\prime}[n,i]=\beta_{n,i}$. Then both $\mathbf{A}^{\prime}$ and $\mathbf{B}^{\prime}$ are both products of three matrices: two of which have recurrence width at most $2$ and the third is the inverse of a matrix that has recurrence width $2$.

D.2.2 HiPPO for General Measures

From the definition of $z$, we see that ${Y}-t=\frac{(z-1)\theta(t)}{2}$. Then

Additionally, given a measure $\omega$ on and OP family $p_{0}({Y}),\ p_{1}({Y}),\dots$ such that for all $i\neq j$,

If we let $\alpha(t)=t-\theta(t)$ and $\beta(t)=t$, then applying the Leibniz rule to (23) we get:

Because $\deg\left({p^{\prime}_{n}(z)}\right)\leq n-1$ and $\deg\left({(z-1)p^{\prime}_{n}(z)}\right)\leq n$, they can be written as a linear combination of $\{p^{\prime}_{i}\}_{i\leq n}$. Let us define $\{\alpha_{n,j}\}$, $\{\beta_{n,j}\}$ such that

where $\mathbf{A}_{1}(t)[n,k]=\begin{cases}2\alpha_{n,k}+\theta^{\prime}(t)\beta_{n,k}&\text{ if }k<n\\ \theta^{\prime}(t)\beta_{n,n}+\theta^{\prime}(t)&\text{ if }k=n\\ 0\ \ &\text{ otherwise }\end{cases}$ for $\alpha_{n,k}$,$\beta_{n,k}$ as defined in (25).

The matrix $\mathbf{A}_{1}$ in Theorem 7 can be re-written as

In particular, both $\mathbf{A}^{\prime}$ and $\mathbf{B}^{\prime}$ both products of three matrices: two of which have recurrence width at most $2$ and the third is the inverse of a matrix that has recurrence width $2$.

Eq. 26 follows from Theorem 7 and defining $\mathbf{A}^{\prime}$ and $\mathbf{B}^{\prime}$ to contain the $\alpha_{n,k}$ and $\beta_{n,k}$ coefficients. ∎

D.2.3 Translated HiPPO (Sliding Windows)

The case when $\theta(t)=\theta$ for all $t$ represents a constant-size sliding window, which Gu et al. denote as the “Translated HiPPO” case with instantiations such as HiPPO-LegT (Translated Legendre) and HiPPO-LagT (Translated Laguerre).

We now state a corollary of Theorem 7 for the case of $\theta(t)=\theta$ for all t.

where $\mathbf{A}_{1}[n,j]=\begin{cases}2\alpha_{n,k}&\text{ if }k<n\\ 0\ \ &\text{ otherwise }\end{cases}$.

to handle the $u(t-\theta)$ term in Corollary D.5.

where $\mathbf{A}=\mathbf{A}_{1}+2\mathbf{A}_{2}$ for $\mathbf{A}_{1}$ as defined in Corollary D.5 and $\mathbf{A}_{2}[n,k]=p_{n}(-1)p_{k}(-1)$.

To approximate $u(t-\theta)$, we note that when ${Y}=t-\theta$, $z=-1$. Then

We now show that the special case of Corollary D.6 for Legendre matches the results from .

Let $p_{n}(z)=\left({\frac{2n+1}{2}}\right)^{1/2}P_{n}(z)$ where $P_{n}(z)$ are the Legendre polynomials. Then

and $\mathbf{b}[n]=\left({\frac{2n+1}{2}}\right)^{\frac{1}{2}}$.

where $\mathbf{A}=\mathbf{A}_{1}+2\mathbf{A}_{2}$ for $\mathbf{A}_{1}$ as defined in Corollary D.5 and $\mathbf{A}_{2}[n,k]=p_{n}(-1)p_{k}(-1)$.

Then $\mathbf{b}[n]=\left({\frac{2n+1}{2}}\right)^{\frac{1}{2}}$ follows from Corollary D.6 and (29).

From the following recurrence relations [1, Chapter 12]:

Recalling that $\mathbf{A}_{1}[n,k]=2\alpha_{n,k}$.

We note that from (28), $\mathbf{A}_{2}[n,k]=\left({\frac{2n+1}{2}}\right)^{\frac{1}{2}}\left({\frac{2k+1}{2}}\right)^{\frac{1}{2}}(-1)^{n}(-1)^{k}=\frac{\left({2n+1}\right)^{\frac{1}{2}}\left({2k+1}\right)^{\frac{1}{2}}}{2}(-1)^{n-k}$.

Recalling $\mathbf{A}=\mathbf{A}_{1}+2\mathbf{A}_{2}$, we get:

D.2.4 Scaled HiPPO: Recovering HiPPO-LegS

We now use Theorem 7 to recover the HiPPO-LegS instantiation for the “Scaled Legendre” measure, the main method from Gu et al. .

Let $p_{n}(z)=\left({\frac{2n+1}{2}}\right)^{1/2}P_{n}(z)$ where $P_{n}(z)$ are the Legendre polynomials and let $\delta(t)=t$ for all t. Then

and $\mathbf{b}[n]=\left({\frac{2n+1}{2}}\right)^{\frac{1}{2}}$.

Let $\theta(t)=t$. By Theorem 7 and noting that $\theta(t)=1$, we get:

for $\alpha_{n,k}$,$\beta_{n,k}$ as defined in (25).

Using the same arguments as in the proof of Corollary D.7, $\mathbf{b}[n]=\left({\frac{2n+1}{2}}\right)^{\frac{1}{2}}$ follows from Corollary D.6 and (29). Also using similar arguments as the proof of Corollary D.7, we have

Including the normalization constant $(2n+1)^{\frac{1}{2}}$, we note that $(z-1)p^{\prime}_{n}(z)=(z+1)p^{\prime}_{n}(z)-2p^{\prime}_{n}(z)$. Then we get

Recalling that the definition for $\mathbf{A}_{1}$ from (30), we get:

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

This section proves Corollary 4.1, showing that the HiPPO matrices for measures corresponding to classical families of orthogonal polynomials are quasiseparable. We define quasi-separability in Section D.3.1. Theorem 8 proves the claimed result for Jacobi polynomials and Lemma D.11 proves the claimed result for Laguerre polynomials.

We note that there is a third family of classical OPs, the Hermite polynomials , which have a two-sided infinite measure. However, since HiPPO is about continuous-time memorization of a function’s history, it requires a one-sided measure and therefore the Hermite polynomials are not appropriate.

Every matrix contained strictly above the diagonal has rank at most $p$.

Every matrix contained strictly below the diagonal has rank at most $q$.

A $(q,q)$-quasiseparable matrix is called $q$-quasiseparable.

We are interested in showing the $\mathbf{A}$ matrices for a broad class of OPs in Corollary D.6 are $O(1)$-quasiseperable. We now state some properties of $q$-quasiseparable matrices:

Let $\mathbf{Q}$ be $q$-quasiseparable. Then:

Any submatrix contained strictly below or above the diagonal in $\mathbf{Q}$ has rank $\leq q$ and its corresponding submatrix in $\mathbf{Q}^{\prime}$ also has rank $\leq q^{\prime}$. This implies that the corresponding submatrix in $\mathbf{Q}\pm\mathbf{Q}^{\prime}$ has rank $\leq q+q^{\prime}$. Therefore $\mathbf{Q}\pm\mathbf{Q}^{\prime}$ is $(q+q^{\prime})$-quasiseparable.

Let the $r=\text{rank}(\mathbf{E})$. Thus any submatrix in $\mathbf{E}$ has rank $\leq r$. Then $\mathbf{E}$ is $r$-quasiseparable.

Multiplication by diagonal matrices only scales the rows and columns, leaving the rank of each submatrix unchanged.

D.3.2 Jacobi Polynomials

The Jacobi polynomial of degree $n$ with parameters $\alpha,\beta>-1$ will be denoted ${J_{n}^{\alpha,\beta}}\left({z}\right)$. The Jacobi polynomials are orthogonal with respect to measure $\omega(z)=(1-z)^{\alpha}(1+z)^{\beta}$. In particular, it is known from (eq. (4.3.3) from ) that

where $\Gamma(\cdot)$ is the gamma function. Let

be our normalization constant. We note that the normalized Jacobi polynomials

We now discuss some useful properties of Jacobi polynomials. It is known that (, eq. (3.100)):

From (4.21.7) in , it is known that the derivative of ${J_{n}^{\alpha,\beta}}\left({z}\right)$ is proportional to ${J_{n-1}^{\alpha+1,\beta+1}}\left({z}\right)$:

Additionally, the Jacobi polynomials ${J_{n-1}^{\alpha+1,\beta}}\left({z}\right)$ and ${J_{n-1}^{\alpha,\beta+1}}\left({z}\right)$ can be written as sums of ${J_{n-1}^{\alpha,\beta}}\left({z}\right)$ polynomials. In particular from (3.112) and (3.115),

Using (35) and (36) in (34) allows us to write $\frac{\partial}{\partial z}{J_{n}^{\alpha,\beta}}\left({z}\right)$ as a sum of $\left\{{J_{k}^{\alpha,\beta}}\left({z}\right)\right\}_{k\leq n}$ as follows:

We use these properties to write $\frac{\partial}{\partial z}{p_{n}^{\alpha,\beta}}\left({z}\right)$ as a sum of $\left\{{p_{k}^{\alpha,\beta}}\left({z}\right)\right\}_{k\leq n}$:

Let ${p_{n}^{\alpha,\beta}}\left({z}\right)$ and $\lambda^{\alpha,\beta}_{n}$ be as defined in (31).

Recall that ${J_{n}^{\alpha,\beta}}\left({z}\right)=\lambda^{\alpha,\beta}_{n}{p_{n}^{\alpha,\beta}}$. Then the claim follows from $(\ref{eq:jnab-sum})$.

D.3.3 HiPPO for Jacobi Polynomials

Let ${p_{n}^{\alpha,\beta}}\left({z}\right)$ be defined as in (31) and $\omega(z)=(1-z)^{\alpha}(1+z)^{\beta}$. Then

where $\mathbf{A}$ is $3$-quasiseperable.

where $\mathbf{A}=\mathbf{A}_{1}+2\mathbf{A}_{2}$ for $\mathbf{A}_{1}$ as defined in Corollary D.5 and $\mathbf{A}_{2}[n,k]={p_{n}^{\alpha,\beta}}(-1){p_{n}^{\alpha,\beta}}(-1)$.

where $\mathbf{D}_{11},\mathbf{D}_{12},\mathbf{D}_{21},\mathbf{D}_{22}$ are the diagonal matrices such that

(39) makes use of the fact that $(-1)^{n+k}=(-1)^{n-k}$ along with the definitions above.

Any submatrix of $\mathbf{Q}_{1}$ below the diagonal contains all $1$s, and submatrix of $\mathbf{Q}_{1}$ above the diagonal contains all s. Then any submatrix above or below the diagonal has rank 1. Therefore $\mathbf{Q}_{1}$ is 1-quasiseparable. Since $\mathbf{Q}_{1}$ is 1-quasiseparable and $\mathbf{D}_{11},\mathbf{D}_{12},\mathbf{D}_{21},\mathbf{D}_{22}$ are all diagonal matrices, part (iii) of Lemma D.9 implies that the matrices $\mathbf{D}_{11}\mathbf{Q}_{1}\mathbf{D}_{12}$ and $\mathbf{D}_{21}\mathbf{Q}_{1}\mathbf{D}_{22}$ are both 1-quasiseparable. Therefore part (i) of Lemma D.9 implies that $\mathbf{A}_{1}$ is 2-quasiseparable.

From (4.1.1) and (4.1.4) in , it is known that

Then $\mathbf{A}_{2}$ can be written $\mathbf{D}_{3}\mathbf{Q}_{2}\mathbf{D}_{4}$ where $\mathbf{D}_{3},\mathbf{D}_{4}$ are the diagonal matrices such that

where $\mathbf{Q}_{2}[n,k]=1$ for all $0\leq n,k<N$. $\mathbf{Q}_{2}$ has rank 1, and $\mathbf{D}_{3},\mathbf{D}_{4}$ are diagonal matrices. Hence by part (ii) and (iii) Lemma D.9, $\mathbf{A}_{2}$ is 1-quasiseparable.

Since $\mathbf{A}_{1}$ is 2-quasiseparable and $\mathbf{A}_{2}$ is 1-quasiseparable, part (i) of Lemma D.9 implies that $\mathbf{A}=\mathbf{A}_{1}+2\mathbf{A}_{2}$ is 3-quasiseparable and the claim follows. ∎

D.3.4 HiPPO-LagT

The Laguerre polynomial of degree $n$ with parameters $\alpha>-1$ will be denoted $L^{\alpha}_{n}\left({z}\right)$. The Laguerre polynomials are orthogonal with respect to measure $z^{\alpha}e^{-z}$. In particular, from (5.1.1) in we know that

Let $\lambda_{n}=\left({\frac{\Gamma(n+1)}{\Gamma(n+\alpha+1)}}\right)^{\frac{1}{2}}$ be our normalization constant. We note that the normalized Laguerre polynomials

Let $p_{n}(z)$ be defined as in $(\ref{eq:lag})$. Then

We now show that $\mathbf{A}$ as defined in Theorem 9 is $1$-quasiseperable.

Let $\mathbf{A}$ be defined as in Theorem 9. Then $\mathbf{A}$ is $1$-quasiseperable.

Below the diagonal, all entries $\mathbf{A}[n,k]=1$. Then any submatrix below the diagonal has rank 1. Similarly, above the diagonal, all entries $\mathbf{A}[n,k]=0$. Then any submatrix above the diagonal also has rank 1. Then by Definition 4, the claim follows.

Appendix E LSSL Algorithms

Section E.1 proves Theorem 2, providing an algorithm to compute the Krylov function efficiently for LSSLs.

Section E.2 shows a further simplification of Corollary 4.1, presenting an even simpler class of structured matrices that we use in our implementation of LSSL.

Section E.3 provides technical details of the implementation of LSSL, in particular for computing the MVM black box (multiplication by $\overline{A}$) and for computing gradients during backpropagation.

This section addresses the computational aspects of the LSSL. In particular, we prove Theorem 2 for the computational speed of computing the Krylov function (7) for quasiseparable matrices $A$, by providing a concrete algorithm in Section E.1.1.

We call (7) the Krylov function following the notation of , since it can be written $\mathcal{K}(A,B)^{T}C$ where $\mathcal{K}(A,B)$ is the Krylov matrix defined in (8). Alternative naming suggestions are welcome.

We follow the similar problem of [17, Lemma 6.6] but track the dependence on $L$ and the log factors more precisely, and optimize it in the case of stronger structure than quasiseparability, which holds in our setting (particularly Theorem 11).

The first step is to observe that the Krylov function $\mathcal{K}_{L}(A,B,C)$ is actually the coefficient vector of $C(I-Ax)^{-1}B\pmod{x^{L}}$ as a polynomial in $x$. (Note that $Ax$ means simply multiplying every entry in $A$ by a scalar variable $x$.) This follows from expanding the power series $(I-Ax)^{-1}=I+Ax+A^{2}x^{2}+\dots$. Thus we first compute $C(I-Ax)^{-1}B$, which is a rational function of degree at most $N$ in the numerator and denominator (which can be seen by the standard adjoint formula for the matrix inverse).

The second step is simply inverting the denominator of this rational function $\pmod{x^{L}}$ and multiplying by the numerator, both of which are operations that need $L\log(L)$ time by standard results for polynomial arithmetic .

For the remainder of this section, we focus on computing the first part. We make two notational changes: First, we transpose $\mathbf{C}$ to make it have the same shape as $\mathbf{B}$. We consider the more general setting where $\mathbf{B}$ and $\mathbf{C}$ have multiple columns; this can be viewed as handling a “batch” problem with several queries for $\mathbf{B},\mathbf{C}$ at the same time.

Let $\mathbf{A}$ be a $q$-quasiseparable matrix. Then

is a $k\times k$ matrix of rational functions of degree at most $N$, which can be computed in $O(q^{3}\log^{4}N)$ operations.

The main idea is that quasiseparable matrices are recursively “self-similar”, in that the principal submatrices are also quasiseparable, which leads to a divide-and-conquer algorithm. In particular, divide $\mathbf{A}=\begin{bmatrix}\mathbf{A}_{00}&\mathbf{A}_{01}\\ \mathbf{A}_{10}&\mathbf{A}_{11}\end{bmatrix}$ into quadrants. Then by Definition 4, $\mathbf{A}_{00},\mathbf{A}_{11}$ are both $q$-quasiseparable and $\mathbf{A}_{01},\mathbf{A}_{10}$ are rank $q$. Therefore the strategy is to view $\mathbf{I}-\mathbf{A}x$ as a low-rank perturbation of smaller quasiseparable matrices and reduce the problem to a simpler one.

Over a commutative ring $\mathcal{R}$, let $\mathbf{A}\in\mathcal{R}^{N\times N}$ and $\mathbf{U},\mathbf{V}\in\mathcal{R}^{N\times p}$. Suppose $\mathbf{A}$ and $\mathbf{A}+\mathbf{U}\mathbf{V}^{T}$ are invertible. Then $\mathbf{I}_{p}+\mathbf{V}^{T}\mathbf{A}^{-1}\mathbf{U}\in\mathcal{R}^{p\times p}$ is invertible and

Suppose we know the expansions of each of

Then the final result can be computed by inverting $\mathbf{I}_{2t}+\mathbf{M}_{3}$ ($O(q^{3}N\log(N))$ operations), multiplying by $\mathbf{M}_{2},\mathbf{M}_{4}$ ($O((kq^{2}+k^{2}q)N\log(N))$ operations), and subtracting from $\mathbf{M}_{1}$ ($O(k^{2}N\log(N))$ operations). This is a total of $O((q^{3}+kq^{2}+k^{2}q)N\log(N))$ operations. Note that when $k=O(q\log N)$, this becomes $O(q^{3}N\log^{3}{N})$; we will use this in the analysis shortly.

To compute $\mathbf{M}_{1},\mathbf{M}_{2},\mathbf{M}_{3},\mathbf{M}_{4}$, it suffices to compute the following:

But to compute those, it suffices to compute the following $(k+t)\times(k+t)$ matrices:

Since $\mathbf{A}_{00}$ and $\mathbf{A}_{11}$ have the same form as $\mathbf{A}$, this is two recursive calls of half the size. Notice that the size of the other input (dimensions of $\mathbf{B},\mathbf{C}$) is growing, but when the initial input is $k=1$, it never exceeds $1+q\log{N}$ (since they increase by $q$ every time we go down a level). Earlier, we noticed that when $k=O(q\log N)$, the reduction step has complexity $O(q^{3}N\log^{3}(N))$ for any recursive call. The recursion adds an additional $\log{N}$ multiplicative factor on top of this. ∎

Suppose that $\mathbf{A}$ is semiseparable instead of quasiseparable, and suppose $q$ is a small constant. Then the cost of Lemma E.2 is $O(N\log^{2}(N))$ operations.

This follows from the fact that in the recursion (45) and (46), the $\mathbf{U},\mathbf{V}$ matrices do not have to be appended if they already exist in $\mathbf{B},\mathbf{C}$. For intuition, this happens in the case when $\mathbf{A}$ is tridiagonal, so that $U,V$ have the structure $(1,0,\dots,0)$, or the case when the off-diagonal part of $\mathbf{A}$ is all $1$ (such as the HiPPO-LegT matrix). The matrices in Section D.3 and Section E.2 (Theorems 8, 9 and 11) actually satisfy this stronger structure, so Corollary E.3 applies.

Combining everything, this proves Theorem 2 with the exact bound $N\log^{2}(N)+L\log(L)$ operations. The memory claim follows similarly, and the depth of the algorithm is $\log^{2}(N)+\log(L)$ from the divide-and-conquer recursions.

E.1.2 Summary of Computation Speed for LSSLs and other Mechanisms

We provide a summary of complexity requirements for various sequence model mechanisms, including several versions of the LSSL. Note that these are over exact arithmetic as in Theorem 2.

First, the self-attention mechanism is another common sequence model that has an $L^{2}$ dependence on the length of the sequence, so it is not suitable for the very long sequences we consider here. (We do note that there is an active line of work on reducing this complexity.)

Second, we include additional variants of the LSSL. In Table 5, LSSL-naive denotes learning $A$ and $\Delta t$ for unstructured $A$; LSSL-fixed denotes not learning $A,\Delta t$ (see Appendix B for details); LSSL denotes the learning $A$ and $\Delta t$ for the structured class $A$.

We include brief explanations of these complexities for the LSSL variants.

Parameters: $O(HN)$ in the matrices $B,C$ and $O(N^{2})$ in the matrix $A$.

Training: $O(HN^{3})$ to invert compute the matrix $\overline{A}$ for all $H$ features. $O(LHN^{2})$ to compute the Krylov matrix $C,CA,\dots$. $O(BL\log(L)HN$ to multiply by $B$ and convolve with $u$.

Memory: $O(HN^{2})$ to store $\overline{A}$. $O(LHN)$ to store the Krylov matrix. $O(BLH)$ to store the inputs/outputs

Inference: $O(HN^{2})$ to for MVM by $\overline{A}$.

Training: $O(BL\log(L)H)$ to convolve with $u$.

Memory: $O(LHN)$ to store the Krylov matrix (but cached, so no backprop). $O(BLH)$ for inputs/outputs.

Inference: $O(HN^{2})$ to for MVM by $\overline{A}$.

Parameters: $O(HN)$ for $A,B,C,\Delta t$.

Memory: $O(BHL)$ to store Krylov (and inputs/outputs).

Inference: $O(HN)$ to multiply $x_{t}[H,N]$ by $\overline{A}[H,N,N]$

E.2 Further Simplification with Tridiagonal Matrices

The algorithm for Theorem 2 for general quasiseparable matrices is still difficult to implement in practice, and we make a further simplification using a particular subclass of quasiseparable matrices.

The class of $N\times N$ matrices $\mathcal{S}_{N}=\{P(D+T^{-1})Q\}$ with diagonal $D,P,Q$ and tridiagonal $T$ includes the original HiPPO-LegS, HiPPO-LegT, and HiPPO-LagT matrices .

Theorem 11 shows that a simple representation involving tridiagonal and diagonal matrices captures all of the original HiPPO matrices. In particular, our LSSL implementation initializes $A$ to be the HiPPO-LegS matrix (Appendix B) and learns within the class defined by Theorem 11.

We note that the matrices in Theorem 11 are all 1-quasiseparable and in particular also contain the HiPPO matrices for Gegenbauer and generalized Laguerre orthogonal polynomials derived in Theorem 9. In fact, the notion of semiseparability, which is closely related to (and actually is the predecessor of) quasiseparability, was originally motivated precisely to capture inverses of tridiagonal matrices. Thus the structured class in Theorem 11 can be viewed as an approximation of $3$-quasiseparable matrices (Corollary 4.1) to $1$-quasiseparable, which still contains many of the HiPPO families of interest.

We simply show that each of these specific matrices can be represented in the proposed form.

Let $A$ denote the HiPPO-LegT transition matrix. Up to row/column scaling (i.e. left- and right- multiplication by diagonal $P$ and $Q$), we can write

This can be written as $-PA^{\prime}Q$ where $P=Q=\operatorname*{diag}((2n+1)^{\frac{1}{2}})$ and

Finally, $A^{\prime}=D+T^{-1}$ where $D=-\operatorname*{diag}(\frac{n}{2n+1})$ and $T$ is the matrix with $1$ on the main diagonal and $-1$ on the subdiagonal.

This can be written as $-P(D+T^{-1})Q$ where $P=Q=I$, $D=-\frac{1}{2}I$, and $T$ is the same tridiagonal matrix as in the HiPPO-LegS case. ∎

E.3 Implementation Details

In this section we provide several implementation details that are useful for implementing LSSLs in practice.

Recall that one of the main primitives of LSSLs is the matrix-vector multiplication $y=\overline{A}x$ (Section 3.1, Appendix B), where $\overline{A}$ is the state matrix $A$ discretized with step size $\Delta t$ using the bilinear method (Section C.1.3). In Section E.3.1, we describe how this MVM can be performed with simpler MVM primitives which we call the “forward difference” and “backward difference”.

However, if these MVM primitives are implemented in a specialized way for particular classes of $A$ matrices (i.e., not using atoms in a standard autograd framework), then we also need to calculate several additional gradients by hand. Section E.3.2 shows that calculating gradients to $A,\Delta t,x$ during backpropagation can actually be reduced to those same forward/backward difference primitives.

Finally, in the case when $A$ is the structured class of matrices in Theorem 11, Section E.2 shows how to efficiently calculate those primitives using a black-box tridiagonal solver. Our codeAvailable at https://github.com/HazyResearch/state-spaces implements all the algorithms in this section, with bindings to the cuSPARSE library for efficient tridiagonal solving on GPU.

The discrete state-space system is given by (4) and (5), re-written here for convenience

where $\overline{A}$ is a function of $A,\delta_{t}$ and $\overline{A}$ is a function of $A,B,\delta_{t}$. In particular, we define $\overline{A}$ to be the matrix discretized using the bilinear method (Section C.1.3), and the system can be written explicitly:

We will call these functions the forward difference and backward difference maps, respectively. (The Euler and backward Euler discretizations (Section C.1.2) are also known as the “forward difference” and “backward difference” methods, which in the case of linear systems reduces down to the maps $F$ and $B$.)

E.3.2 Gradients through the Forward/Backward Difference Primitives

In this section we will let $y=F(A,\Delta t,x)$ or $y=B(A,\Delta t,x)$ denote the computation of interest, $L(y)$ denote a generic loss function, and $dx,dy,\dots$ denote gradients to $x,y,\dots$ (e.g., $dx=\frac{\partial L(y)}{\partial x}$).