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

Authors' response to Reviewer BxH3.

We thank the reviewer for the valuable comments and suggestions. In the following, we address the weakness and questions raised. Weakness 1. *It would be beneficial to include practical applications that demonstrate the utility of the proposed methods.* * We appreciate this suggestion. In light of your comment, we have expanded our empirical evaluation to include the Long Range Arena (LRA) benchmark, which provides a robust set of tasks to validate our approach under practical conditions. Specifically, we trained 6-layer S4D models on this benchmark, initializing 10% of the real part to zero and configuring the timescale for this component based on Theorem 3.1. Our findings, detailed in the second table of the **General Responses**, show that our method consistently outperforms the default S4D model across all tasks, underscoring its practical effectiveness. For more detailed experimental results and analysis, we invite the reviewer to consult the **General Responses**. **We have also included the LRA benchmark results in Section 3.2, Appendix B, and Appendix B.2 of the revised manuscript to provide a more comprehensive evaluation of our approach.** Weakness 2 \& 3. A stronger connection among the theoretical components and more emphasis on distilling key messages would enhance the paper's accessibility. * Thank you for the helpful feedback. In response, we have elaborated on the interconnections among the three theoretical components in the **General Responses**. The central message is that these components collectively form a cohesive, data-dependent initialization scheme for S4D models. Specifically, given an input sequential dataset, we can estimate its autocorrelation to initialize the model timescale $\Delta$. This allows us to harness the memory function of the input-output mapping to set the real part $\Re(w)$, which helps address the curse of memory challenges associated with long sequences. Lastly, by accurately recovering the memory function from the data, we can balance the approximation-estimation tradeoff by initializing the imaginary part $\Im(w)$ based on the dominant frequencies of the memory function. This structured approach enhances the model's efficacy and aligns theoretical insights with practical application. **To make the paper more accessible, we have expanded on these connections in the Introduction and Conclusion sections of the revised manuscript.**

Authors' response to Reviewer BxH3 (continue).

Q4. *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$?* * Given a specific task with a training sequential dataset, we can estimate the maximal eigenvalue of the data autocorrelation matrix $\mathbb{E}\_x[x x^\top]$. Then we can initialize the model timescale $\Delta$ as $1/\sqrt{L \lambda\_{\max}(\mathbb{E}\_x[x x^\top])}$ to ensure the model's stability as specified in Theorem 3.1. **In light of the reviewer's feedback, we have included the above discussion in the end of Theorem 3.1 in the revised manuscript.** Q5. *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?* * Yes, the real part is trained but *without* any reparameterization method. This approach means $\Re(w)$ is likely to become positive during training. Our goal is to allow the model to learn directly from the data without introducing new variables, such as reparameterization methods, which could complicate the analysis. Otherwise, it would be unclear whether the improvements are due to the zero real part or the introduced reparameterization method. The experimental results showcased in the **General Responses** illustrate that even with a portion of the real part initialized to zero, the model successfully learns from the data, and some real parts of the state vector become non-negative after training. In light of the reviewer's question, we have added a discussion on why not constraining the parameter values during training in Section 3.2 of the revised manuscript. **For the behavior of the real part of $w$ after training, we provide a figure (Figure 8) in Appendix B of the revised manuscript, which shows the distribution of the real part of $w$ after training on the synthetic task and the sMNIST dataset.** **For the LRA benchmark, we show the ratio of non-negative real parts to the total real parts in Table 1 and 3 in the revised manuscript.** Q6. *The assumption for $c$ is real-valued. However, many SSM implementations use a complex $c$. What necessitates this assumption?* * In the current version of our work, this assumption is necessary because we utilize the Gershgorin circle theorem to bound the eigenvalues of the Gram matrix, as explained in the proof of Theorem 3.3. This theorem is applicable only when the matrix has dominant diagonal entries. If the vector $c$ is complex-valued, the resulting Gram matrix would not be diagonal-dominant, rendering the Gershgorin circle theorem ineffective. However, we believe that this assumption could be relaxed by employing alternative techniques to bound the spectrum of the Gram matrix, and we plan to explore this possibility in future work. **In light of the reviewer's feedback, we have included the above discussion in Section 3.3 of the revised manuscript.** Q7. *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?* * For a specific task, if we manage to accurately recover the memory function from the sequential data, we can then apply a Fourier transform to identify the dominant frequencies of the memory function. Given a state size $m$, we can greedily select $m$ frequencies that exhibit the largest separation distance from these dominant frequencies to initialize the imaginary part. Our theory suggests that this approach achieves an optimal balance between the tradeoff of approximation and estimation. However, practically, recovering the memory function accurately from the data is challenging, and hyperparameter tuning might be needed to find the optimal balance. **In light of the reviewer's question, we have included the above discussion in the end of Section 3.3 of the revised manuscript.**

Authors' response to Reviewer vKbW.

