How to Train Your HiPPO: State Space Models with Generalized Orthogonal Basis Projections

Introduction

The Structured State Space model (S4) is a recent deep learning model based on continuous-time dynamical systems that has shown promise on a wide variety of sequence modeling tasks . It is defined as a linear time-invariant (LTI) state space model (SSM), which give it multiple properties : as an SSM, S4 can be simulated as a discrete-time recurrence for efficiency in online or autoregressive settings, and as a LTI model, S4 can be converted into a convolution for parallelizability and computational efficiency at training time. These properties give S4 remarkable computational efficiency and performance, especially when modeling continuous signal data and long sequences.

Despite its potential, several aspects of the S4 model remain poorly understood. Most notably, Gu et al. claim that the long range effects of S4 arise from instantiating it with a particular matrix they call the HiPPO matrix. However, this matrix was actually derived in prior work for a particular time-varying system , and the use of this matrix in a time-invariant SSM did not have a mathematical interpretation. Consequently, the mechanism by which S4 truly models long-range dependencies is actually not known. Beyond this initialization, several other aspects of parameterizing and training S4 remain poorly understood. For example, S4 involves an important timescale parameter $\Delta$, and suggests a method for parameterizing and initializing this parameter, but does not discuss its meaning or provide a justification.

This work aims to provide a comprehensive theoretical exposition of several aspects of S4. The major contribution of this work is a cleaner, more intuitive, and much more general formulation of the HiPPO framework. This result directly generalizes all previous known results in this line of work . As immediate consequences of this framework:

We prove a theoretical interpretation of S4’s state matrix $\bm{A}$, explaining S4’s ability to capture long-range dependencies via decomposing the input with respect to an infinitely long, exponentially-decaying measure (Fig. 1 (Left)).

We derive new HiPPO matrices and corresponding S4 variants that generalize other nice basis functions. For example, our new method S4-FouT produces truncated Fourier basis functions. This method thus automatically captures sliding Fourier transforms (e.g. the STFT and spectrograms) which are ubiquitous as a hand-crafted signal processing tool, and can also represent any local convolution, thus generalizing conventional CNNs (Fig. 1 (Middle)).

We provide an intuitive explanation of the timescale $\Delta$, which has a precise interpretation as controlling the length of dependencies that the model captures. Our framework makes it transparent how to initialize $\Delta$ for a given task, as well as how to initialize the other parameters (in particular, the last SSM parameter $\bm{C}$) to make a deep SSM variance-preserving and stable.

𝑨𝑥𝑩𝑢x^{\prime}=\bm{A}x+\bm{B}u that produce particular basis functions, leading to several generalizations and new methods. (Left: LegS) We prove that the particular $\bm{A}$ matrix chosen in S4 produces Legendre polynomials under an exponential re-scaling, resulting in smooth basis functions with a closed form formula. This results in a simple mathematical interpretation of the method as orthogonalizing against an exponentially-decaying measure, granting the system better ability to model long-range dependencies. (Middle, Right: FouT) We derive a new SSM that produces approximations to the truncated Fourier basis, perhaps the most intuitive and ubiquitous set of basis functions. This method generalizes sliding Fourier Transforms and local convolutions (i.e. CNNs), and can also encode spike functions to solve classic memorization tasks. Empirically, we validate our theory on synthetic function reconstruction and memorization tasks, showing that empirical performance of state space models in several settings is predicted by the theory. For example, our new S4-FouT method, which can provably encode a spike function as its convolution kernel, performs best on a continuous memorization task compared to other SSMs and other models, when $\Delta$ is initialized correctly. Finally, we show that the original S4 method is still best on very long range dependencies, achieving a new state of the art of 86% average on Long Range Arena, with 96% on the most difficult Path-X task that even the other S4 variants struggle with.

Background

The state space model (SSM) is defined by the simple differential equation (1) and (2). It maps a 1-D input signal $u(t)$ to an $N$-D latent state $x(t)$ before projecting to a 1-D output signal $y(t)$.

For the remainder of this paper, we will assume $\bm{D}=0$ and omit it for simplicity, unless explicitly mentioned.

SSMs can in general have dynamics that change over time, i.e. the matrices $\bm{A},\bm{B},\bm{C},\bm{D}$ are a function of $t$ in (1) and (2). However, when they are constant the system is linear time invariant (LTI), and is equivalent to a convolutional system (3). The function $K(t)$ is called the impulse response which can also be defined as the output of the system when the input $u(t)=\delta(t)$ is the impulse or Dirac delta function. We will call these time-invariant state space models (TSSM). These are particularly important because the equivalence to a convolution makes TSSMs parallelizable and very fast to compute, which is critical for S4’s efficiency.

Our treatment of SSMs will consider the $(\bm{A},\bm{B})$ parameters separately from $\bm{C}$. We will refer to an SSM as either the tuple $(\bm{A},\bm{B},\bm{C})$ (referring to (3)) or $(\bm{A},\bm{B})$ (referring to Definition 1) when the context is unambiguous. We also drop the T in TSSM when the context is clearly time-invariant.

Given a TSSM $(\bm{A},\bm{B})$, $e^{t\bm{A}}\bm{B}$ is a vector of $N$ functions which we call the SSM basis. The individual basis functions are denoted $K_{n}(t)=\bm{e}_{n}^{\top}e^{t\bm{A}}\bm{B}$, which satisfy $x_{n}(t)=(u\ast K_{n})(t)=\int_{-\infty}^{t}K_{n}(t-s)u(s)\mathop{}\!ds$. Here $\bm{e}_{n}$ is the one-hot basis vector.

This definition is motivated by noting that the SSM convolutional kernel is a linear combination of the SSM basis controlled by the vector of coefficients $\bm{C}$, $K(t)=\sum_{n=0}^{N-1}\bm{C}_{n}K_{n}(t)$.

To be applied on a discrete input sequence $(u_{0},u_{1},\dots)$ instead of continuous function $u(t)$, (1) must be discretized by a step size $\Delta$ that represents the resolution of the input. Conceptually, the inputs $u_{k}$ can be viewed as sampling an implicit underlying continuous signal $u(t)$, where $u_{k}=u(k\Delta)$. Analogous to the fact that the SSM has equivalent forms either as an dynamical system (1) or a continuous convolution (3), the discrete-time SSM can be computed either as a recurrence or a discrete convolution. The mechanics to compute the discrete-time SSM has been discussed in previous works . For our purposes, we only require the following fact: for standard discretization methods used in prior work, discretizing the state space $(\bm{A},\bm{B})$ at a step size $\Delta$ is exactly equivalent to discretizing the state space $(\Delta\bm{A},\Delta\bm{B})$ at a step size $1$. This allows thinking of $\Delta$ simply as modulating the SSM parameters $(\bm{A},\bm{B})$ instead of representing a step size.

A poorly understood question from prior work is how to interpret and choose this $\Delta$ parameter, especially when the input $u_{k}$ does not actually arise from uniformly sampling an underlying continuous signal. S4 specifies to log-uniformly initialize $\Delta$ in the range $(\Delta_{min},\Delta_{max})=(0.001,0.1)$, but do not provide a concrete justification. In Section 3.3 we show a simpler interpretation of $\Delta$ directly in terms of the length of dependencies in a discrete input sequence.

2 HiPPO: High-order Polynomial Projection Operators

S4 is defined as a TSSM where $(\bm{A},\bm{B})$ is initialized with a particular formula (4). This was called the HiPPO matrix in , but is actually just one of several such special matrices derived in . To disambiguate other variants of S4, we refer to the full S4 method using this HiPPO SSM as S4-LegS. Other cases considered in this work include LegT from prior work (5) and FouT that we introduce (6).

These matrices were originally motivated by the question of ‘online memorization’ of an input signal. The key idea is that for a suitably chosen SSM basis $\bm{A},\bm{B}$, then at any time $t$, the current state $x(t)$ can be used to approximately reconstruct the entire input $u$ up to time $t$ (Fig. 2).

