PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Albert Gu, Isys Johnson, Aman Timalsina, Atri Rudra, Christopher Ré

arXiv:2206.12037v2cs.LG

TL;DR

S4’s HiPPO initialization had lacked a mathematical interpretation when used as a time-invariant system, leaving its long-range mechanism and timescale choice unclear. The paper generalizes HiPPO to explain S4 through exponentially warped Legendre projections, derives other basis variants, and reports strong long-range results. Its analysis also identifies scope boundaries among variants and initialization choices.

  • Problem

    The time-invariant use of S4’s HiPPO matrix lacked a mathematical interpretation, and the meaning and initialization of its important timescale parameter Δ were unclear.

  • Method

    The paper generalizes HiPPO to derive orthogonal-basis state-space models, interpreting S4 as exponentially warped Legendre projections and deriving Fourier-based variants.

  • Results

    S4-LegS remains the best S4 variant for long-range dependencies, while S4-FouT performs best on a continuous memorization task when Δ is initialized correctly.

  • Takeaways & Limitations

    The framework explains S4’s long-range behavior and provides principled guidance for choosing dependency lengths and initializing SSM parameters.

  • Takeaways & Limitations

    The authors do not recommend HiPPO-LagT because it does not represent orthogonalization against an exponentially decaying measure, and LegT produces a weaker spike function than FouT.

Abstract

from arXiv · show

Linear time-invariant state space models (SSM) are a classical model from engineering and statistics, that have recently been shown to be very promising in machine learning through the Structured State Space sequence model (S4). A core component of S4 involves initializing the SSM state matrix to a particular matrix called a HiPPO matrix, which was empirically important for S4's ability to handle long sequences. However, the specific matrix that S4 uses was actually derived in previous work for a particular time-varying dynamical system, and the use of this matrix as a time-invariant SSM had no known mathematical interpretation. Consequently, the theoretical mechanism by which S4 models long-range dependencies actually remains unexplained. We derive a more general and intuitive formulation of the HiPPO framework, which provides a simple mathematical interpretation of S4 as a decomposition onto exponentially-warped Legendre polynomials, explaining its ability to capture long dependencies. Our generalization introduces a theoretically rich class of SSMs that also lets us derive more intuitive S4 variants for other bases such as the Fourier basis, and explains other aspects of training S4, such as how to initialize the important timescale parameter. These insights improve S4's performance to 86% on the Long Range Arena benchmark, with 96% on the most difficult Path-X task.

1 Introduction

S4 is efficient for long sequences but its HiPPO initialization and training choices were not fully understood. This work gives HiPPO a general mathematical interpretation, derives new basis-function variants, and reports strong benchmark performance.

  • Motivation: S4 combines recurrent simulation with convolutional computation, providing efficiency for online inference and parallel training on long sequences.Its linear time-invariant state-space structure supports both forms.
  • Open questions: The prior use of the HiPPO matrix in a time-invariant SSM lacked a mathematical interpretation, leaving S4’s long-range mechanism unresolved.The meaning and initialization of the timescale parameter Δ were also poorly understood.
  • Contributions: The generalized HiPPO framework proves that S4’s state matrix decomposes inputs using exponentially warped Legendre polynomials and an exponentially decaying measure.This supplies a theoretical explanation for the smooth basis functions associated with long dependencies.
  • Contributions: The framework derives S4-FouT, which uses truncated Fourier basis functions and can represent sliding Fourier transforms and local convolutions.The construction also supports spike-function representations for memorization tasks.
  • Training insights: The framework interprets Δ as controlling dependency length and clarifies how to initialize it and other SSM parameters.The paper also discusses variance-preserving and stable initialization through the final SSM parameter C.
  • Empirical validation: 86% average performance on Long Range Arena and 96% on Path-X are reported for the original S4 method.The paper presents these results as a new state of the art for the benchmark evaluation.

2 Background

