Paper: Robustifying State-space Models for Long Sequences via Approximate Diagonalization (ICLR-2024)

Reviewer comment (Experiments):
I am not convinced by the claim the S4-PTD model outperforms the S4D models on LRA. The LRU paper (https://arxiv.org/abs/2303.06349) reports results for S4D that are much better than the original reported S4D paper results. In addition, the appendix of the current paper under review states that mild hyperparameter tuning was performed for the S4-PTD models. Would mild hyperparameter tuning improve the S4D results as well? To be clear, I am not claiming that S4-PTD must outperform S4D on all tasks to be a valuable contribution, simply that I think the claim about improving performance is too strong, and I would expect to see more results in a fair comparison to support this claim.

Author response:
We thank the reviewer for pointing the new results out. First, we clarify that by saying "mild tuning," we do not mean a laborious grid search. Instead, what we really want to say is that we did not take the exact hyperparameters from S4D/S5 because the models are initialized differently so the S4-PTD/S5-PTD models are expected to have a different set of good hyperparameters. The S4D/S5 models must have also been tuned and we did not encounter a better set of hyperparameters that can significantly improve the S4D/S5 performance.

The LRU paper has performed a more careful tuning than what we have done. In addition, the result reported for the Image task is based on colorful sCIFAR instead of the grayscaled one as used in the original S4D/S5 paper and ours. In the appendix of the LRU paper, the authors also report some architectural differences, e.g., they used a different activation function for the PathX task and it seems that the number of model parameters is also different from what is used in the original S4D/S5 paper and our manuscript. We will reach out to the authors of the LRU paper to get more information on the details of the S4D model they have tuned. We will make sure that we have an unbiased claim about the performance of our models in the camera-ready version should the paper be accepted.

Reviewer comment (Experiments):
More experiment results on a broader range of robustness tasks would help strengthen the paper. At a minimum, perhaps modified versions of the speech and BIDMC tasks from previous papers in this line of work could have allowed for further evaluation of this method on some form of real world data.

Author response:
We have done some experiments on the Speech Command task (see Table 6). We are working on the BIDMC task and will be sure to include the result in the camera-ready version should the manuscript be accepted. We remark that due to the limited time, we could not tune the models for Speech Command. The hyperparameters are exactly copied from the S4D/S5 models. Hence, the results reported may not be optimal.

Reviewer comment (Experiments):
There seems to be a missing piece related to the training dynamics that I think would strengthen the paper. The theory is related to the dynamics matrix $A$ at initialization, but the $A$ matrices are trained during the experiments. How do the trajectories of the eigenvalues differ during training such that the PTD models are more robust during the perturbation experiments? This should be relatively to simple to track during training.

Author response:
We have added an empirical evaluation of the evolution of the transfer function over training (see Figure 13). We find that the instabilities in the S4D model are preserved throughout training while the S4-PTD model does not develop any instability. We hope this adds a new perspective to our paper and makes it more complete.

Reviewer comment (Experiments):
The paper doesn't really seem to support the claim that the S5 models suffer from the same perturbation issue. I believe this is likely the case, but no empirical results are presented that support this. Can an example of S5 and S5-PTD be included in the perturbation experiments also?

Author response:
It is very likely that the S5 initialization suffers from the same issue, as suggested by Figure 4. We have implemented an experiment to test the robustness of the S5/S5-PTD models. We will update our manuscript when the results come in.

Reviewer comment (Writing):
**W.1.: The core method**: Please can the authors clarify what the method actually entails. As far as I can tell, the “method” essentially specifies an initialization for S4D, given in (10) followed by diagonalization. This is used instead of the diagonalization of the Normal component of the HiPPO matrix in S4D. If this is true, this should be absolutely, concretely clearly stated somewhere. I think right now it is tied up in the introduction with spectral analysis, PTD, backwards/forwards approximations etc.

Author response:
**Response:** Yes, this is indeed a main practical contribution of the paper. We have highlighted this in several places:
* In the last paragraph of the introduction (right above the contributions), we have explained the S4-PTD and S5-PTD models in a succinct way and pointed out their advantage in terms of robustness.

* We have also highlighted this method in Contribution \#4.

* In section 4, we explained the model concretely in (9) and the text around it. In addition, at the end of the first paragraph, we stated how the S4-PTD and S5-PTD models are different from the S4D and S5 models.

* Finally, we also highlighted our contribution in our conclusion.

Reviewer comment (Presentation):
**W.2.: Choice of what the paper emphasizes**: The authors have clearly done _a lot_ of theoretical analysis. However, I think the net result is an incredibly simple result – adding a random perturbation makes models more robust. I really think this result, and any usable empirical analysis of this result, should be highlighted. Some of the theoretical points should be relegated to the appendix. If the results are to be believed, then the banner result should be “ignore DPLR/diagonal state spaces and use random perturbations”. This core message is not highlighted.

Author response:
**Response:** We agree that more empirical analysis of our theory and model can be informative. In addition to highlighting the main practical message of the paper, as explained above, we have made the following changes:

* We have put Lemma 1 and its discussion into the appendices. In addition, we have also moved the experiment on the interpolation/extrapolation into the appendix; instead, we now use Figure 4 to more efficiently deliver the same message.

* We have added Figure 2 (previously as Figure 8-9), in which the diagonal and the DPLR systems are simulated given both a smooth input and a non-smooth input. This helps the readers to visualize and better understand Theorem 1.

* We have also added Figure 4, which is a new experiment added to the manuscript. The main message is that we have verified the spikes in the transfer function of the S4D initialization impair its robustness. This is also related to our answer to Q1 raised by the reviewer (see below).

Reviewer comment (Presentation):
Similarly, there are some nice experiments in the supplement demonstrating the pathology experimentally, e.g. fig 7-10. These should be brought up to the main, as they help flesh out the arguments far more than a detailed introduction to the intricacies of the conjugation of HiPPO or some of the theorems. I think this will help improve the reach and impact of the paper.

Author response:
First, we have put this experiment into the appendix because we believe our new result (see Figure 4) helps better explain our main message. It is not coincidental that the S4D model is much worse than the S4 and S4-PTD models in this case. This can be understood from Figure 4. The main reason is that the S4D model is not stable near certain Fourier modes. Hence, it is hopeless to do interpolation or extrapolation near those frequencies if they are not included in the training dataset. We believe the fact that the S4 model and the S4-PTD model behave similarly is not merely coincidental either. If one looks into the low-frequency domains in Figure 9 (previously as Figure 2), then the S4 and the S4-PTD models behave very differently. We believe the main reason that the two models behave similarly in the high-frequency domain is that we observe, empirically, that most changes in the LTI systems during training only happen in the low-frequency domain. (See Figure 13 and our answer to Q1.) We think this explains why the two models predict similarly in the high-frequency domains but not the low-frequency ones.

Reviewer comment (Evaluation):
**W.3.: Figure 2 results**: I find the results in Figure 2 very hard to parse out. The core point (to my understanding) is that S4-PTD is a faithful approximation of S4. However, I do not necessarily see why S4D’s predictions _should_ be reminiscent of S4 – it’s a different model! – or that it should be worse in the way predicted. I think it is coincidental that S4D is so much worse than S4(-PTD) on this task. I ask the authors to clarify the intuition and results of this experiment.

Author response:
First, we have put this experiment into the appendix because we believe our new result (see Figure 4) helps better explain our main message. It is not coincidental that the S4D model is much worse than the S4 and S4-PTD models in this case. This can be understood from Figure 4. The main reason is that the S4D model is not stable near certain Fourier modes. Hence, it is hopeless to do interpolation or extrapolation near those frequencies if they are not included in the training dataset. We believe the fact that the S4 model and the S4-PTD model behave similarly is not merely coincidental either. If one looks into the low-frequency domains in Figure 9 (previously as Figure 2), then the S4 and the S4-PTD models behave very differently. We believe the main reason that the two models behave similarly in the high-frequency domain is that we observe, empirically, that most changes in the LTI systems during training only happen in the low-frequency domain. (See Figure 13 and our answer to Q1.) We think this explains why the two models predict similarly in the high-frequency domains but not the low-frequency ones.

Reviewer comment (Reproducibility):
I think it is also somewhat coincidental that S4-PTD’s predictions look like S4 – surely different random seeds should result in different predictions?

Author response:
First, we have put this experiment into the appendix because we believe our new result (see Figure 4) helps better explain our main message. It is not coincidental that the S4D model is much worse than the S4 and S4-PTD models in this case. This can be understood from Figure 4. The main reason is that the S4D model is not stable near certain Fourier modes. Hence, it is hopeless to do interpolation or extrapolation near those frequencies if they are not included in the training dataset. We believe the fact that the S4 model and the S4-PTD model behave similarly is not merely coincidental either. If one looks into the low-frequency domains in Figure 9 (previously as Figure 2), then the S4 and the S4-PTD models behave very differently. We believe the main reason that the two models behave similarly in the high-frequency domain is that we observe, empirically, that most changes in the LTI systems during training only happen in the low-frequency domain. (See Figure 13 and our answer to Q1.) We think this explains why the two models predict similarly in the high-frequency domains but not the low-frequency ones.

Reviewer comment (Writing):
**W.4.: Missing definitions and exposition**: Lots of terms are basically undefined, e.g.:  
- The introduction of “backward stability” is deferred to the appendix, but is really the core of the method. There is no discussion in the main really of why forward/backwards are even relevant, or why existing methods fall down here.  
- Pseudospectral theory is undefined, nor why adding a random perturbation ($E$) helps.  
- Why spectral information with a poor condition number is “meaningless”.

Author response:
**Response:** We have added some definition/explanation to clarify these terminologies/statements:

* We have added a paragraph in section 4 (see the second paragraph that starts with "Consider the process of diagonalizing...") to define the meaning of the backward and forward errors in the context of SSMs. In the same paragraph, we have also explained from a high-level perspective why it suffices to consider the backward error and not the forward error. We believe this is a good place to introduce the two types of errors in detail because they are not used until Theorem 3.

* We have added a footnote (see footnote 2) to explain what the pseudospectral theory is mainly about. The pseudospectral theory studies a very broad class of questions and is not just a single theorem. For example, Theorem 4 about the eigenvector condition number under perturbation is a result that falls under the pseudospectral theory and explains why adding the perturbation matrix $\mathbf{E}$ helps.

* The true spectral information of a highly non-normal matrix is meaningless because it is very unstable. Hence, using the spectral information, one would suffer from a very large variance (in terms of the bias-variance tradeoff). The other way to view it is that when $\mathbf{A}$ is highly non-normal, the eigenvector matrix $\mathbf{V}$ is very ill-conditioned. Hence, for example, one cannot project a vector $\mathbf{v}$ onto the eigenbasis and use the spectral information of $\mathbf{A}$ to study $f(\mathbf{A})\mathbf{v}$, where $f$ is an analytic function. We have added a brief version of this discussion as a footnote (see footnote 1).

Reviewer comment (Presentation):
This lack of definitions makes the contribution a little less self-contained. The reviewer would rather drop some of the theoretical analysis from the main to the supplement and bring up more of the core definitions. This would dramatically improve the impact of the work.

Author response:
**Response:** We have added some definition/explanation to clarify these terminologies/statements:

* We have added a paragraph in section 4 (see the second paragraph that starts with "Consider the process of diagonalizing...") to define the meaning of the backward and forward errors in the context of SSMs. In the same paragraph, we have also explained from a high-level perspective why it suffices to consider the backward error and not the forward error. We believe this is a good place to introduce the two types of errors in detail because they are not used until Theorem 3.

* We have added a footnote (see footnote 2) to explain what the pseudospectral theory is mainly about. The pseudospectral theory studies a very broad class of questions and is not just a single theorem. For example, Theorem 4 about the eigenvector condition number under perturbation is a result that falls under the pseudospectral theory and explains why adding the perturbation matrix $\mathbf{E}$ helps.

* The true spectral information of a highly non-normal matrix is meaningless because it is very unstable. Hence, using the spectral information, one would suffer from a very large variance (in terms of the bias-variance tradeoff). The other way to view it is that when $\mathbf{A}$ is highly non-normal, the eigenvector matrix $\mathbf{V}$ is very ill-conditioned. Hence, for example, one cannot project a vector $\mathbf{v}$ onto the eigenbasis and use the spectral information of $\mathbf{A}$ to study $f(\mathbf{A})\mathbf{v}$, where $f$ is an analytic function. We have added a brief version of this discussion as a footnote (see footnote 1).

Reviewer comment (Presentation):
**W.5.: No conclusion or discussion**: Not really a weakness, but there is no conclusion, discussion or outlined future work. It would be nice to see some critical self-reflection on the work. It would also help reinforce the core messages the authors are trying to convey in the paper.

Author response:
**Response:** We have added a conclusion to summarize the manuscript and suggest future research avenues.

Reviewer comment (Writing):
(Minor weakness) “section 3.4” should be capitalized.

Author response:
**Response:** We have made the following changes:

* In the ICLR style guideline, sections are not capitalized. We follow this convention in our manuscript.

* We have improved our explanation in the caption of the figure. Note that we have moved the figure into the appendix (see Figure 9).

* We have corrected all misuses of \citet and \citep.

* We have made the names of the authors consistent in our bibliography.

Reviewer comment (Presentation):
(Minor weakness) The explanation of Figure 2 on Page 6 is very poor.

Author response:
**Response:** We have made the following changes:

* In the ICLR style guideline, sections are not capitalized. We follow this convention in our manuscript.

* We have improved our explanation in the caption of the figure. Note that we have moved the figure into the appendix (see Figure 9).

* We have corrected all misuses of \citet and \citep.

* We have made the names of the authors consistent in our bibliography.

Reviewer comment (Writing):
(Minor weakness) Use \citet and \citep where appropriate.

Author response:
**Response:** We have made the following changes:

* In the ICLR style guideline, sections are not capitalized. We follow this convention in our manuscript.

* We have improved our explanation in the caption of the figure. Note that we have moved the figure into the appendix (see Figure 9).

* We have corrected all misuses of \citet and \citep.

* We have made the names of the authors consistent in our bibliography.

Reviewer comment (Presentation):
(Minor weakness) Author style in the bibliography is inconsistent.

Author response:
**Response:** We have made the following changes:

* In the ICLR style guideline, sections are not capitalized. We follow this convention in our manuscript.

* We have improved our explanation in the caption of the figure. Note that we have moved the figure into the appendix (see Figure 9).

* We have corrected all misuses of \citet and \citep.

* We have made the names of the authors consistent in our bibliography.

Reviewer comment (Experiments):
**Q.1.** This is a very general question, but something that the reviewer cannot wrap their head around. Throughout the text this is an analysis of (e.g.) HiPPO initialization, but the results presented are for trained models. The reviewer is surprised that the pathology predicted by the initialization is preserved through training quite so neatly. For instance, in Figure 3, how are the frequencies that the model should be perturbed at obtained? Is the trained model inspected, or the a-priori specified HiPPO matrix? Each channel in the S4D model should have different frequencies, and so are all of them being corrupted, or just one?

Author response:
**Response:** We thank the reviewer for this very insightful question. This question has inspired some discussion that we did not include in our manuscript initially. We have added a more detailed explanation of the experiment setup in the main text and Appendix J.2 and a fair amount of empirical discussion (see Appendix K.2) which we believe will clarify and strengthen the paper. Hopefully, the following points we make (partially) answer the reviewer's question:

* In the experiment shown in Figure 3(a/b), we fixed the state matrix $\mathbf{A}$. The main reason for fixing the matrix is that, as explained later on, training slightly shifts the locations of the spikes in the transfer function of the S4D LTI systems, which requires us to perform a case study after each epoch to find a frequency $s$ at which the S4D model is not stable. We have made this experimental setup clear in our manuscript to avoid confusion.

* We have added a new appendix (See Appendix K). In Figure 4 and Appendix K.2, we show an experiment where we contaminate the test dataset with Fourier noises at varying frequencies $s$. This is still done with a fixed state matrix $\mathbf{A}$. In Figure 4, we plot the transfer function of the S4D initialization on top of the test accuracy, which reveals a strong correlation between the "spikes" in the transfer function and the effect of the perturbation at the corresponding frequencies. This connects our theoretical study of the S4D transfer functions to the discussion of the robustness of the model.

* Now, we consider the case where we indeed train the state matrix $\mathbf{A}$. In this new experiment, we train our SSMs as everyone does in practice (see Appendix K.2 for more details). Given our observation in Figure 4, we use the transfer function as a protocol to evaluate the robustness of our model. We watch the evolution of the transfer function (randomly selected from many transfer functions in the SSM) throughout training. The results are shown in Figure 13 and here is a summary of what is going on:

   * The spikes in the transfer function of the S4D model have slightly decayed but never vanished. We believe this slight decay is possibly due to the weight decay rate we set for the LTI matrices.

   * The locations of the spikes shift very slightly. This is the main reason why when conducting the experiment in Figure 3(a/b), we fixed the matrix $\mathbf{A}$: otherwise, we would have to eyeball the location of the spike after each epoch.

   * The transfer function of the S4-PTD model never develops any spike or instability throughout training.

In conclusion, this shows training does not completely solve the robustness issue of the S4D model. Neither does it impair the robustness of the S4-PTD model.

Since this paper is mainly about analyzing the initialization, we chose to put all these discussions into the appendices and leave them as a remark in the main manuscript. Understanding the training dynamics of the LTI systems is certainly an interesting future research avenue. We do remark that in our conclusion. We are happy to answer any follow-up questions that the reviewer may have and also include a more thorough empirical evaluation along this line in our camera-ready version should this manuscript be accepted.

Reviewer comment (Theory):
Follow-up to Q.1: Why does the approximate diagonalization on initialization lead to more robust trained models? Presumably $E$ isn’t learned? Once the model is diagonalized, the initialization is essentially thrown away. As a result, there is a real disconnect in Figure 3. It looks like S4-PTD might be _slightly_ worse than S4, but the training curves are ostensibly identical. But the robustness to perturbations is clear. Why do these pathological frequencies survive training in S4?

Author response:
**Response:** We thank the reviewer for this very insightful question. This question has inspired some discussion that we did not include in our manuscript initially. We have added a more detailed explanation of the experiment setup in the main text and Appendix J.2 and a fair amount of empirical discussion (see Appendix K.2) which we believe will clarify and strengthen the paper. Hopefully, the following points we make (partially) answer the reviewer's question:

* In the experiment shown in Figure 3(a/b), we fixed the state matrix $\mathbf{A}$. The main reason for fixing the matrix is that, as explained later on, training slightly shifts the locations of the spikes in the transfer function of the S4D LTI systems, which requires us to perform a case study after each epoch to find a frequency $s$ at which the S4D model is not stable. We have made this experimental setup clear in our manuscript to avoid confusion.

* We have added a new appendix (See Appendix K). In Figure 4 and Appendix K.2, we show an experiment where we contaminate the test dataset with Fourier noises at varying frequencies $s$. This is still done with a fixed state matrix $\mathbf{A}$. In Figure 4, we plot the transfer function of the S4D initialization on top of the test accuracy, which reveals a strong correlation between the "spikes" in the transfer function and the effect of the perturbation at the corresponding frequencies. This connects our theoretical study of the S4D transfer functions to the discussion of the robustness of the model.

* Now, we consider the case where we indeed train the state matrix $\mathbf{A}$. In this new experiment, we train our SSMs as everyone does in practice (see Appendix K.2 for more details). Given our observation in Figure 4, we use the transfer function as a protocol to evaluate the robustness of our model. We watch the evolution of the transfer function (randomly selected from many transfer functions in the SSM) throughout training. The results are shown in Figure 13 and here is a summary of what is going on:

   * The spikes in the transfer function of the S4D model have slightly decayed but never vanished. We believe this slight decay is possibly due to the weight decay rate we set for the LTI matrices.

   * The locations of the spikes shift very slightly. This is the main reason why when conducting the experiment in Figure 3(a/b), we fixed the matrix $\mathbf{A}$: otherwise, we would have to eyeball the location of the spike after each epoch.

   * The transfer function of the S4-PTD model never develops any spike or instability throughout training.

In conclusion, this shows training does not completely solve the robustness issue of the S4D model. Neither does it impair the robustness of the S4-PTD model.

Since this paper is mainly about analyzing the initialization, we chose to put all these discussions into the appendices and leave them as a remark in the main manuscript. Understanding the training dynamics of the LTI systems is certainly an interesting future research avenue. We do remark that in our conclusion. We are happy to answer any follow-up questions that the reviewer may have and also include a more thorough empirical evaluation along this line in our camera-ready version should this manuscript be accepted.

Reviewer comment (Presentation):
The reviewer invites the authors to comment on Q.1 and notes that these thoughts, while a little jumbled, may be emblematic of a shortcoming in the presentation or cross-linking of the ideas. They offer to provide further clarification if needed.

Author response:
**Response:** We thank the reviewer for this very insightful question. This question has inspired some discussion that we did not include in our manuscript initially. We have added a more detailed explanation of the experiment setup in the main text and Appendix J.2 and a fair amount of empirical discussion (see Appendix K.2) which we believe will clarify and strengthen the paper. Hopefully, the following points we make (partially) answer the reviewer's question:

* In the experiment shown in Figure 3(a/b), we fixed the state matrix $\mathbf{A}$. The main reason for fixing the matrix is that, as explained later on, training slightly shifts the locations of the spikes in the transfer function of the S4D LTI systems, which requires us to perform a case study after each epoch to find a frequency $s$ at which the S4D model is not stable. We have made this experimental setup clear in our manuscript to avoid confusion.

* We have added a new appendix (See Appendix K). In Figure 4 and Appendix K.2, we show an experiment where we contaminate the test dataset with Fourier noises at varying frequencies $s$. This is still done with a fixed state matrix $\mathbf{A}$. In Figure 4, we plot the transfer function of the S4D initialization on top of the test accuracy, which reveals a strong correlation between the "spikes" in the transfer function and the effect of the perturbation at the corresponding frequencies. This connects our theoretical study of the S4D transfer functions to the discussion of the robustness of the model.

* Now, we consider the case where we indeed train the state matrix $\mathbf{A}$. In this new experiment, we train our SSMs as everyone does in practice (see Appendix K.2 for more details). Given our observation in Figure 4, we use the transfer function as a protocol to evaluate the robustness of our model. We watch the evolution of the transfer function (randomly selected from many transfer functions in the SSM) throughout training. The results are shown in Figure 13 and here is a summary of what is going on:

   * The spikes in the transfer function of the S4D model have slightly decayed but never vanished. We believe this slight decay is possibly due to the weight decay rate we set for the LTI matrices.

   * The locations of the spikes shift very slightly. This is the main reason why when conducting the experiment in Figure 3(a/b), we fixed the matrix $\mathbf{A}$: otherwise, we would have to eyeball the location of the spike after each epoch.

   * The transfer function of the S4-PTD model never develops any spike or instability throughout training.

In conclusion, this shows training does not completely solve the robustness issue of the S4D model. Neither does it impair the robustness of the S4-PTD model.

Since this paper is mainly about analyzing the initialization, we chose to put all these discussions into the appendices and leave them as a remark in the main manuscript. Understanding the training dynamics of the LTI systems is certainly an interesting future research avenue. We do remark that in our conclusion. We are happy to answer any follow-up questions that the reviewer may have and also include a more thorough empirical evaluation along this line in our camera-ready version should this manuscript be accepted.

Reviewer comment (Evaluation):
**Q.2:** Why is it relevant that the final peaks in Figure 1 are all (roughly) the same height?

Author response:
**Response:** The fact that the peaks are all roughly the same height means that $G_{\text{Diag}}$ does not converge to $G_{\text{DPLR}}$ uniformly. Hence, given any fixed state space size $n$, there is always a Fourier mode $\mathbf{u}$ on which the diagonal system and the DPLR system produce very different results. In other words, if the peaks vanish as $n \rightarrow \infty$, then by taking $n$ large enough, the transfer function $G_{\text{Diag}}$ will become a smooth function so it is robust to any Fourier-mode perturbation. However, due to the non-vanishing peaks, the S4D initialization is always not robust to _some_ Fourier-mode perturbation no matter how large $n$ is. This can be seen from Figure 4.

We hope this answers the reviewer's questions and concerns. We are happy to answer any follow-up question(s) that the reviewer may have.

Reviewer comment (Theory):
Is the constant magnitude at the end of Section 3.1 a conjecture, or easy to show from Lemma 1? How were the final spike locations derived? It seems possible to find the exact location for every \(n\) and calculate the magnitude of the spikes from Lemma 1, although the calculations may be difficult.

Author response:
**Response:** The constant magnitude of the last peak is a proved statement in our manuscript. It is proved in Appendix D: Proof of Theorem 2. (See the part around ``If we can show that $s_n=\Omega(n^2)$, then we have that $|G_{\text{DPLR}}(s_ni)-G_{\text{Diag}}(s_ni)| = \Omega(1)$...") In particular, we explicitly show that the height of the spikes stay constant and then use this fact to show that $G_{\text{DPLR}} -G_{\text{Diag}}$ does not converge to zero uniformly (i.e. Theorem 2). The calculation of the exact location of the last spike is difficult, involving solving a complicated trigonometric equation. However, we derived an asymptotic expression for it based on the angle of the denominator of the expression in Lemma 1 (see (24) in Appendix D).

Reviewer comment (Presentation):
Clarification is needed on interpreting the left side of the plot in relation to the right: are the final yellow/orange/blue spikes supposed to actually reach magnitude \(\approx 1\) (same as the purple spike), but the coarse sampling of the plot's x-axis misses the exact spike location, thus not rendering it precisely?

Author response:
**Response:** The reviewer is correct about the interpretation of the right panels. The big panel on the left does not have enough resolution to reveal the correct height of the spikes, and in the right panels, we zoom in to see that the height of the spike remains constant. We have added a clarification in the figure caption.

Reviewer comment (Theory):
The text claims the last spike is located at \(|s| = \Theta(n^2)\), but it seems like larger \(n\) correspond to earlier spikes. This appears important as it relates to a later question.

Author response:
**Response:** We apologize for the incorrect legend in Figure 1. The correct order from top to bottom should correspond to $n = 10000, 1000, 100$, and $10$, respectively. We thank the reviewer for pointing out this issue. We have corrected the legend.

Reviewer comment (Theory):
Regarding Section 3.3, there is uncertainty about the accuracy of the statement that “the outputs of the two systems diverge given the unit impulse (i.e., the Dirac delta function) as the input.” Based on the S4D paper, the outputs should be the impulse response or continuous convolution kernel \(K(t) = C e^{tA} B\), which S4D shows converges pointwise. This seems to conflict with Appendix F.3, Figure 9 in this submission, as the results look different. Referencing S4D Figure 2 and its corresponding reproduction.

Author response:
**Response:** The reviewer made a very good point here. First, we confirm that the output from the Dirac delta is equivalent to the convolution kernel used in the S4D paper. The real difference here is that the S4D paper mainly concerns the discretized system while our theory considers the continuous-time systems. For example, in S4D Figure 2, there is a distinction between the ZOH and the bilinear discretizations. However, from a continuous perspective, this distinction should not exist.

Now, we explain why there seems to be a dichotomy between our result and the S4D result. We are unable to post pictures on OpenReview, but we kindly advise the reviewer to run the following three commands in the Jupyter Notebook (https://github.com/HazyResearch/state-spaces/blob/main/notebooks/ssm_kernels.ipynb):

```
plot(N=128, measure='legsd', T=2.5, dt=0.01, yticks=[-1.0, 0.0, 1.0, 2.0], disc='bilinear')
plot(N=128, measure='legsd', T=2.5, dt=0.001, yticks=[-1.0, 0.0, 1.0, 2.0], disc='bilinear')
plot(N=128, measure='legsd', T=2.5, dt=0.0001, yticks=[-1.0, 0.0, 1.0, 2.0], disc='bilinear')
```

Then, the reviewer will notice that the convolution kernel of the S4D initialization under the discretization step $dt = 0.01$ is almost identical to that of the S4 initialization, whereas the convolution kernel under the discretization step $dt = 0.0001$ is wildly different.

Why does this happen? The reason is the so-called aliasing error [Trefethen, 2019, Ch. 4]: under a coarse discretization, the high frequencies cannot be observed. For example, one cannot tell the difference between $\cos(10\pi t)$ and $\cos(1000\pi t)$ by samples at $t = k\pi$ ($k \in \mathbb{Z}$). The same story applies here. If a discretization step $dt$ is fixed a priori, then there is a threshold $S_{dt}$ such that any frequency $s > S_{dt}$ collapse into a frequency lower than $S_{dt}$. Hence, if $n$ is sufficiently large, then by our theory, the S4 initialization and the S4D initialization differ mostly in a region of high frequencies. If this region goes beyond $S_{dt}$, then the two convolution kernels will appear almost identical. However, since our analysis is on a continuous level, this pathology will not appear because our "discretization step" $dt$ is truly infinitesimal. In conclusion, from a continuous perspective, the impulse response of the S4D initialization never converges to that of the S4 initialization. However, convergence of the discretized convolution kernel happens if a discretization step $dt$ is fixed a priori. We have also left this as a remark in our main text. We hope this addresses the reviewer's concern.

Lloyd N. Trefethen. Approximation Theory and Approximation Practice, Extended Edition. Society for Industrial and Applied Mathematics, Philadelphia, PA, 2019.

Reviewer comment (Reproducibility):
In Section 3.4, the experiment description does not specify which \(n\) is used. Based on Figure 1, it seems fixed at \(n=10000\).

Author response:
**Response:** We fixed $n = 32$ in this experiment, which is the default state-space size in running many ablation studies and examples from previous SSMs works. We have added this detail to Appendix G (previously as section 3.4). Again, it is our mislabeled legend in Figure 1 that gave the reviewer the false impression that $n = 10000$.

Reviewer comment (Experiments):
For Theorem 3, how tight is the upper bound empirically? The result suggests that as \(n\) grows, even with the error matrix norm controlled, the total output deviation still increases with larger \(n\). Is this behavior observed in practice—does the output deviation increase with state size?

Author response:
**Response:** We have conducted an empirical simulation to test the perturbation of the transfer function. The results have been added to Appendix I.1 (See Figure 11). In summary, we observed that as $n$ increases, if we fix the norm of $\mathbf{E}$, then in the average case, the transfer function perturbation decays with respect to $n$. As the reviewer noted, in this theorem, we fixed the norm of $\mathbf{E}$, so it is not surprising that $\\|G_{\text{DPLR}} - G_{\text{Pert}}\\|_{\infty}$ decreases as $n \rightarrow \infty$ in the average case because the relative perturbation size $\\|\mathbf{E}\\|/\\|\mathbf{A}\\|$ also decreases. On the other hand, Theorem 3 is a worst-case error estimate and we have added a theoretical discussion about the dichotomy between the worst-case behavior and the average-case behavior in Appendix I.1 as well.

Reviewer comment (Experiments):
How is the hyperparameter \(\gamma\) set for each dataset? It is not in the hyperparameter table. From Figure 3(c), results may be somewhat sensitive to this parameter.

Author response:
**Response:** In Appendix J.1, we revealed that for the LRA dataset, we choose $\gamma$ so that the size of the perturbation is about 10% of the size of $\mathbf{A}_H$. By Figure 3(b), this still empirically guarantees that our initialization is robust under the worst-case Fourier mode perturbation. We adopt the same convention for the interpolation/extrapolation experiment (now moved into the appendices).

Reviewer comment (Writing):
The writing is quite dense and can be improved if the authors make it more self-contained.

Author response:
**Response:** We have improved the presentation of our paper. While a list of major changes we have made is summarized in our general comment above, we highlight some of them that particularly made the manuscript less dense:

* We have moved the discussion of interpolation/extrapolation into the appendices. It used to be slightly cumbersome; we have replaced it with Figure 4, which verifies the potential non-robustness of the S4D initialization.

* We have relegated some technical discussion (e.g., Lemma 1) also into the appendices. Instead, we brought a numerical example (see section 3.2) into the main text, which helps elaborate the theory of the diagonal vs DPLR initialization.

* We have made our main practical contribution (perturbing the HiPPO-LegS matrix makes the model more robust) clear in both the introduction and section 4. Moreover, we have added a self-contained discussion about the backward vs forward errors in section 4 in the context of SSMs.

Reviewer comment (Theory):
Can the author elaborate more on why the proposed method can outperform the ones without diagonalization? This seems to be quite interesting.

Author response:
We find two ways of interpreting this question and we will provide our answers to both.

* On one hand, the reviewer may intend to ask why a diagonal structure (or a DPLR structure) of the matrix $\mathbf A$ is used instead of a dense one. In this case, the answer is that a diagonal/DPLR structure of the matrix $\mathbf A$ enables fast evaluation of the model and computation of the gradient. In comparison, evaluating an LTI system with a dense matrix $\mathbf A$ takes $\mathcal{O}(Tn^2)$, where $T$ is the length of the input and $n$ is the size of the state space. This makes training infeasible. We refer interested readers to [Gu et al., 2021] for more information.

* On the other hand, the reviewer may wonder why our S4-PTD (having the diagonal structure) model outperforms the S4 model (having the DPLR structure) on many LRA tasks.  Here are some factors that might be relevant and can possibly explain this advantage of the diagonal structure.

   * The diagonal structure requires fewer parameters to be trained and fewer architectural configurations to be tuned. These possibly make the training easier.

   * A diagonal structure of $\mathbf A$ corresponds to writing the transfer function $G$ into partial fractions while a DPLR structure of $\mathbf A$ corresponds to writing $G$ into the so-called rational barycentric formula. Partial fractions are usually more unstable to parameter perturbations compared to the barycentric formula. In the context of training, this might be an advantage of using a diagonal matrix $\mathbf A$ because it allows us to escape from the local minima even with a small learning rate.

   * This is not an explanation of why the S4-PTD model sometimes outperforms the S4 model, but we want to remark that from an expressiveness perspective, enforcing $\mathbf A$ to be DPLR does not have any advantage over enforcing it to be diagonal, because the transfer function of either $\mathbf A$ is always a rational function with $\leq n$ poles and any such transfer function can be represented using both a diagonal $\mathbf A$ or a DPLR $\mathbf A$. Hence, one should not expect that the diagonal model is naturally inferior to the DPLR model.

We hope this answers the reviewer's questions and concerns. We are happy to answer any follow-up question(s) that the reviewer may have.