The main theoretical idea is as follows. Suppose that the basis functions satisfy Definition 2.

We call an SSM $(\bm{A}(t),\bm{B}(t))$ an orthogonal SSM (OSSM) for the basis $p_{n}(t,s)$ and measure $\omega(t,s)\geq 0$ if the functions $K_{n}(t,s)=p_{n}(t,s)\omega(t,s)$ satisfy, at all times $t$,

In the case of a time-invariant OSSM (TOSSM), $K_{n}(t,s)=:K_{n}(t-s)$ (depends only on $t-s$) giving us Definition 1 with measure $\omega(t-s):=\omega(t,s)$ and basis $p_{n}(t-s):=p_{n}(t,s)$.

To be more specific about terminology, $p_{n}(t)$ and $\omega_{n}(t)$ are called the basis and measure for orthogonal SSMs (Definition 2), while $K_{n}(t)$ are called the SSM basis kernels which applies more generally to all SSMs (Definition 1). The distinction will be made clear from context, notation, and the word “kernel” referring to $K_{n}(t)$.

For OSSMs, $(p,\omega)$ and $K$ are uniquely determined by each other, so we can refer to an OSSM by either. One direction is obvious: $(p,\omega)$ determine $K$ via $K_{n}(t,s)=p_{n}(t,s)\omega(t,s)$.

If a set of kernel functions satisfies $K_{n}(t,s)=p_{n}(t,s)\omega(t,s)$ where the functions $p_{n}$ are complete and orthogonal w.r.t. $\omega$ (equation (7) right), $p$ and $\omega$ are unique.

Equation (7) is equivalent to saying that for every fixed $t$, $\langle p_{n},p_{m}\rangle_{\omega}=\delta_{n,m}$, or that $p_{n}$ are an orthonormal basis with respect to measure $\omega$. More formally, defining $p_{n}^{(t)}(s)=p_{n}(t,s)$ and $\omega^{(t)}$ similarly, then $p_{n}^{(t)}$ are orthonormal in the Hilbert space with inner product $\langle p,q\rangle=\int p(s)q(s)\omega^{(t)}(s)\mathop{}\!ds$). By equation (7), $x_{n}(t)=\int_{-\infty}^{t}u(s)K_{n}(t,s)\mathop{}\!ds=\langle u,p_{n}^{(t)}\rangle_{\omega^{(t)}}$ where $p_{n}^{(t)}(s)=p_{n}(t,s)$. Thus at all times $t$, the state vector $x(t)$ is simply the projections of $u\mid_{\leq t}$ onto a orthonormal basis, so that the history of $u$ can be reconstructed from $x(t)$. HiPPO called this the online function approximation problem .

Consider an OSSM that satisfies (7) and fix a time $t$. Furthermore suppose that in the limit $N\to\infty$, the $p_{n}^{(t)}$ are a complete basis on the support of $\omega$. Then $u(s)=\sum_{n=0}^{\infty}x_{n}(t)p_{n}(t,s)$ for all $s\leq t$.

1𝑡𝑨𝑥1𝑡𝑩𝑢x^{\prime}=\frac{1}{t}\bm{A}x+\frac{1}{t}\bm{B}u, $x(t)$ represents the projection of the entire history of $u$ onto the Legendre polynomials. It was previously unknown how to interpret the time-invariant version of this ODE using the same $(\bm{A},\bm{B})$ matrices. The main barrier to using Proposition 2 for function reconstruction is that SSMs are in general not OSSMs. For example, even though we will show that (4) is an TOSSM, its diagonalization is not.

There is no TOSSM with the diagonal state matrix $\bm{A}=\operatorname*{diag}\{-1,-2,\dots\}$.

HiPPO can be viewed as a framework for deriving specific SSMs that do satisfy (7). The original HiPPO methods and its generalizations primarily focused on the case when the $p_{n}$ are orthogonal polynomials, and specifically looked for solutions to (7), which turn out to be SSMs. We have rephrased the HiPPO definition in Definition 2 to start directly from SSMs.

We discuss the two most important cases previously introduced.

(5) is a TOSSM that approximates the truncated Legendre polynomials (Fig. 3).

As $N\to\infty$, the SSM with $(\bm{A},\bm{B})$ in (5) is a TOSSM with

This particular system was the precursor to HiPPO and has also been variously called the Legendre Delay Network (LDN) or Legendre Memory Unit (LMU) . The original motivation of this system was not through the online function approximation formulation of HiPPO, but through finding an optimal SSM approximation to the delay network that has impulse response $K(t)=\delta(t-1)$ representing a time-lagged output by 1 time unit (Fig. 3). We state and provide an alternate proof of this result in Section 3.2, Theorem 9.

Unlike the HiPPO-LegT case which is an LTI system (1) (i.e. TOSSM), the HiPPO-LegS matrix (4) was meant to be used in a time-varying system $x^{\prime}(t)=\frac{1}{t}\bm{A}x(t)+\frac{1}{t}\bm{B}u(t)$ . In contrast to HiPPO-LegT, which reconstructs onto the truncated Legendre polynomials in sliding windows $[t-1,t]$, HiPPO-LegS reconstructs onto Legendre polynomials on “scaled” windows $[0,t]$; since the window changes across time, the system is not time-invariant (Fig. 2).

The SSM $(\frac{1}{t}\bm{A},\frac{1}{t}\bm{B})$ for $(\bm{A},\bm{B})$ in (4) is an OSSM with

However, the S4 model applies the exact same formula (4) inside the time-invariant SSM (1), i.e. dropped the $\frac{1}{t}$ term, which had no mathematical interpretation. In other words, while $(\frac{1}{t}\bm{A},\frac{1}{t}\bm{B})$ is an OSSM, it was not known whether the TSSM $(\bm{A},\bm{B})$ is a TOSSM. Given that the performance of SSM models is very sensitive to these matries $\bm{A}$ , it remained a mystery why this works. In Section 3 we will prove that (4) actually does correspond to a TOSSM.

While HiPPO originally showed just the above two cases involving Legendre polynomials (and another case called LagT for Laguerre polynomials, which will not be a focus of this work), follow-up work showed that there exist OSSMs corresponding to all families of orthogonal polynomials $\{p_{n}(t)\}$. Our more general framework will also subsume these results.

We use HiPPO-[SSM] to refer to a fixed OSSM $(\bm{A},\bm{B})$ suitable for online function approximation, where [SSM] is a suffix (e.g. LegS, LegT) that abbreviates the corresponding basis functions (e.g. scaled Legendre, truncated Legendre). S4-[SSM] refers to the corresponding trainable layer $(\bm{A},\bm{B},\bm{C})$ with randomly initialized $\bm{C}$, trained with S4’s representation and computational algorithm .

Generalized HiPPO: General Orthogonal Basis Projections

In Section 3.1, we prove that the LTI HiPPO-LegS is actually a TOSSM and show closed formulas for its basis functions. In Section 3.2, we include more specific results on finite-window SSMs, including introducing a new method HiPPO-FouT based on truncated Fourier functions, and proving previously established conjectures. Section 3.3 shows more general properties of TOSSMs, which establish guidelines for interpreting and initializing SSM parameters such as the timescale $\Delta$.

Our main, fully general, result is Theorem 12 in Section C.2, which describes a very general way to derive OSSMs for various SSM basis functions $K_{n}(t,s)$. This result can be instantiated in many ways to generalize all previous results in this line of work.

We show the matrices $(\bm{A},\bm{B})$ in (4) are deeply related to the Legendre polynomials $L_{n}$ defined in Theorem 5.

Define $\sigma(t,s)=\exp(a(s)-a(t))$ for any differentiable function $a$. The SSM $(a^{\prime}(t)\bm{A},a^{\prime}(t)\bm{B})$ is an OSSM with

As more specific corollaries of Corollary 3.1, we recover both the original time-varying interpretation of the matrix in (4), as well as the instantiation of LegS as a time-invariant system. If we set $a^{\prime}(t)=\frac{1}{t}$, then we recover the scale-invariant HiPPO-LegS OSSM in Theorem 5,