State-space models map signals into latent dynamical states and can become convolutional systems when time invariant. HiPPO constructs systems whose states encode online projections of signal history onto orthogonal bases, with Legendre variants providing the main background cases.

  • State-space models: An SSM maps a one-dimensional input signal u(t) to an N-dimensional latent state x(t), then projects it to an output signal y(t).The paper assumes D = 0 unless stated otherwise.
  • State-space models: Constant SSM matrices form a linear time-invariant system equivalent to a convolution, enabling parallel and efficient computation.The associated impulse response is the output produced by an impulse input.
  • State-space models: The SSM convolution kernel is a linear combination of SSM basis functions controlled by the output vector C.The paper distinguishes these general SSM basis kernels from the bases and measures used for orthogonal SSMs.
  • Discrete SSMs: For discrete inputs, the timescale Δ sets the discretization resolution and can equivalently be viewed as scaling the continuous-time parameters A and B.The paper motivates interpreting Δ through the dependency length in a discrete sequence rather than only as a sampling step.
  • HiPPO variants: S4 uses a particular HiPPO initialization, while this paper distinguishes the corresponding S4-LegS, prior LegT, and new FouT variants.The HiPPO matrices were originally derived in several special cases rather than representing a single unique matrix.
  • HiPPO framework: HiPPO states are designed to reconstruct the input history from online projections onto an orthogonal basis under a specified measure.This is formulated as an online function-approximation problem.
  • HiPPO framework: Complete orthogonal bases permit reconstruction in the infinite-state limit, but general SSMs need not be orthogonal SSMs and diagonalizing a valid system may break that structure.The paper gives the stronger limitation that no TOSSM has diagonal state matrix diag{−1, −2, …}.
  • Legendre variants: HiPPO-LegT approximates truncated Legendre polynomials on sliding windows and is also known as the Legendre Delay Network or Legendre Memory Unit.Its original motivation was approximating a one-unit delay through a spike-like impulse response.

3 Generalized HiPPO: General Orthogonal Basis Projections

The generalized HiPPO framework derives time-invariant orthogonal SSMs from broad choices of basis functions, clarifying S4-LegS and enabling Fourier-based variants. It also connects timescale and normalization to dependency length, while establishing approximation and empirical results for delay and convolutional behavior.

  • General framework: Theorem 12 provides a general construction of orthogonal SSMs for varied basis functions, subsuming earlier HiPPO results.The framework covers both polynomial and non-polynomial bases.
  • S4-LegS: S4-LegS is interpreted as orthogonal projection onto Legendre polynomials with an exponential warping of the time axis.This gives the time-invariant system a mathematical interpretation connected to its basis functions.
  • S4-FouT: S4-FouT derives a time-invariant SSM from truncated Fourier functions and can represent any local convolution.The Fourier basis is orthonormal on [0, 1], and its linear combinations can represent any function on that interval.
  • Delay networks: Finite-window LegT and FouT approximate delay networks, with their kernels converging to δ(t −1) as N →∞.For LegT, the transfer function is also the [N −1/N] Padé approximant to e−s.
  • Timescales and normalization: The timescale ∆ controls dependency length, with normalized systems modeling expected length 1 and ∆modulating it to length 1 ∆.In discrete time, LegT and FouT represent sliding windows of 1/∆elements; HiPPO-LegS motivates log-uniform initialization over expected lengths 10 to 1000.
  • Normalization and empirical behavior: The framework yields variance-preserving SSMs and identifies a limitation of LegT: despite Padé-based delay approximation, it produces a weaker spike than FouT and performs slightly worse on synthetic tasks.A random vector C with independent mean 0, variance 1 entries produces outputs matching input variance.

4 Experiments

