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

Authors: Fusheng Liu, Qianxiao Li
Primary area: other topics in machine learning (i.e., none of the above)
Keywords: State Space Model; Optimization;

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.

Decision: accept (poster)
Reviewer scores: 6, 8, 6

Reviews and discussion:
<thread>
creation date:
1729469942179

thread_type:
official_review

<thread_object>
creator: reviewer_bxh3

<content>
summary:
This paper provides a theoretical analysis of a state-space model from three perspectives: discretisation size ($\Delta$), the real components of the diagonal weights, and their imaginary counterparts. It discusses the scaling laws of $\Delta$, highlights how zero real parts contribute to long memory, and introduces an approximation-estimation tradeoff influenced by the imaginary components.

strengths:
**(I)** The analyses presented are rigorous, with clear explanations connecting the results to different initialization strategies, which is crucial for understanding SSMs.

**(II)** The paper is well-structured, with illustrative examples that help to clarify the theoretical concepts.

weaknesses:
**(I)** The paper primarily focuses on theoretical aspects, with limited empirical evaluation beyond basic synthetic examples. It would be beneficial to include practical applications that demonstrate the utility of the proposed methods.

**(II)** The explanation of the three theoretical components sounds a bit separated. While they all fall under the umbrella of initializations, a stronger connection would give the paper a more cohesive narrative.

**(III)** Although the theoretical results are clearly presented, they are dense. More emphasis on distilling key messages and offering practical guidelines would enhance the paper's accessibility.

questions:
**(I)** The analysis focuses on the single-input single-output case. How would these results extend to the multiple-input multiple-output (S5) case? High-level insights would suffice.

**(II)** The paper assumes ZOH discretisation. How would the results differ with bilinear discretisation?

**(III)** Several questions regarding Theorem 4.1:

   (a) The upper bound applies to the expected output. Is there a corresponding probabilistic bound?

   (b) Can the upper bound be shown to be tight?

   (c) How does the real part of $w_j$ affect this bound? Specifically, how would increasing or decreasing it impact the upper bound?

   (d) The analysis focuses on the final output. How would the results change if the sum of outputs $\sum_{\ell = 1}^L y_\ell^2$  were considered instead? This pooling is adopted by many SSMs.

**(IV)** The paper explores various settings for how $\Delta$ should scale with $L$. Given a specific task, what method would you use to determine the appropriate $\Delta$?

**(V)** When the real part of $w$ is set to zero at initialization, is it still trained in practice? If so, do you enforce non-positivity? In this case, how do you handle reparameterization, especially since you can’t use the logarithm of it? Could you also show the behavior of the real part of $w$ after training? If it becomes negative, what is the benefit of zero initialization?

**(VI)** On line 376, the assumption is made that $c$ is real-valued. However, many SSM implementations use a complex $c$. What necessitates this assumption?

**(VII)** Given a specific task, how would you determine the appropriate balance between approximation and estimation tradeoff? Would this be driven by hyperparameter tuning or derived from theoretical insights?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730473992913

thread_type:
official_review

<thread_object>
creator: reviewer_vkbw

<content>
summary:
An analysis of the effects of some parameters of the so-called SSM layer is provided. Specifically, the authors analyze 1) the relation between the timescale $\Delta$ and the sequence length $L$, 2) the effect of the (zero) real part of the time evolution matrix, and 3) the effect of the imaginary part of the time evolution matrix, both theoretically and empirically.

strengths:
- The claims are clear, and most of them are supported both theoretically and empirically.
- The topic is indeed important as the SSM-based models are increasingly used in various problems.
- Although the analyses may not be surprising per se in terms of dynamical systems theory, their implication especially in using SSM architecture sounds new and useful.

weaknesses:
The empirical evaluation is limited to synthetic data or somewhat simple benchmark data. This fact does not much diminish the value of the paper, but reporting the applicability of the proposed theory to more real-world datasets would certainly be appreciated.

A thing remained a bit unclear to me is that how the discussions are necessarily relevant only to the initialization. Doesn't it make sense to constrain the parameter values (e.g., $\Re(W)$) as guided in the given theory, not only at the initialization but also during training iterations? As the paper's title emphasizes the role of *initialization*, a bit more discussion in this regard might be helpful.

questions:
No questions that may change my evaluation.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730663813413

thread_type:
official_review

<thread_object>
creator: reviewer_yhuj

<content>
summary:
The paper studies the initialization strategy used in SSM from various point of view, namely the time scale, the real and imaginary part of the state matrix. It demonstrates the dependency of the SSM timescale on sequence length based on sequence autocorrelation. Further it is shown that having a zero real part for the eigenvalues of the SSM state matrix mitigates the curse of memory. Finally, the paper demonstrates that the imaginary part of the eigenvalues of the SSM state matrix determine the conditioning of SSM optimization problems, and present a approximation-estimation tradeoff when training SSMs with a specific class of target functions.

