State Space Models advanced 7 min read 12 flashcards

HiPPO and Structured State Initialisation

Why a randomly initialised state space layer fails on long sequences, what optimal polynomial projection of history gives you instead, and how the structure that makes it work also makes it computable.

The first linear state space models did not work. The architecture was right, the duality was right, and on the Long Range Arena benchmark they performed around chance. The fix was not a change to the architecture but to the initialisation of a single matrix, and the resulting jump in performance is one of the clearer demonstrations that initialisation can be the entire content of a method.

The question HiPPO answers

Suppose you must compress an arbitrarily long input history into a fixed \(N\)-dimensional state, updated online, and you want the state at any moment to be the best possible summary of everything seen so far. "Best" needs a definition: project the history onto the first \(N\) orthogonal polynomials, under a measure that says how much each past time step matters.

HiPPO (Gu et al., 2020, arXiv:2008.07669) derives, in closed form, the matrix \(A\) whose linear recurrence maintains exactly those projection coefficients as new input arrives. Different measures give different matrices. The scaled Legendre measure weights all history uniformly up to the present and yields HiPPO-LegS, the variant most used; an exponentially decaying measure yields a matrix that behaves like a leaky integrator with a fixed timescale.

The result is that the state is not an arbitrary learned summary. It is the coefficient vector of an optimal polynomial approximation to the entire input history, and the recurrence updates it correctly as the window grows.

Why random initialisation fails

A random \(A\) gives a kernel \(C\bar{A}^j\bar{B}\) whose magnitude decays roughly geometrically with \(j\) at a rate set by the spectral radius. Any such kernel has an effective memory horizon of at most a few hundred steps, and gradients through the convolution beyond that horizon are numerically zero. The model cannot learn long dependencies because the initialisation places it in a region of parameter space where long-range gradients do not exist to guide it out.

The HiPPO matrix instead produces a kernel with structured oscillatory components at many timescales simultaneously, so signal from thousands of steps back arrives at the output with non-negligible magnitude from the very first gradient step. Training then refines a model that already sees far, rather than trying to teach one that sees nothing.

Structure as a computational gift

The HiPPO-LegS matrix has a specific algebraic form: normal plus low rank. That property is what makes the kernel computable at all. Decomposing \(A\) this way lets the required powers be evaluated through a Cauchy kernel in \(O((N+L)\log^2(N+L))\) rather than \(O(N^2 L)\), which is the technical core of S4.

Later work found that most of the benefit survives with a purely diagonal \(A\), provided the eigenvalue initialisation approximates the HiPPO spectrum. This is the S4D result, and it matters practically: a diagonal state matrix is a few lines of code, runs in complex arithmetic, and drops the implementation burden that kept S4 confined to specialist repositories. Mamba and its successors use diagonal state matrices for this reason.

When it breaks

Initialisation is not a constraint. \(A\) is trained, so nothing keeps it in the HiPPO family after the first optimiser step. Training can and does move eigenvalues, sometimes toward instability. Parameterising eigenvalues so that stability holds by construction, rather than hoping training respects it, is standard in robust implementations.

The optimality claim is about a specific problem. HiPPO optimally reconstructs the input signal under a chosen measure. Optimal reconstruction is not the same as optimal representation for a downstream task; a summary that discards irrelevant detail may be more useful than one that preserves everything. The theory explains why the initialisation is good, not why it is best.

Timescale initialisation matters as much as \(A\). The step size \(\Delta\) sets each channel's timescale, and the usual practice initialises \(\log \Delta\) uniformly over a range spanning several orders of magnitude so the layer covers short and long horizons simultaneously. Initialising all channels at one timescale reproduces the original failure even with a correct HiPPO matrix.

Diagonal is not free everywhere. The equivalence between diagonal-plus-low-rank and purely diagonal parameterisations depends on the initialisation approximating the right spectrum. Diagonal models initialised carelessly regress toward random-initialisation behaviour, which is easy to miss because the model still trains and still produces reasonable short-range results.

Check yourself

12 flashcards for this concept

Click a card to reveal the answer.

Drill the whole track