Autocorrelation Matters: Understanding the Role of Initialization Schemes for State Space Models

TL;DR

This paper investigates how different initialization schemes for state space models affect their performance, emphasizing the importance of input sequence autocorrelation and eigenvalue properties for stability and optimization.

Contribution

It introduces a new perspective on SSM initialization by analyzing autocorrelation effects and eigenvalue roles, providing guidelines for improved stability and training.

Findings

Eigenvalues' real part influences memory retention and stability.

Eigenvalues' imaginary part affects optimization conditioning.

Proper timescale selection mitigates the curse of memory.

Abstract

Current methods for initializing state space model (SSM) parameters primarily rely on the HiPPO framework \citep{gu2023how}, which is based on online function approximation with the SSM kernel basis. However, the HiPPO framework does not explicitly account for the effects of the temporal structures of input sequences on the optimization of SSMs. In this paper, we take a further step to investigate the roles of SSM initialization schemes by considering the autocorrelation of input sequences. Specifically, we: (1) rigorously characterize the dependency of the SSM timescale on sequence length based on sequence autocorrelation; (2) find that with a proper timescale, allowing a zero real part for the eigenvalues of the SSM state matrix mitigates the curse of memory while still maintaining stability at initialization; (3) show that the imaginary part of the eigenvalues of the SSM state matrix determines the conditioning of SSM optimization problems, and uncover an approximation-estimation tradeoff when training SSMs with a specific class of target functions.