The experiments compare S4 variants across long-range, reconstruction, and memorization tasks, testing hypotheses about basis choice and timescale initialization. Results support theory-dependent specialization: S4-LegS handles sparse long dependencies, S4-FouT handles dense memorization, and correct timescales are important.

  • Experimental hypotheses: S4-LegS is hypothesized to excel on sparse memorization and very long dependencies, whereas S4-FouT targets dense memorization but finite local windows.LegS smooths signals against an infinitely long measure, while FouT represents spike functions over a finite window.
  • Long Range Arena: S4 improves the previous Long Range Arena state of the art by 6 points.
  • Long Range Arena: S4-LegS is extremely strong on Path-X, whose sparse dependencies have length 16384, while FouT cannot solve it because it is a finite-window method.
  • Long Range Arena: Path-X performance depends on timescale initialization: lowering ∆min to 10^4, on the order of task length L = 16384, improves results, while spreading ∆ across timescales can capture different resolutions.The experiments report that a wider ∆max range can support features at different resolutions, which are combined across network layers.
  • Function reconstruction: LegT and FouT reconstruct uniform-measure histories far better than other methods when timescales are appropriate, but perform worst when timescales are mis-specified.The reconstruction task uses white-noise inputs and requires the final output to reconstruct the last 1000 input elements.
  • Memorization: On the Delay Task, S4-FouT outperforms S4-LegT, and both outperform other methods when ∆ is set correctly.The task requires sequence outputs to reproduce inputs shifted by 1000 steps; LSTM and Transformer baselines perform no better than random guessing, with RMSE 0.43.

5 Summary: How to Train Your HiPPO

The HiPPO framework derives state-space matrices from prescribed basis functions, clarifying SSM kernels and enabling Legendre- and Fourier-based S4 variants. These variants characterize long-range and local dependencies through different kernel families and timescales.

  • SSMs represent convolution kernels as linear combinations of basis functions parameterized by A, B, and C.
  • HiPPO derives matrices A and B for prescribed families of well-behaved basis functions.
  • HiPPO-LegS produces a smooth, long-range kernel family and remains the strongest method for long-range dependencies among S4 variants.
  • HiPPO-FouT uses truncated Fourier functions to subsume local convolutions and capture sliding DFT or STFT transforms.
  • The timescale ∆ controls dependency length, with finite-window FouT kernels having length exactly 1/∆.

A Related Work

The paper places HiPPO within prior Legendre, delay-network, and LSSL work, then presents a framework that simplifies and generalizes orthogonal state-space models. Its formulation covers polynomial and non-polynomial bases while handling finite-support discontinuities uniformly.

  • Delay networks: For LegT, the transfer function of the resulting SSM is the [N −1/N] Padé approximant to e^-s, completing the reverse connection to delay networks.
  • Parameterization: The timescale parameter θ in the LMU is redundant with SSM step size ∆, which also controls constant scaling and timescale normalization.
  • HiPPO and LSSL: Prior HiPPO work developed LegS, LegT, and LagT systems for specific polynomial bases and time-varying or time-invariant settings.
  • HiPPO and LSSL: LSSL generalized the framework to orthogonal polynomial families by projecting input history onto rescaled polynomials over a variable window.
  • HiPPO and LSSL: The paper’s framework simplifies and generalizes prior results through a natural orthogonal SSM definition that subsumes LSSL as a corollary.
  • Generalization: The framework permits flexible polynomial transformations, including the exponentially changed LegS basis L(e^-t).
  • Generalization: It also supports non-polynomial bases such as truncated Fourier functions and handles finite-support cases through basis-function discontinuities.
  • S4: S4 and LSSL make the full (A, B, C) state-space layer trainable, addressing computational issues in applying these models to deep learning.

B.1 Delay (Continuous Copying) Task

The Delay Task evaluates whether a linear SSM can reproduce a long-delayed version of a bandlimited white-noise signal under a fixed training configuration.

  • The task shifts a bandlimited white-noise input by 1000 steps, with input length 4000 and bandwidth 1000 Hz.
  • Models use single-layer linear SSMs with H = 4 hidden units and state size N = 1024.
  • Training uses Adam with learning rate 0.001 for 20 epochs.
  • The LRA settings reuse the hyperparameters from prior work, with a detailed parameter table provided for self-containment.

C Proof Details

The appendices supply missing proofs, describe the general framework and results, and prove the main results from Sections 3.1–3.3.

  • Appendix C.1 furnishes missing proofs from Section 2.
  • Appendix C.2 describes the general framework and results, while Appendices C.3–C.5 prove the results from Sections 3.1–3.3.

C.1 Proofs from Background

