How to Train Your HiPPO: State Space Models with Generalized Orthogonal Basis Projections
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.