And if $a^{\prime}(t)=1$, this shows a new result for the time-invariant HiPPO-LegS TOSSM:

The SSM $(\bm{A},\bm{B})$ is a TOSSM with

This explains why removing the $\frac{1}{t}$ factor from HiPPO-LegS still works: it is orthogonalizing onto the Legendre polynomials with an exponential “warping” or change of basis on the time axis (Fig. 1).

2 Finite Window Time-Invariant Orthogonal SSMs

For the remainder of this section, we restrict to the time-invariant SSM setting (3). A second important instantiation of Theorem 12 covers cases with a discontinuity in the SSM basis functions $K_{n}(t)$, which requires infinite-dimensional SSMs to represent. The most important type of discontinuity occurs when $K_{n}(t)$ is supported on a finite window, so that these TSSMs represent sliding window transforms.

We first derive a new sliding window transform based on the widely used Fourier basis (Section 3.2.1). We also prove results relating finite window methods to delay networks (Section 3.2.2)

Using the more general framework (Theorem 12) that does not necessarily require polynomials as basis functions, we derive a TOSSM that projects onto truncated Fourier functions.

This SSM corresponds to Fourier series decompositions, a ubiquitous tool in signal processing, but represented as a state space model. The basis is visualized in Fig. 1 (middle) for state size $N=1024$.

A benefit of using these well-behaved basis functions is that we can leverage classic results from Fourier analysis. For example, it is clear that taking linear combinations of the truncated Fourier basis can represent any function on $$, and thus S4-FouT can represent any local convolution (i.e. the layers of modern convolutional neural networks).

Let $K(t)$ be a differentiable kernel on $$, and let$\hat{K}(t)$be its representation by the FouT system (Theorem 6) with state size$N.$If$K$is$L-$Lipschitz, then for$\epsilon>0,N\geq\left(\frac{L}{\pi\epsilon}\right)^{2}+2$, we have$\lVert{K(t)-\hat{K}(t)}\rVert\leq\epsilon.$If$K$has$k-$derivatives bounded by$L$, then we can take$N\geq\left(\frac{L}{\pi^{k}\epsilon}\right)^{\frac{2}{2k-1}}+2.$

2.2 Approximating Delay Networks

An interesting property of these finite window TSSMs is that they can approximate delay functions. This is defined as a system with impulse response $K(t)=\delta(t-1)$: then $y(t)=(K\ast u)(t)=u(t-1)$, which means the SSM outputs a time-lagged version of the input. This capability is intuitively linked to HiPPO, since in order to do this, the system must be remembering the entire window $u([t-1,t])$ at all times $t$, in other words perform an approximate function reconstruction. Any HiPPO method involving finite windows should have this capability, in particular, the finite window methods LegT and FouT.

For the FouT system $\bm{A}$ and $\bm{B}$, let $\bm{C}$ be (twice) the vector of evaluations of the basis functions $\bm{C}_{n}=2\cdot p_{n}(1)$ and let $\bm{D}=1$. For the LegT system $\bm{A}$ and $\bm{B}$, let $\bm{C}$ be the vector of evaluations of the basis functions $\bm{C}_{n}=p_{n}(1)=(1+2n)^{\frac{1}{2}}(-1)^{n}$ and let $\bm{D}=0$.

Then the SSM kernel $K(t)=\bm{C}e^{t\bm{A}}\bm{B}+\bm{D}\delta(t)$ limits to $K(t)\to\delta(t-1)$ as $N\to\infty$.

Theorem 8 is visualized in Figs. 1 and 3 (right). Further, the result for LegT can be characterized even more tightly for finite $N$. In fact, this was the original motivation for the LDN/LMU , which worked backward from the transfer function of the desired delay function impulse response $K(t)=\delta(t-1)$, and noticed that the SSM for Padé approximations to this were linked to Legendre polynomials. This was not fully proven, and we state it here and provide a full proof in Section C.4.

For $\bm{A},\bm{B},\bm{C},\bm{D}$ in the LegT system described in Theorem 8, the transfer function $\mathcal{L}\{K(t)\}(s)$ is the $[N-1/N]$ Padé approximant to $e^{-s}=\mathcal{L}\{\delta(t-1)\}(s)$.

We remark that although LegT (LMU) is designed to be an “optimal” approximation to the delay function via Padé approximants, it actually produces a weaker spike function than FouT (Fig. 3 vs. Fig. 1) and empirically performs slightly worse on synthetic tasks testing this ability (Section 4.3). This may be because Padé approximation in the Laplace domain does not necessarily translate to localization in the time domain.

3 Properties of Time-invariant Orthogonal SSMs: Timescales and Normalization

We describe several general properties of TOSSMs, which let us answer the following questions:

How should all parameters $(\bm{A},\bm{B},\bm{C})$ be initialized for an SSM layer to be properly normalized?

What does $\Delta$ intuitively represent, and how should it be set in an SSM model?

It turns out that for TOSSMs, these two questions are closely related and have intuitive interpretations.

First, several basic transformations preserve the structure of TOSSMs. Consider a TOSSM $(\bm{A},\bm{B})$ with basis functions $p_{n}(t)$ and measure $\omega(t)$. Then, for any scalar $c$ and unitary matrix $\bm{V}$, the following are also TOSSMs with the corresponding basis functions and measure (Section C.5, Proposition 13):

A standard aspect of training deep learning models, in general, concerns the scale or variance of activations. This has been the subject of much research on training deep learning models, touching on deep learning theory for the dynamics of training such as the exploding/vanishing gradient problem , and a large number of normalization methods to ensure properly normalized methods, from the simple Xavier/He initializations to BatchNorm and LayerNorm to many modern variants and analyses of these .

The following proposition follows because for a TOSSM, $x(t)$ can be interpreted as projecting onto orthonormal functions in a Hilbert space (Proposition 2).

Consider an (infinite-dimensional) TOSSM. For any input $u(t)$, $\|x(t)\|_{2}^{2}=\|u\|_{\omega}^{2}=\int_{-\infty}^{t}u(s)^{2}\omega(t-s)\mathop{}\!dt$.

For a TOSSM with a probability measure (i.e. $\int\omega(t)=1$) and any constant input $u(t)=c$, the state has norm $\|x(t)\|^{2}=c^{2}$ and the output $y(t)$ has mean 0, variance $c^{2}$ if the entries of $\bm{C}$ are mean and variance $1$.

Note that the probability measure requirement can be satisfied by simply rescaling $\bm{B}$. Corollary 3.4 says the TOSSM preserves the variance of inputs, the critical condition for a properly normalized deep learning layer. Note that the initialization of $\bm{C}$ is different than a standard Linear layer in deep neural networks, which usually rescale by factor depending on its dimensionality such as $N^{-\frac{1}{2}}$ .

As discussed in Section 2, converting from continuous to discrete time involves a parameter $\Delta$ that represents the step size of the discretization. This is an unintuitive quantity when working directly with discrete data, especially if it is not sampled from an underlying continuous process.

We observe the following fact: for all standard discretization methods (e.g. Euler, backward Euler, generalized bilinear transform, zero-order hold ), the discretized system depends on $(\bm{A},\bm{B})$, and $\Delta$ only through their products $(\Delta\bm{A},\Delta\bm{B})$. This implies that the SSM $(\bm{A},\bm{B})$ discretized at step size $\Delta$ is computationally equivalent to the SSM $(\Delta\bm{A},\Delta\bm{B})$ discretized at step size $1$.

Therefore, $\Delta$ can be viewed just as a scalar scaling of the base SSM instead of changing the rate of the input. In the context of TOSSMs, this just scales the underlying basis and measure (Scalar Scaling). More broadly, scaling a general SSM simply changes its timescale or rate of evolution.