The proof establishes uniqueness of the orthogonal basis and shows that the SSM kernels cannot satisfy the TOSSM conditions under the stated assumption.

  • Proposition 1: Completeness and orthogonality force the alternative basis functions q_n to equal the original basis functions p_n.The proof derives c_{ℓ,j}=δ_{ℓ,j}, implying p_ℓ=q_ℓ for every index.
  • Proposition 3: The kernels K_n(t)=e^−t(n+1)B_n are assumed nondegenerate when B_n≠0.This sets up the contradiction argument for showing they do not define a TOSSM under the proposed measure.
  • Proposition 3: Substituting n=1,m=1 and n=0,m=2 yields a contradiction with the required TOSSM relations.The contradiction completes the proof that the assumed TOSSM representation cannot hold.

C.2 General theory

The generalized HiPPO framework constructs orthogonal state space models by combining time warping and basis tilting, then differentiating projection coefficients into an SSM.

  • Construction: Time warping reshapes the modeled interval, while tilting constructs new orthogonal basis functions under the transformed measure.The resulting basis remains orthogonal for admissible warping functions and nonzero tilting functions.
  • Construction: The framework starts with orthonormal basis functions over a measure and defines HiPPO coefficients as projections of the input signal.These coefficients are inner products with tilted basis functions under the warped measure.
  • Derivation: Differentiating the projection integrals decomposes the state dynamics into basis, boundary, input, and warping or tilting contributions.The vectorized terms produce the state transition and input matrices of the resulting OSSM.
  • Derivation: Theorem 12 gives an OSSM with transition matrix A0+(c(t)+d(t))I−cD(A′)^⊤ and input matrix B−dD.The construction applies when the stated basis, warping, tilting, and limiting conditions hold.
  • Special cases: The model behavior depends on the lower interval endpoint t0, with separate corollaries for t0 independent of t and t0=t−θ.These choices yield distinct OSSM forms and associated measures.

C.3.1 Explanation of S4-LegS

The Legendre constructions instantiate generalized HiPPO with exponential or finite-window warping, yielding OSSMs and explaining HiPPO-LegS as exponential warping of Legendre projections.

  • General Legendre construction: Corollary C.4 constructs OSSMs from a time-warping function and a nonnegative tilting parameter η.Its transition matrix is −a′(t)(A+(η+1)I), with input matrix a′(t)B.
  • HiPPO-LegS: For Legendre polynomials, exponential warping σ(t,s)=exp(a(s)−a(t)) produces the time-scaled SSM (a′(t)A,a′(t)B).The Legendre basis is orthogonal under unit measure, allowing the general corollary to apply.
  • HiPPO-LegS: The time-invariant construction uses basis functions L_n(e^−t) with measure I[t0,t]e^−t.The associated state equation is x′(t)=(A+I)x(t)+Bu(t).
  • Interpretation: Removing the 1/t factor from HiPPO-LegS still works because the system orthogonalizes onto Legendre polynomials with exponential warping.This supplies the paper’s interpretation of the time-invariant LegS construction.

C.4.3 Function Approximation Error

The appendix bounds approximation error for truncated basis expansions and derives recursive transfer-function relations, including connections to Padé approximants and TOSSM closure properties.

  • Fourier approximation: Fourier-basis approximation uses Parseval’s identity and integration by parts to bound high-order coefficients.Periodicity removes the boundary term in the integration-by-parts argument.
  • Function approximation: For kernels with k-bounded derivatives, the approximation error satisfies ∥K−K̂∥²≤Lπ^k(N−2)^−k+1/2.Increasing the state size N improves the bound for sufficiently smooth kernels.
  • Transfer functions: The transfer-function derivation tracks four related functions simultaneously because block inversion introduces dependencies beyond the desired recurrence.A block LDU factorization is used to obtain the full recurrence.
  • Transfer functions: The resulting recurrence matches the fundamental recurrence for continuants, whose denominators are those of Padé approximants.The connection follows because the relevant polynomial recurrences are identical.
  • Closure properties: TOSSMs are closed under constant scaling, identity shifts, and unitary basis changes.These operations respectively change timescale, exponentially tilt the basis and measure, or preserve the measure under a unitary transformation.
Loading 2206.12037v2…