strengths:
1. The paper has presented some very novel an interesting insights into the initialization scheme for SSM.
2. The analysis is well grounded and thorough.
3. The experiments are well though out, extensive and results are convincing of the claims made in the paper.

weaknesses:
While the experimental results are convincing on the datasets considered in the paper, how are the expected to hold in larger scale experiments? Can authors comment on that.

What do the authors think of extending these insights or draw parallel to setting(s) where the sequence length can be varied?

questions:
Please see above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732352444326

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General responses.

comment:
We thank each reviewer for their careful reading of our work and their thoughtful feedback.
 We have revised the paper with the main changes in *red* according to the comments and suggestions.
In this unified response, first we would like to depict the whole picture and illustrate the connections among each theory in this paper.
Then we address a common concern shared by all reviewers: the absence of practical applications involving more real-world datasets.

The main goal of this paper is to provide a theoretical understanding on the effects of three specific S4D hyperparameters: the model timescale $\Delta$, the real part of the state vector $\Re(w)$, and the imaginary part of the state vector $\Im(w)$.
These components are intricately linked as a **data-dependent** initialization scheme for S4D models.
**First**, for any given sequential dataset, we can estimate its autocorrelation. 
Using this information, we can apply Theorem 3.1 to initialize $\Delta$, taking into account both data autocorrelation and sequence length.
**Second**, if the true input-output mapping is represented by an underlying linear functional, often referred to as a memory function, that exhibits a long memory pattern, our second theory, detailed in Section 3.2, suggests that initializing with a zero real part $\Re(w)$ can help mitigate the challenges posed by long sequences.
**Finally**, the third theory introduced in Section 3.3 discusses an approximation-estimation tradeoff that arises when the true memory function $\rho^*$ features closely spaced frequencies. 
If we can accurately recover $\rho^*$ from the sequential data, we can then initialize the imaginary part $\Im(w)$ based on the dominant frequencies of $\rho^*$, thereby finding an optimal balance informed by theoretical insights.
**To make the paper more accessible, we have elaborated on these connections in the Introduction and Conclusion sections of the revised manuscript.**

Building on the feedback from reviewers, we are expanding our experiments to include larger scale experiments. 
Specifically, we are applying our first two theories to the Long Range Arena (LRA) benchmark, which features six diverse tasks ranging from text to image processing. 
However, the third theory cannot be directly validated using the LRA benchmark, as it necessitates knowledge of the exact memory function for the input-output mapping—a challenge in practical scenarios.

To proceed, we begin with ablation studies on the gray-sCIFAR dataset to evaluate the benefit of initializing the real part to zero in multi-layer S4D models, while making minimal modifications. Based on our theory, zeroing the real part can alleviate the *curse of memory* in scenarios where the memory function exhibits a long memory pattern. 
Given that we lack precise knowledge of the memory function for the gray-sCIFAR dataset, we opt to initialize a fraction of the real part as zero and compare this setup to the default S4D model.
In particular, for each layer of the S4D model, which has a feature dimension of $d$ and a state size of $m$, there are $d$ state vectors $w \in \mathbb{C}^m$. 
At initialization, a fraction $p \in [0, 1]$ of these state vectors is randomly selected to have their real parts set to zero. 
For these selected vectors, their corresponding timescale $\Delta_0 \in \mathbb{R}^{p \cdot d}$ is initialized to a constant, as suggested by our first theory. 
From Figure 1 (Right), we observe that the maximal eigenvalue of $\mathbb{E}[xx^\top]$ on the sCIFAR dataset scales with $L$. 
Therefore, in accordance with Theorem 3.1, we initialize $\Delta_0$ as a constant $1/L$. 
Other hyperparameters are maintained as per the default S4D model settings.
When $p = 0$, the training proceeds following the baseline setup. 
As $p$ increases, the model starts with more zero real parts. 
To ensure credible results, we **exclude** any reparameterization strategies on the zero real part, allowing the model to adapt from the data during training. 
This approach isolates the impact of the zero real part on performance without confounding variables introduced by reparameterization.
In our ablation studies, a 4-layer S4D model is used, with $p$ varied across $[0, 0.1, 0.2, 0.3, 0.4, 0.5]$. 
We present both the test accuracy and the ratio of **non-negative** real part parameters to the total $Ldm$ real part parameters upon completion of training in the subsequent table.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper 'Autocorrelation Matters: Understanding the Role of Initialization Schemes for State Space Models' was reviewed by 3 reviewers who gave it an average score of 6.67 (final scores: 6+6+8). The reviewers found this work interesting, clear, and the claims supported both theoretically and empirically. All reviewers argue for accepting the paper.