The ODE $y^{\prime}=c\bm{A}y+c\bm{B}u$ evolves at a rate $c$ times as fast as the SSM $x^{\prime}=\bm{A}x+\bm{B}u$, in the sense that the former maps $u(ct)\mapsto x(ct)$ if the latter maps $u(t)\mapsto x(t)$.

The most intuitive example of this is for a finite window TOSSM such as LegT or FouT. Discretizing this system with step size $\Delta$ is equivalent to considering the system $(\Delta\bm{A},\Delta\bm{B})$ with step size $1$, which produces basis functions supported exactly on $[0,\frac{1}{\Delta}]$. The interpretation of the timescale $\Delta$ lends to simple discrete-time corollaries of the previous continuous-time results. For example, LegT and FouT represent sliding windows of $1/\Delta$ elements in discrete time.

By Theorem 8, as $N\to\infty$, the discrete convolutional kernel $\bm{\overline{K}}\to\bm{e}_{\lceil\Delta^{-1}\rceil}$, i.e. the discrete delay network with lag $\frac{1}{\Delta}$.

For HiPPO-FouT matrices $(\bm{A},\bm{B})$, by Theorem 6, as $N\to\infty$, the discrete convolutional kernel $\bm{\overline{K}}$ (over the choice of $\bm{C}$) can represent any local convolution of length $\lfloor\Delta^{-1}\rfloor$.

This discussion motivates the following definition. Properly normalized TOSSMs $(\bm{A},\bm{B})$ will model dependencies of expected length $1$, and $\Delta$ modulates it to model dependencies of length $\frac{1}{\Delta}$, allowing fine-grained control of the context size of a TOSSM.

By this definition, HiPPO-LegS is timescale normalized. This motivates S4’s initialization of $\Delta$ log-uniformly in $(0.001,0.1)$, covering a geometric range of sensible timescales (expected length $10$ to $1000$). In Section 4 we show that the timescale can be chosen more precisely when lengths of dependencies are known.

4 Discussion

Table 1 summarizes the results for TOSSMs presented in this section, including both original HiPPO methods defined in Gu et al. as well as our new methods.

In summary, we do not recommend using the original HiPPO-LagT, which despite the original motivation does not represent orthogonalizing against an exponentially decaying measure. Instead, HiPPO-LegS (as a time-invariant SSM) actually represents an exponentially decaying measure.

For a timescale normalized orthogonal SSM (i.e. $\int_{0}^{\infty}\omega(t)=1$ and $\int_{0}^{\infty}t\omega(t)=1$):

$\frac{1}{\Delta}$ exactly represents the range of dependencies it captures. For example, S4-FouT can represent any finite convolution kernel of length $\frac{2}{\Delta}$ (so the expected length of a random kernel is $\frac{1}{\Delta}$).

A random vector $\bm{C}$ with independent mean , variance $1$ entries is a variance-preserving SSM, i.e. produces outputs matching the variance of the input.

The Linear State Space Layer succeeded HiPPO by incorporating it into a full deep SSM model, and also generalized the HiPPO theory to show that all orthogonal polynomials can be defined as the SSM kernels for some $(\bm{A},\bm{B})$. Our framework is even stronger and immediately produces the main result of LSSL as a corollary (Appendix), and can also work for non-polynomial methods (e.g. FouT).

These results show that all orthogonal polynomial bases, including truncated and scaled variants, have corresponding OSSMs with polynomial kernels. If we define this special case as polynomial OSSMs (POSSMs), we have therefore deduced that all of the original HiPPOs are POSSMs.

Experiments

We study the empirical tradeoffs of our proposed S4 variants. We compare several S4 variants based on the TOSSMs introduced in this work, as well as to simpler diagonal SSMs called S4D that are not orthogonal SSMs . Corresponding to our main contributions, we hypothesize that

S4-LegS excels at sparse memorization tasks because it represents very smooth convolution kernels that memorize the input against an infinitely-long measure (Corollary 3.3, Fig. 1). Conversely, it is less appropriate at short-range tasks with dense information because it smooths out the signal.

S4-FouT excels at dense memorization tasks because it can represent spike functions that pick out past elements at particular ranges (Section 3.2.2). However, it is less appropriate at very long range tasks because it represents a finite (local) window.

$\Delta$ can be initialized precisely based on known time dependencies in a given task to improve performance.

The Long Range Arena (LRA) benchmark is a suite of sequence classification tasks designed to stress test sequence models on modeling long sequences. We improve S4’s previous state of the art by another 6 points (Table 2). Validating our hypothesis, S4-LegS is extremely strong at the hardest long-range task (Path-X) involving sparse dependencies of length 16384, which FouT cannot solve because it is a finite window method.

The Path-X task also serves as a validation of the theory of timescales in Section 3.3. To set these results, we lowered the initialization of $\Delta$ in accordance with known length of dependencies in the task. Fig. 4 illustrates the importance of setting $\Delta$ correctly.

2 Theory: Function Reconstruction, Timescales, Normalization

Fig. 5 confirms the HiPPO theory of online function reconstruction (Proposition 2) for the proposed TOSSMs LegS and FouT.

Fig. 6 shows that S4-LegT and S4-FouT, the methods that theoretically reconstruct against a uniform measure, are far better than other methods. We include the new diagonal variants (S4D) proposed in , which are simpler SSM methods that generally perform well but do not learn the right function on this task. We also include a method S4-(LegS/FouT) which combines both LegS and FouT measures by simply initializing half of the SSM kernels to each. Despite having fewer S4-FouT kernels, this still performs as well as the pure S4-FouT initialization.

Finally, we validate the theory of normalization in Section 3.3, which predicts that for a properly normalized TOSSM, the projection parameter $\bm{C}$ should be initialized with unit variance, in contrast to standard initializations for deep neural networks which normalize by a factor related to the size of $N$ (in this case $N=64$). Table 3 shows classification results on datasets Sequential CIFAR (sCIFAR) and Speech Commands (SC), using models of size at most 150K parameters. This replicates the setup of the “ablation models” of [8, Section 5]. Results show that using standard deviation $1.0$ for $\bm{C}$ is slightly better than alternatives, although the difference is usually minor.

3 Memorization: the Delay (continuous copying) Task

Next, we study how the synthetic reconstruction ability transfers to other tasks. The Delay Task requires models to learn a sequence-to-sequence map whose output is the input lagged by a fixed time period (Fig. 7(a)). For recurrent models, this task can be interpreted as requiring models to maintain a memory buffer that continually remembers the latest elements it sees. This capability was the original motivation for the Legendre Memory Unit, the predecessor to HiPPO-LegT, which was explicitly designed to solve this task because it can encode a spike kernel (Fig. 3). In Fig. 7(b), we see that our new S4-FouT actually outperforms S4-LegT, which both outperform all other methods when the timescale $\Delta$ is set correctly. We note that this task with a lag of just 1000 time steps is too hard for baselines such as an LSTM and Transformer, which empirically did not learn better than random guessing (RMSE 0.43).

Summary: How to Train Your HiPPO

SSMs represent convolution kernels that are linear combinations (parameterized by $\bm{C}$) of basis functions (parameterized by $\bm{A}$ and $\bm{B}$).

HiPPO is a general mathematical framework for producing matrices $\bm{A}$ and $\bm{B}$ corresponding to prescribed families of well-behaved basis functions. We derive HiPPO matrices corresponding to exponentially-scaled Legendre families (LegS) and the truncated Fourier functions (FouT).

HiPPO-LegS corresponds to the original S4 method and produces a very smooth, long-range family of kernels (Fig. 1) that is still the best method for long-range dependencies among all S4 variants

HiPPO-FouT is a finite window method that subsumes local convolutions (e.g. generalizing vanilla CNNs, Corollary 3.6) and captures important transforms such as the sliding DFT or STFT

Independently of a notion of discretization, the timescale $\Delta$ has a simple interpretation as controlling the length of dependencies or “width” of the SSM kernels. Most intuitively, for a finite window method such as FouT, the kernels have length exactly $\frac{1}{\Delta}$, and generalize standard local convolutions used in deep learning.