We thank the reviewer for the valuable comments and suggestions. In the following, we address the 2 weaknesses raised. Weakness 1. *The empirical evaluation is limited to synthetic data or somewhat simple benchmark data. More real-world datasets would certainly be appreciated.* * We appreciate the reviewer's suggestion. In light of your comment, we have conducted experiments using the Long Range Arena (LRA) benchmark to further support our theory. We specifically trained 6-layer S4D models on this benchmark, initializing 10% of the real part to zero and setting the timescale for this zero real part according to Theorem 3.1. The results, which can be found in the second table of the **General Responses**, indicate that our method consistently surpasses the default S4D model across all tasks. We encourage the reviewer to refer to the **General Responses** for more comprehensive experiment details. **We have also included the LRA benchmark results in Section 3.2, Appendix B, and Appendix B.2 of the revised manuscript to provide a more comprehensive evaluation of our approach.** Weakness 2. *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?* * We appreciate the reviewer's insightful question. In this study, we opted not to constrain the parameter values during training. This decision was made to maintain better control over the experiments and to focus on the effects of initialization schemes without introducing additional variables, such as reparameterization methods, which could complicate the analysis. For instance, we initialized the real part $\Re(w)$ to zero and allowed the model to learn from the data during the training process without incorporating any reparameterization techniques. Introducing such methods would make it unclear whether improvements were due to the initialization or the reparameterization itself. Our results on the LRA benchmark, as detailed in the tables in the **General Responses**, demonstrate that the model can effectively learn from the data with the real part initialized to zero. Even after training, some **non-negative** values remain in the state vector's real part $\Re(w)$. **In light of the reviewer's feedback, we have included the above discussion on why not constraining the parameter values during training in Section 3.2 of the revised manuscript.**

Authors' response to Reviewer YHUj.

We thank the reviewer for the valuable comments and suggestions. In the following, we address the 2 weakness raised. Weakness 1. *How are the expected results to hold in larger scale experiments? Can authors comment on that.* * We appreciate the reviewer's insightful question. In light of your comment, we conducted experiments using the Long Range Arena (LRA) benchmark to further validate our theory. Specifically, we trained 6-layer S4D models on this benchmark with 10% of the real part initially set to zero. Additionally, we initialized the timescale for this zero real part as outlined in Theorem 3.1. The results, presented in the second table of the **General Responses**, demonstrate that our method consistently outperforms the default S4D model across all tasks. For more detailed information, we invite the reviewer to consult the **General Responses**. **In light of the reviewer's feedback, we have added the LRA benchmark results to Section 3.2, Appendix B, and Appendix B.2 of the revised manuscript to provide a more comprehensive evaluation of our approach.** Weakness 2. *What do the authors think of extending these insights or draw parallel to setting(s) where the sequence length can be varied?* * That's an excellent question. Let's start by considering a scenario where the sequence length alternates between two values, $L_1$ and $L_2$. To address this, we can double the model's feature dimension and initialize it separately for the first and second halves, allowing the model to accommodate two fixed-length datasets simultaneously. If the sequence lengths vary even more, we could cluster them into several groups, increasing the model's feature dimension according to the same strategy. Ultimately, the key idea is to expand the model's feature dimension to manage varying sequence lengths, which we believe is a promising direction for future investigation. **In light of the reviewer's question, we have included a discussion on extending our insights to scenarios with varying sequence lengths in Remark 3.2 of the revised manuscript.**

General responses (continue).

| Initialize a fraction $p$ of $\Re(w)$ to be 0 | $p=0$ | $p=0.1$ | $p=0.2$ | $p=0.3$ | $p=0.4$ | $p=0.5$ | |:-----------:|:-----------:|:-----------:|:-----------:|:-----------:|:-----------:|:-----------:| |Accuracy| 84.09 (0.47) | **84.60** (0.38) | 84.23 (0.49) | 83.77 (0.46) | 83.50 (0.42) | 83.19 (0.39)| |Ratio for $\Re(w) \geq 0$ after training| 0\% | 3.90\% | 7.62\% | 10.85\% | 14.34\% | 17.25\% | From our ablation study, we found that initializing an appropriate fraction of state vectors with zero real parts enables the model to outperform the default S4D configuration. Importantly, even post-training, some non-negative real parts persist, suggesting that the model retains stability and effectively adapts to the data. Using this setup, we proceeded to train S4D models on the entire LRA benchmark. For this, we initialized 10% of the real parts to zero, without further tuning the fraction $p$. Correspondingly, we set the timescale $\Delta_0$ to a constant value $\Delta_{\min}$, adhering to the default training settings. All models were trained with a 6-layer architecture, maintaining the original S4D training conditions as specified in the work by Gu et al. (2022). | S4D | ListOps | Text | Retrieval | Image | Pathfinder | PathX | Avg | |:-----------:|:-----------:|:-----------:|:-----------:|:-----------:|:-----------:|:-----------:|:-----------:| |Baseline| 60.47 | 86.18 | 89.46 | 88.19 | 93.06 | 91.95| 84.89 | | Initialize 10\% of $\Re(w)$ to be $0$ | **61.44** | **88.05** | **90.73** | **89.11** | **95.58** | **97.55** | **87.08**| |Ratio for $\Re(w) \geq 0$ after training| 1.29\% | 1.99\% | 2.58\% | 4.31\% | 4.31\% | 4.04\% | 3.09\% | Given time constraints, we conducted the experiments using a single random seed. The experiment results demonstrate that initializing a portion of the real parts to zero consistently improves performance across all six tasks when compared to the default S4D model. We hope these recent findings effectively address the reviewers' concerns and would like to express our gratitude to the reviewers for their time and thoughtful consideration. **In light of the reviewers' suggestion to include more real-world datasets, we have expanded our empirical evaluation to include the Long Range Arena (LRA) benchmark in Section 3.2, Appendix B, and Appendix B.2 of the revised manuscript.**