How to Train Your HiPPO: State Space Models with Generalized Orthogonal Basis Projections
It addresses a foundational gap in understanding state space models for long sequences in machine learning, leading to improved performance on benchmarks.
The paper tackles the lack of theoretical understanding of how the Structured State Space sequence model (S4) captures long-range dependencies by generalizing the HiPPO framework, providing a mathematical interpretation and deriving new variants. This improves S4's performance to 86% on the Long Range Arena benchmark, with 96% on the Path-X task.
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.