A companion paper to this work builds on the theory introduced here to define a simpler version of S4 using diagonal state matrices (S4D), which are approximations to the orthogonal SSMs we introduce and can inherit S4’s strong modeling abilities . It also includes experiments on more datasets comparing various state space models, including the S4 variants (S4-LegS and S4-FouT) introduced here.

We gratefully acknowledge the support of DARPA under Nos. FA86501827865 (SDH) and FA86501827882 (ASED); 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); 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, Swiss Re, Brown Institute for Media Innovation, Department of Defense (DoD) through the National Defense Science and Engineering Graduate Fellowship (NDSEG) Program, Fannie and John Hertz Foundation, National Science Foundation Graduate Research Fellowship Program, Texas Instruments, and members of the Stanford DAWN project: Teradata, Facebook, Google, Ant Financial, NEC, VMWare, and Infosys. Atri Rudra and Isys Johnson are supported by NSF grant CCF-1763481. 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 DARPA, NIH, ONR, or the U.S. Government.

References

Appendix A Related Work

We discuss in more detail the differences between this work and the previous results in this line of work.

The HiPPO-LegT matrix (5) was first introduced as the LMU . The original motivation was to produce a state space model that approximates the Delay Network, which can be defined as the LTI system that transforms $u(t)$ into $u(t-1)$, i.e. lags the input by 1 time unit. This can also be defined as the system with impulse response $K(t)=\delta(t-1)$, i.e. convolves by the convolutional kernel with a $\delta$ spike at time $1$.

The connection between the Delay Network and Legendre polynomials was made in two steps. First, the transfer function of the ideal system is $\mathcal{L}[\delta(t-1)](s)=e^{-s}$ and must be approximated by a proper rational function to be represented as an SSM. Taking Padé approximants of this function yields “optimal” approximations by rational functions, which can then be distilled into a SSM $(\bm{A},\bm{B},\bm{C})$ whose transfer function $\bm{C}(s\bm{I}-\bm{A})^{-1}\bm{B}$ matches it. Second, the SSM basis $e^{t\bm{A}}\bm{B}$ for this system can be computed and found to match Legendre polynomials. However, despite making this connection and writing out formulas for this SSM, Voelker did not provide a complete proof of either of these two connections.

The preceding two steps that motivated the LDN can be informally written as the chain of transformations (i) transfer function $e^{-s}$ $\to$ (ii) SSM $(\bm{A},\bm{B},\bm{C})$ $\to$ (iii) Legendre polynomials $e^{t\bm{A}}\bm{B}$. The HiPPO framework in a sense proceeded in the opposite direction. Gu et al. started by defining the system that convolves with truncated Legendre polynomials, and with a particular differentiation technique showed that it could be written as a particular SSM which they called HiPPO-LegT. This SSM turned out to be the same (up to a minor change in scaling) as the original $(\bm{A},\bm{B})$ defined by the LMU, thus proving the second of the two steps relating this particular SSM to the Legendre polynomials.

In this work, we show the final piece in this reverse chain of equivalences. In particular, we start from the LegT SSM $(\bm{A},\bm{B},\bm{C})$ and directly prove that its transfer function produces Padé approximants of the exponential. Our proof introduces new techniques in an inductive argument that can be applied to HiPPO SSMs beyond the LegT case, and relates them to continued fraction expansions of the exponential.

We comment on a minor difference between the parameterization of HiPPO-LegT and the LMU. The LMU is originally defined as

where $\theta$ is a hyperparameter that controls the length of the window. However, we point out that such constant scaling of the SSM is also controlled by the step size $\Delta$ as discussed in Section 3.3. Therefore $\theta$ is redundant with $\Delta$, so the LegT matrices defined in and in this work do not have a concept of $\theta$. Additionally, in this work we redefine the LegT matrices $(\bm{A},\bm{B})$ to be scaled by a factor of $2$ to make them properly timescale normalized, using the theory developed in Section 3.3.

As discussed in Section 2.2, HiPPO can be thought of as a framework for deriving state space models corresponding to specific polynomial bases. The original paper did not explicitly draw the connection to state space models, and also developed systems only for a few particular cases which were called LegS (a time-varying system involving Legendre polynomials), LegT (a time-invariant system with the truncated Legendre polynomials), and LagT (involving Laguerre polynomials).

A follow-up paper on Linear State Space Layers (LSSL) generalized these results to all orthogonal polynomial families, and also generalized the flexibility of the time-varying component. They produced SSMs $x^{\prime}(t)=\bm{A}(t)x(t)+\bm{B}(t)u(t)$ where at all times $t$, $x(t)$ can be viewed as the projection of the history of $u(s)\mid_{s\leq t}$ onto orthogonal polynomials $p_{n}$ rescaled onto the interval $[t-\theta(t),t]$, where $\theta(t)$ is an arbitrary factor. This generalized all 3 cases of the original HiPPO paper.

Compared to these works, our framework (Definition 2) simplifies and generalizes the concepts directly in terms of (time-varying) state space models. We define a more natural concept of orthogonal SSM, derive very general instantiations of it (Section 3.1), and flesh out its properties (Section 3.3). Our general result subsumes all prior cases including all cases of the LSSL as a direct corollary. Some concrete advantages include:

It allows more flexible transformations of polynomial bases, such as including a change-of-basis inside the polynomials. The previously expained case of LegS is an instance of this, which has basis functions $L(e^{-t})$ with an exponential change of basis, instead of vanilla polynomials.

It can be applied to non-polynomial bases, such as the truncated Fourier basis FouT.

It does not require considering multiple cases depending on where the basis functions are supported. Instead, we handle this by considering discontinuities in the basis functions.

While the preceding discussion covers theoretical interpretations of SSMs, S4 (and its predecessor LSSL) are the application of these SSMs to deep learning. In comparison to prior works such as the LMU and HiPPO which require a pre-determined system $(\bm{A},\bm{B})$ and incorporate them naively into an RNN, LSSL and S4 use a full state space model $(\bm{A},\bm{B},\bm{C})$ as a completely trainable deep learning layer. Doing this required resolving computational problems with the SSM, which was the main focus of S4. In this work, we make a distinction between HiPPO, which is the theoretical derivation and interpretation of particular SSMs $(\bm{A},\bm{B})$, and S4, which is the incorporation of those SSMs as a trainable deep learning layer with a particular algorithm.

Appendix B Experiment Details and Additional Experiments

The Delay Task consists of input-output pairs where the input is a white noise signal of length 4000 bandlimited to 1000 Hz. The output is the same signal shifted by 1000 steps (Fig. 7(a)). We use single layer linear SSMs with $H=4$ hidden units and state size $N=1024$. Models are trained with the Adam optimizer with learning rate 0.001 for 20 epochs.

B.2 Long Range Arena

The settings for LRA use the same hyperparameters in . A more detailed protocol can be found in . To be self-contained, we recreate the same table of parameters in Table 4.

Appendix C Proof Details

We furnish the missing proofs from Section 2 in Section C.1. We will describe our general framework and results in Section C.2, and prove the results in Sections 3.1, 3.2 and 3.3 in Sections C.3, C.4 and C.5 respectively.

This corresponds to results from Section 2.

The SSM kernels are $K_{n}(t)=e^{-t(n+1)}\bm{B}_{n}$. Assume $\bm{B}_{n}\neq 0$ so that the kernels are not degenerate.

Suppose for the sake of contradiction that this was a TOSSM with measure $\omega(t)$. Then we must have

Plugging in $n=1,m=1$ and $n=0,m=2$ gives

C.2 General theory

where the integrals in this paper are over the range $[-\infty,\infty]$, unless stated otherwise.

This is sufficient to derive an OSSM based on the HiPPO technique. The generalized HiPPO framework demonstrates how to build (T)OSSMs utilizing time warping to shape the time interval and tilting to construct new sets of orthogonal basis functions.

