UnHiPPO: Uncertainty-aware Initialization for State Space Models

TL;DR

This paper introduces UnHiPPO, an uncertainty-aware initialization method for state space models that accounts for measurement noise, improving robustness in noisy sequence data scenarios.

Contribution

It extends the HiPPO framework to noisy data by deriving a new initialization that infers the latent system's posterior without additional runtime overhead.

Findings

Improves noise robustness during training and inference.

Provides a practical implementation available online.

Reformulates HiPPO as a stochastic control problem.

Abstract

State space models are emerging as a dominant model class for sequence problems with many relying on the HiPPO framework to initialize their dynamics. However, HiPPO fundamentally assumes data to be noise-free; an assumption often violated in practice. We extend the HiPPO theory with measurement noise and derive an uncertainty-aware initialization for state space model dynamics. In our analysis, we interpret HiPPO as a linear stochastic control problem where the data enters as a noise-free control signal. We then reformulate the problem so that the data become noisy outputs of a latent system and arrive at an alternative dynamics initialization that infers the posterior of this latent system from the data without increasing runtime. Our experiments show that our initialization improves the resistance of state-space models to noise both at training and inference time. Find our implementation at https://cs.cit.tum.de/daml/unhippo.