We will need the notion of a “time warping” function $\overline{\sigma}$ as follows:

We will be using a special case of time-warping function, which we say has a discontinuity at $t_{0}$ for some $t_{0}\in(-\infty,t]$:

We allow for $t_{0}=-\infty$, in which case we think of the interval $[t_{0},t]$ as $(-\infty,t]$.

Before proceeding, let us clarify our notation. We will use $\sigma_{t}$ and $\sigma_{s}$ to denote the partial derivatives $\frac{\partial}{\partial t}\sigma(t,s)$ and $\frac{\partial}{\partial s}\sigma(t,s)$ respectively. We will drop the parameters $(t,s)$ and use $f$ instead of $f(t,s)$ when it is clear from context to reduce notational clutter. Further, we will extend this notation to function composition, i.e. write $g\circ f(t,s))$ as $g(f)$ and function product, i.e. use $fgh$ instead of $f(t,s)g(t,s)g(t,s)$. Finally, we’ll shorten $fgh\circ\phi(t,s)$ as $fgh(\phi)$.

We also define the tilting $\chi$ and show that regardless of warping, we can construct a new orthogonal basis (note that the result holds for warping functions as in (9) and not just those as in (10)).

For the set of orthonormal functions $\{p_{n}\}_{n=0}^{N-1}$ orthogonal over measure $\omega I$, the set of basis functions

for time-warping function $\sigma$ satisfying (9) and any $\chi(t,s)$ that is non-zero in its support.

Consider the following sequence of equalities:

In the above, the second equality follows from the substitution $y\leftarrow\sigma(t,s)$ and hence $dy=\sigma_{s}ds$ and the final equality follows from (8). Then since $\chi(t,s)$ is always non-zero, we have

Without loss of generality, we can split $\chi$ into a product

of one part that depends on $\sigma$ and another arbitrary component.

Since we have an orthonormal basis and measure, we can try to derive the (T)OSSM. For a given input signal $u(t)$, the HiPPO coefficients are defined as the projections.

defined as inner product of $u(t)$ with the tilted basis functions $\chi p_{n}$ with respect to the measure $\mu$ as defined in Lemma C.1. For additional convenience, we use the decomposition $\chi=\psi^{-1}\phi^{-1}$ from (11) to get:

The HiPPO technique is to differentiate through this integral in a way such that it can be related back to $x_{n}(t)$ and other $x_{k}(t)$. We require for every $n$, we require that there are a set of coefficients $\{\gamma_{nk}\}_{k=0}^{N-1}$ such that

Consider a set of basis functions $p_{n}$ orthogonal over $\omega$, time warping $\overline{\sigma}(t,s)$ as in (9), (10), and tilting $\chi$ as in (11) and (14) with the functions $\sigma,p_{n},\omega,\psi$ obeying (13). If $\frac{dt_{0}}{dt}\neq 0$, further assume that for some vector $\bm{A^{\prime}}$, we have as $N\to\infty$,

Applying the Leibniz rule to (12), we get

Let us consider each term separately. The first term

corresponds to the differentiation of the basis functions and measure. In order to relate this to $\{x_{k}(t)\}$, it suffices that $\sigma_{t}(p_{n}\omega\psi)^{\prime}(\sigma)$ satisfies (13) which implies that when we vectorize this, we get $x^{(0)}(t)=\bm{A}^{0}\cdot x(t)$.

For additional warping and tilting terms, we consider

To reduce this term to $x_{n}(t)$, recall from (10) that

where $c(t),d(t)$ are defined as in (10) and (14).

We will end up with $x^{(1)}_{n}(t)=(c(t)+d(t))x_{n}(t)$. This leads to the the vectorized form $x^{(1)}(t)=(c(t)+d(t))\bm{I}x(t)$.

where $H(x)$ is the “heaviside step function.” It is know that $H^{\prime}(x)=\delta(x)$, which implies

Using the above in RHS of (17), we separate out $x^{(2)}_{n}(t)$ and $x^{(3)}_{n}(t)$ as follows. First, define

In the last equality, we have used the fact that $\sigma(t,t)=\overline{\sigma}(t,1)=1$ by definition. It follows that in vectorized form we have $x^{(2)}(t)=\bm{B}u(t)$.

If $\frac{dt_{0}}{dt}=0$, then we have $\bm{D}=\bm{0}$ and hence we have $x^{(3)}(t)=\bm{0}=-\overline{c}\bm{D}\left(\bm{A^{\prime}}\right)^{\top}x(t)-\overline{d}\bm{D}u(t)$

If $\frac{dt_{0}}{dt}\neq 0$, then as $N\to\infty$, from (15), the above comes out to

It follows that in vectorized form we have $x^{(3)}(t)=-\overline{c}\bm{D}\left(\bm{A^{\prime}}\right)^{\top}x(t)-\overline{d}\bm{D}u(t)$. The result follows after combining the terms.

We see that the behavior of is the model is dictated by $t_{0}$. In particular, in this paper, we will consider two special cases.

Follows from Theorem 12. Since $t_{0}$ is independent of $t$, then $\frac{\mathop{}\!dt_{0}}{\mathop{}\!dt}=0$, and $\bm{D}=\bm{0}$. ∎

This follows directly from Theorem 12 by setting $t_{0}=t-\theta$. ∎

C.3 LegS (and LSSL?)

i.e. the measure is completely “tilted” away, and let

Let’s consider the special case of (18) where $b(t)=0$. This is most generally satisfied by

Note that the condition $\sigma(t,t)=1$ forces $z=-a$. Hence, we have

We now consider the following special case of Corollary C.2:

Let $\eta\geq 0$. The SSM $(-a^{\prime}(t)(\bm{A}+(\eta+1)\bm{I}),a^{\prime}(t)\bm{B})$, where $t_{0}$ is independent of $t$, is an OSSM for basis functions and measure

Given a orthonormal basis $p_{0},p_{1},\dots,p_{N-1}$ with respect to a measure $\omega$. Note that time-warping function $\sigma$ satisfying (19) implies that $\sigma_{s}=a^{\prime}(s)\sigma$.

We fix tilting $\chi(t,s)=\frac{\omega(\sigma)}{\sigma^{\eta}}$, which in turn follows by setting

We show shortly that we satisfy the pre-conditions of Corollary C.2, which implies (with our choice of $\chi$ and $\sigma$) that we have an OSSM with basis functions $p_{n}(t,s)=\frac{\omega(\sigma)}{\sigma^{\eta}}p_{n}(\sigma)$ and measure

To complete the proof, we show that out choice of paramters above satisfies the conditions of Corollary C.2 (by showing they satisfy the conditions of Theorem 12). We verify that $\sigma$ and $\phi$ satisfy (10) and (14), noting that

This implies that setting $c(t)=-a^{\prime}(t)$ and $d(t)=-\eta a^{\prime}(t)$ is enough to satisfy (10) and (14).

Further, note that (19) and the fact that $\psi=\omega^{-1}$ imply that

It follows that (13) is satisfied as long as

for some set of coefficients $\{\alpha_{nk}\}_{k=0}^{N-1}$, which is exactly (21). This implies the $\gamma_{nk}$ in Corollary C.2 satisfy.

Let $\bm{A}$ be the matrix such that $\bm{A}_{nk}=-\alpha_{nk}$ and then note that $-a^{\prime}(t)(\bm{A}+(\eta+1)\bm{I})$ is exactly the first parameter of the SSM in Corollary C.2. Similarly, recall in Corollary C.2

where the final equality follows since in our case, $\sigma_{s}(t,t)=a^{\prime}(t)\exp(a(t)-a(t))=a^{\prime}(t)$. Overloading notation and letting $\bm{B}_{n}=p_{n}(1)$, all conditions of Corollary C.2 hold, from which the claimed result follows. ∎

We are particularly interested in the following two special cases of Corollary C.4.

Letting $a^{\prime}(t)=\frac{1}{t}$ implies that $a(t)=\ln{t}$. Then we can observe that is a case of Corollary C.4 with time warping

We set $\eta=0$ in Corollary C.4, which in turn sets $\phi=\sigma^{0}=1$. This gives the tilting

Then by Corollary C.4, it follows that that we can use $\sigma$ and $\chi$ to build an OSSM with basis functions

This is a case of Corollary C.4 where $a^{\prime}(t)=1$, $\sigma=\exp(s-t)$, and we pick $\eta=0$, implying that $\phi=\sigma^{0}=1$. It follows that

Utilizing Corollary C.4, we can use $\sigma$ and $\chi$ to build an OSSM with basis functions

Next we instantiate Corollary C.4 to prove Corollary 3.1. (Even though strictly not needed, we instantiate Corollary C.6 and Corollary C.5 to prove Theorem 5 and Corollary 3.3.) To that end, we will need the following result:

Let the Legendre polynomials orthonormal over the interval $$be denoted as$L_{n}$. Then

The Legendre polynomials satisfy the following orthogonality condition over $$:

Let us denote the normalized Legendre polynomials orthogonal over $$as$\lambda_{n}P_{n}(z)$where$\lambda_{n}=\sqrt{\frac{2n+1}{2}}$. To orthogonalize them over$$, let$y=\frac{1+z}{2}$. It follows that$z=2y-1$,$\mathop{}\!dz=2\mathop{}\!dy$. Note that we then have

then we have an a set of functions over $$ such that

From [2, (2.8), (2.9)], note that $P_{n}(-1)=(-1)^{n}$ and $P_{n}(1)=1$. This implies that

(25) implies that $P_{n}(z)=\frac{L_{n}(y)}{\sqrt{2n+1}}$, thus

Let $L_{n}$ be the Legendre polynomials orthonormal over the interval $$. Define$\sigma(t,s)=\exp(a(s)-a(t))$. The SSM$(a^{\prime}(t)\bm{A},a^{\prime}(t)\bm{B})$ is an OSSM with

where $\bm{A}$ and $\bm{B}$ are defined as in (4).

We consider our basis functions, the Legendre polynomials, which are orthogonal with respect to unit measure. This allows us to invoke Corollary C.4 with $\omega=1$. Further, here we have $t_{0}=-\infty$ and $\eta=0$. Now we have an SSM:

where $\bm{A}^{0}_{nk}=\alpha_{nk}$ as in (21) and $\bm{B}_{n}=L_{n}(1)$.

From (24) observe that $\bm{B}_{n}=(2n+1)^{\frac{1}{2}}$. From (22), we have

We write that $\bm{A}=-(\bm{A}^{0}+\bm{I})$. Indeed,

Thus the $\bm{A}$ and $\bm{B}$ match those in (4), which completes our claim. ∎

We consider our basis functions, the Legendre polynomials, which are orthogonal with respect to unit measure. This allows us to invoke Corollary C.5 with $\omega=1$. Now we have

where $\bm{A}^{0}_{nk}=\alpha_{nk}$ as in (21) and $\bm{B}_{n}=L_{n}(1)$.

From (24) observe that $\bm{B}_{n}=(2n+1)^{\frac{1}{2}}$. From (22), we have

We write that $\bm{A}=-(\bm{A}^{0}+\bm{I})$. Indeed,

We consider our basis functions, the Legendre polynomials, which are orthogonal with respect to unit measure, warping function $\sigma=\exp(s-t)$, and with tilting $\chi=\omega$. We note that $\sigma=\exp(s-t)$ satisfies (19) with, $a^{\prime}(t)=1$. This allows us to invoke Corollary C.5.

This explains why removing the $\frac{1}{t}$ factor from HiPPO-LegS still works: it is orthogonalizing onto the Legendre polynomials with an exponential “warping”.

C.4 Finite Windows

Out plan is to apply Corollary C.3, for which we must show that the basis functions $L_{n}(t,s)$, time warping $\sigma(t,s)$, and tilting $\chi(t,s)=\psi^{-1}\phi^{-1}(t,s)$ satisfy (13), (10), and (14), respectively. We first set some parameters– note that because $\omega=1$ and set $\psi=\phi=1$.

The above along with (23), we see that the Legendre polynomials satisfy (13) with

We also note that $\sigma_{s}=\frac{1}{\theta}$.

satisfying (10) trivially by setting $c(t)=0$. Similarly, since $\phi=1$ (14) is also satisfied trivially by setting $d(t)=0$. Finally we note that the $L_{n}$ forms a complete basis over $$, hence as$N\to\infty$, we have

The above defines $\bm{A^{\prime}}$ by setting $\bm{\bm{A^{\prime}}}_{n}=L_{n}(0)$ (as well as $\overline{c}=1$ and $\overline{d}=0$.) Now by Corollary C.3, we have an SSM

where $\bm{D}_{n}=\frac{1}{\theta}L_{n}(0)$, and by (22) $\bm{A}^{0}_{nk}=\gamma_{nk}$ (as in (26)) and $\bm{B}^{\prime}_{n}=\frac{1}{\theta}L_{n}(1)$.

From (24), we have $\bm{D}_{n}=\frac{1}{\theta}(2n+1)^{\frac{1}{2}}(-1)^{n}$ and $\bm{B}_{n}=\frac{1}{\theta}(2n+1)^{\frac{1}{2}}$.

The proof is complete by noting that $\bm{A}^{0}-\bm{D}\left(\bm{A^{\prime}}\right)^{\top}=\frac{1}{\theta}\bm{A}$ and $\bm{B}^{\prime}=\frac{1}{\theta}\bm{B}$.

We first give a proof of Theorem 6. Then, we prove Theorem 7 as a function approximation result pertaining to S4-FouT.

C.4.2 Explanation of S4-FouT

We seek to derive $\bm{A}$ and $\bm{B^{\prime}}$ from (6) using Corollary C.3:

We use the time-warping function $\sigma(t,s)=1-(t-s)$, which implies that we have

We also have $\omega(\sigma)=1,$ and we order our bases in the form $p_{n}=(1,c_{1}(t),s_{1}(t),c_{2}(t),s_{2}(t),\ldots)$Note that this is 0-indexed., where the basis functions have derivatives:

Consequently, we can define $\gamma_{nk}$ as follows:

Further, the discontinuity is at $t_{0}=t-\theta,\ \theta=1$ which implies that $\frac{dt_{0}}{dt}=1.$ We now seek to use the stored approximation to $u$ at time $t$ to compute $u(t-1)$.

First, denote the latent state $x(t)$ with coefficients $x=(x^{1}(t),x^{c}_{1}(t),x^{s}_{1}(t),x^{c}_{2}(t),x^{s}_{2}(t),\ldots)$ and define the functions $v(s)$ and $w(s)$ such that we have

Now, let $\hat{u},\hat{v},$ and $\hat{w}$ denote the reconstruction of $u,v$ and $w,$ where we have

Towards that end, we examine the sine and cosine coefficients of $u$ and $v$ as follows:

Here, for (38) and (39), we use the change of variables $s^{\prime}\leftarrow 2t-s-1,$ which gives us

Then, we use the fact that $c_{n}(1-\sigma(t,s^{\prime}))=c_{n}(\sigma(t,s^{\prime}))$ but $s_{n}(1-\sigma(t,s^{\prime}))=-s_{n}(\sigma(t,s^{\prime})).$ That is, both $u$ and $v$ have the same cosine coefficients but negated sine coefficients of each other. But, we know that both $s_{n}(\sigma(t,t-1))=s_{n}(1-(t-(t-1)))=s_{n}(0)=0$ and $s_{n}(\sigma(t,t))=s_{n}(1-(t-t))=s_{n}(1)=0$, and hence, the reconstruction of $\hat{u}$ at the endpoints $\sigma(t,t-1)=0$ and $\sigma(t,t)=1$ depends only on the cosine coefficients, whence we assert that the reconstruction $\hat{u}$ agrees with $\hat{v}$ at both endpoints. Therefore, we have $\hat{u}(t,t)=\hat{v}(t,t)$ implying that $\hat{w}(t,t)=\hat{u}(t,t)$.

Note that $w$ is continuous and periodic, for which the basis $\{1,c_{n},s_{n}\}_{n}$ is complete, and hence, we know that as $N\to\infty,$ $\hat{w}\to w.$ Thus, at $s=t,$ we have $\hat{u}(t,t)=\hat{w}(t,t)=w(t)=\frac{u(t)+v(t)}{2}=\frac{u(t)+u(t-1)}{2},$ which completes the proof of the claim in (37).

Recall from (34) that we can express the stored approximation of $u(t)$, given by $\hat{u}(t,s),$ as follows:

For the value at $t,$ the approximation $\hat{u}(t,t)$ is then given by

Due to (37), we know $u(t-1)=2\hat{u}(t,t)-u(t)$, which combined with the above yields:

Finally, with regards to Corollary C.3, for Theorem 12, (29) satisfies (10) and (32) satisfies (14) with (33) satisfying (13) for $\bm{A}^{0}.$ Moreover, from (40), we can take $\overline{c}=1,\overline{d}=-1,$ and $\bm{A^{\prime}}_{k}:=\begin{cases}2&k=0\\ 2\sqrt{2}&k\text{ odd}\\ 0&\text{ otherwise}\end{cases}$ to satisfy (15).

Invoking Corollary C.3 now yields the following OSSM: Recall that, like the coefficients, the matrices are 0-indexed.

where $\bm{A}^{0}_{nk}=\gamma_{nk}$ with $\bm{D}_{n}$ and $\bm{B}_{n}$ specified as follows:

Here, the values are derived from the expressions of Corollary C.3:

Recall that we have $p_{n}\in\{1,c_{n},s_{n}\},\omega(t,s)=1$, and from (27) and (30), $\sigma_{s}(t,s)=1$ with $\psi(t,s)=\phi(t,s)=1.$ Thus, (41) is due to $1(0)\cdot 1=1,s_{n}(0)\cdot 1=0$ but $c_{n}(0)\cdot 1=\sqrt{2}.$ Similarly, (42) is because $1(0)\cdot 1=1,s_{n}(1)\cdot 1=0$ but again $c_{n}(1)\cdot 1=\sqrt{2}.$

As $c(t)=d(t)=0,$ we define $\bm{A}\leftarrow\bm{A}^{0}-\overline{c}\bm{D}\left(\bm{A^{\prime}}\right)^{\top}$ and $\bm{B}\leftarrow\bm{B}-\overline{d}\bm{D}$, given by

C.4.3 Function Approximation Error

First, the state size being $N$ dictates that there are $\lfloor N/2\rfloor$ $s_{n}$ and $c_{n}$ basis functions each. We fix time $t$ and denote $x_{n}^{c}$ and $x_{n}^{s}$ to be the respective coefficients for $s_{n}$ and $c_{n}$ basis corresponding to S4-Fou. Since $\{s_{n},c_{n}\}_{n\geq 0}$ forms an orthonormal basis, by Parseval’s identity, we have

Thus, in order to bound the error, it suffices to bound the high-order coefficients by integration by parts as follows:

The quantity in the bracket vanishes as $s_{n}$ is periodic. Therefore

where we use the fact that $K$ is $L-$Lipshitz. For $x_{n}^{s}$, a similar argument holds and we get:

We use (44) to get the following estimate on $\lVert{K-\hat{K}}\rVert:$

Thus, it suffices for $N$ to satisfy the following inequality:

We now use the same argument as above to the fact that $K$ has order-$k$ bounded derivative. By iteration, we get:

Again, due to (43), this then gives us the following estimate on the square error:

If $K$ has order $k-$bounded derivatives, then we use (45) to get the following estimate on $\lVert{K-\hat{K}}\rVert:$

Again, it suffices for $N$ to satisfy the following inequality:

C.4.4 Delay Network

Finally, we prove Theorem 9. Note that this is a stronger version of the LegT portion of Theorem 8, while the FouT portion is a corollary of the proof of Theorem 6.

We start by working out some calculations concretely to provide an example. The SSM corresponding to HiPPO-LegT is

(In the RHS and for the rest of this part, we will redefine $\bm{A}$ to be the $\pm 1$ matrix found above for convenience.)

We have $\bm{A}=-1,\bm{B}=\bm{C}=1$, and the transfer function is $\bm{C}(s\bm{I}-\bm{A})^{-1}\bm{B}=\frac{1}{1+s}$.

It can be verified that this is indeed $[1/2]_{\exp}(-s)$.

We will now sketch out a method to relate these transfer functions recursively.

We will redefine $\bm{Z}$ to be the vector that ENDS in $+1$.

Now we can use the block matrix inversion formula.https://en.wikipedia.org/wiki/Block_matrix#/Block_matrix_inversion Ideally, this will produce a recurrence where the desired transfer function $\bm{Z}_{n}(s\bm{P}_{n}^{-1}-\bm{A}_{n})^{-1}\bm{1}_{n}$ will depend on $\bm{Z}_{n-1}(s\bm{P}_{n}^{-1}-\bm{A}_{n-1})^{-1}\bm{1}_{n-1}$. However, looking at the block matrix inversion formula, it becomes clear that there are also dependencies on terms like $\bm{1}_{n-1}^{\top}(s\bm{P}_{n-1}^{-1}-\bm{A}_{n-1})^{-1}\bm{1}_{n-1}$ and $\bm{Z}_{n-1}(s\bm{P}_{n-1}^{-1}-\bm{A}_{n-1})^{-1}\bm{Z}_{n-1}^{\top}$.

The solution is to track all of these terms simultaneously.

Instead of using the explicit block matrix inversion formula, it will be easier to work with the following factorization used to derive it (block LDU decompositionhttps://en.wikipedia.org/wiki/Schur_complement#/Background).

Using Lemma C.12, we can factor the inverse as

Now we can derive the full recurrence for all these functions.

Now we’ll define a few transformations which will simplfy the calculations. Define

This will be the most important term, as it determines the denominator of the expressions. Simplifying the recurrence slightly gives

Now let $G_{n}^{1z}(s)=\frac{P_{n}^{1z}(s)}{Q_{n}^{1z}(s)}$ where $P,Q$ are polynomials. Clearing the denominator $Q$ yields

But this is exactly the fundamental recurrence formula for continuants of the continued fraction

Therefore $Q_{n-1}^{1z}(s)$ are the denominators of the Pade approximants.

First note that $G_{n}^{11}(s)=G_{n}^{zz}(s)$ is straightforward from the fact that their recurrences are identical. The recurrence is

But note that this is exactly satisfied by the Padé approximants, by the determinantal formula of continued fractions. This shows that $G_{n-1}^{1z}(s)$ are the Padé approximants of $e^{-s}$, as desired.

C.5 Normalization and Timescales

Consider a TOSSM $(\bm{A},\bm{B})$ for basis functions $p_{n}(t)$ and measure $\omega(t)$. Then, the following are also TOSSMs with the corresponding basis functions and measure:

Constant scaling changes the timescale: $(c\bm{A},c\bm{B})$ is a TOSSM with basis $p(ct)$ and measure $\omega(ct)c$.

Identity shift tilts by exponential: $(\bm{A}+c\bm{I},\bm{B})$ is a TOSSM with basis $p(t)e^{-ct}$ and measure $\omega(t)e^{2ct}$.

Unitary change of basis preserves measure: $(\bm{V}\bm{A}\bm{V}^{*},\bm{V}\bm{B})$ is a TOSSM with basis $\bm{V}p(t)$ and measure $\omega(t)$.

We define $p(t)$ to be the vector of basis functions for the OSSM $(\bm{A},\bm{B})$,

Recall that the SSM kernels are $K_{n}(t)=p_{n}(t)\omega(t)$ so that $p(t)\omega(t)=e^{t\bm{A}}\bm{B}$.