Paper: From generalization analysis to optimization designs for state space models (ICLR-2024)

Reviewer comment (Theory):
My biggest concern is the assumption that the inputs in `Theorem 1` are sampled from a Gaussian process. Perhaps the Gaussian assumption is more reasonable in a non-temporal setting, such as linear regression. However, most time series inputs that we encounter in practice cannot be ''sampled'' from a Gaussian process. For example, you cannot find a single GP that accounts for the flattened MNIST images, because images representing different numbers may have their own unique features, and such distinct features cannot be captured solely by the randomness in your GP. If you fix a GP and randomly sample your sequential pixels, then most figures you obtain won't represent any number. I understand that this Gaussian assumption is crucial in proving your generalization error bound and there is perhaps no way out, but this indeed results in a gap between your theory and the methodologies you proposed.

Author response:
* We agree with the reviewer that the Gaussian process assumption is a strong assumption. To address this concern, we have relaxed the Gaussian process assumption to a more general class of random processes that are sub-Gaussian and Hölder continuous.
This is a more general assumption than the Gaussian process assumption in the sense that it covers a larger class of random processes including Gaussian processes, Hölder continuous bounded processes.
In particular, as we mentioned in the first point of the response *To all reviewers*, when the range of the random process is a finite set, then the Hölder continuity condition can be dropped.
This relaxation includes the case of image classification tasks, where the input image is flattened as a sequence and the range for each pixel value is a finite set, e.g., $\{0,1,\ldots,255\}$.
**In the revised paper, we have substantially improved the generality of our theorem (see the revised statement in Theorem 1 and proof in Appendix E) by dropping the Gaussian process assumption and replacing it with a much weaker one in Assumption 1.**

Reviewer comment (Reproducibility):
The proposed regularization method (9) combines a normalized $\ell^2$ loss and a regularization term. In an SSM, one usually chains multiple LTI systems; however, only the target output of the entire SSM is known (e.g., whether the maze is solvable, what the number in the MNIST figure is). In that case, it is unclear how the ''target outputs'' $y_i$ of the intermediate LTI systems are defined.

Author response:
* We would like to mention that the regularized loss by the regularization method (9) is given by the original loss plus the regularization measure $\tau(\theta)$.
When training multi-layer SSMs, we calculate the complexity $\tau(\theta)$ at each layer and add them together as a total regularization measure.
Notice that the complexity for each single layer only depends on the input sequence statistics and the model parameters of that layer.
Therefore, we do not need to know how the target outputs $y_i$ of the intermediate systems are defined.
**To make the statement more clear, we have changed the notation for the regularized loss (9) to $\tilde{R}_n(\theta) := R_n(\theta) +  \lambda  \cdot \tau (\theta)$ and added the above explanations in Section 4.3.**

Reviewer comment (Evaluation):
The evaluation of the model does not show clear evidence of why the scaling of the initialization makes the model more robust. For example, in `Table 1`, comparing the cases `w/o (8), (9)` to `w/ (8)`, it seems that adding the scaling improves the training accuracy but makes the generalization accuracy even worse. This actually contradicts the claim that the model is made more robust by scaling the initialization.

Author response:
* There is a misunderstanding here due to the ambiguity of the term `robustness`.
We would like to emphasize that 'robustness' in the paper refers to the robustness of the *output value scales* (instead of the test performance) on SSMs across different temporal dependencies of the input sequence.
This is desirable because one does not a priori have a good sense of the input sequence and the temporal dependency of the input sequence may vary a lot.
Therefore, a good way to initialize SSMs should be one that is not sensitive to the variations of the temporal structures.
This is not true for the original HiPPO based initialization scheme, which is independent of the input sequence.
By looking at the left-hand side of Figure 2, we can see that the SSM output scales at initialization without rescaling (dashed lines) are very sensitive to the variations of temporal dependencies, while the SSM output scales at initialization with rescaling  (solid lines) are much more robust.
In Table 1, we can also see that adding the scaling improves the training performance but the generalization performance is not guaranteed to be improved.
The above empirical findings also match Proposition 1, which shows that the scale of the SSM output scales after rescaling does not depend on the variations of the temporal structures.
**To make the statement more clear, we have added the above discussions in Section 4.4.**

Reviewer comment (Experiments):
Since the regularization involves a hyperparameter $\lambda$, it is a good practice to perform an ablation study to demonstrate the effect of changing $\lambda$.

Author response:
* **We have added more numerical results on the sensitivity of the regularization coefficient $\lambda$ for unidirectional S4-Legs, bidirectional S4-Legs, bidirectional S4D-Legs in Section A.2**.
The results in **Table 6, 7, 8** show that the test performance is much more sensitive to $\lambda$ for the Pathfinder and PathX tasks compared to other tasks.
This is because the generalization measure for the Pathfinder and PathX tasks are much larger than other tasks (as shown in **Figure 3, 4, and 5**), making the generalization measure very sensitive to the magnitude of $\lambda$.
The ablations study on the sensitivity of $\lambda$ mathches the intuition that the Pathfinder and PathX tasks are more challenging than the other tasks.

Reviewer comment (Reproducibility):
The setup of this paper does not consider the matrix $\mathbf{D}$ in an LTI system. How easy is it to incorporate that matrix and do you need to scale $\mathbf{D}$ in initialization?

Author response:
* As we mentioned in the third point of the response *To all reviewers*, our generalization bound can be extended to the case with skip connection $D$.
This is because that our generalization bound does not have a restriction on the kernel function $\rho_\theta(s)$. When adding $D$, we can incorporate it into the kernel function to get a new kernel $\rho_\theta(s) + D \delta (s)$, where $\delta(\cdot)$ is a delta function.
However, the initialization scheme may not work well when $D x(t)$ is much larger than $\int_0^t \rho_\theta(s) x(t-s)$.
This is because that the initialization scheme only adjusts $C$ and requires the kernel function to be linear in $C$.
In that case, $D$ should also be rescaled.
**We have added the above discussions in Section 5.**

Reviewer comment (Theory):
In `Theorem 1`, can you show the explicit dependency of $C_T$ on $T$? This is important because in training an SSM, the discretization size $\Delta$ is usually trainable, making the final time $T$ in the continuous world change from time to time. Hence, in order to apply your theory, it is better if we can understand the role of $T$.

Author response:
* Yes, we have modified the generalization bound in Theorem 1 to show the explicit dependency of $C_T$ on $T$, which is given by $\sqrt{\log T}$.

Reviewer comment (Theory):
The presence of `Proposition 1` seems a bit abrupt. How does that relate to your ''robustness'' discussion? In addition, what kind of ''stability'' are you referring to? This is a fairly ambiguous term, which can represent the stability of a numerical algorithm, the asymptotic stability of your LTI system (i.e., if your eigenvalues of $\mathbf{A}$ are all in the left half-plane), or something else.

Author response:
* In view of the comment, we have changed 'robustness' to 'robustness of the output value scales on SSMs' in the revised paper.
The 'stability' in the original paper refers to the stability of the training loss value.
**To make the statement more clear, we add the above explanations after Section 4.2.**

Reviewer comment (Evaluation):
In your experiments, it is shown that the regularizer improves the training accuracy, which is a bit counter-intuitive. Do you have a justification for that?

Author response:
* Maybe there is a misunderstanding here. Our experiments do not show that the regularizer improves the training accuracy. By looking at Table 1, when comparing w/o (8, 9) vs w (9) and w (8) vs w (8, 9), we can see that the final training loss is not improved by adding the regularizer (9).
However, with the initialization scheme (8), the training performance is improved.
This is shown in Table 1 by comparing w/o (8, 9) vs w (8) and w (9) vs w (8, 9).

Reviewer comment (Theory):
The main techniques adopted in the proof are sub-exponential property of r.v. and Borell-TIS inequality. How did they yield temporal dependency generalization bounds since both of them are temporal independent?

Author response:
* We would like to emphasize that the sub-Gaussian property and the Hölder continuity condition in Assumption 1 are applied to sequences in a time interval $[0, T]$ instead of a single time point (see equation (14) for the terms that we need to bound) in our proof of Theorem 1. This is where the temporal dependency comes from.  
 Combining the $\varepsilon$-net argument and the covering number tools, we can apply the sub-Gaussian property to a sequence in a time interval $[0, T]$.  
 Since we have relaxed the Gaussian process assumption to a more general class of random processes that are sub-Gaussian and Hölder continuous, the proof of Theorem 1 has also been revised accordingly.  
 **In light of the question, we added more details in the proof sketch of Theorem 1.**

Reviewer comment (Experiments):
There is a considerable amount of work presented to achieve better generalization only on a near-pathological synthetic example and marginal improvement on LRA (see Q.1. as well).

Author response:
* We would like to emphasize that this paper is not just about getting better generalization. The theory is interesting by itself as it is the first data-dependent generalization bound for SSMs that is based on the temporal dependency of the input sequence.
  For the improvement on LRA, **we have added more numerical results for bidirectional models such as bidirectional S4-Legs and bidirectional S4D-Legs in Table 2 and Appendix A.2.**
  The results in the updated Table 2 show that our proposed regularization method consistently improves the generalization performance in most tasks for all three models. 
  In that sense, the improvement on LRA is not marginal.
  * Our original version used the unidirectional S4-Legs model (Gu 2022a), which processes a sequence in one direction from past to future.
  Compared to recent commonly used bidirectional models (Gu 2023), the unidirectional model is closest to our problem setting even though it is not the state-of-the-art model.
  Therefore, the numbers in Table 2 should be compared with the results in the earliest paper (Gu 2022a), and we can see that our numbers match with the results in (Gu 2022a).

Reviewer comment (Experiments):
The reviewer would like to see more evaluation of the initialization across different model sizes, sensitivity to hyperparameters, and experimental repeats to ensure result reliability.

Author response:
* **We have added more numerical results on the sensitivity of the regularization coefficient $\lambda$ for unidirectional S4-Legs, bidirectional S4-Legs and bidirectional S4D-Legs in Section A.2**.
 The results in **Table 6, 7, 8** show that the test performance is much more sensitive to $\lambda$ for the Pathfinder and PathX tasks compared to other tasks.
 This is because the generalization measure for the Pathfinder and PathX tasks are much larger than other tasks (as shown in **Figure 3, 4, and 5**), making the generalization measure very sensitive to the magnitude of $\lambda$.
 The ablations study on the sensitivity of $\lambda$ matches the intuition that the Pathfinder and PathX tasks are more challenging than the other tasks.
 For the repeated experiments, we have added the standard deviation (in 3 independent runs) of the test accuracy for the PathX task in **Table 2** for the bidirectional models.

Reviewer comment (Experiments):
While the additional time complexity is theoretically similar to the original S4 model, the reviewer would like to see a concrete comparison of runtimes to confirm this.

Author response:
* **We have added the running time per epoch in the updated Table 2 in Section 4.4.**
The running time comparison shows that the proposed regularization method brings a little extra computational cost compared to the original training method.

Reviewer comment (Experiments):
The reviewer would like a more thorough comparison to the initialization and metrics proposed by Orvieto et al. [2023], and/or evaluation of whether this initialization/regularization scheme can be applied to S4-adjacent methods (e.g., S4D, S5, Liquid-S4, MEGA).

Author response:
* **We have added more evaluation on the bidirectional S4D-Legs model in Table 2 and Appendix A.2.**
The results in **Table 2** show that our proposed regularization method and initialization scheme can improve the generalization performance.

Reviewer comment (Theory):
The reviewer questions the reasonableness of the assumptions and the tightness of the bounds in practice. They are unsure if the GP assumption is sensible in practice and note the lack of validation. They also ask how the fidelity of the GP approximation impacts the regularizer’s performance.

Author response:
* We have relaxed the Gaussian process assumption to a more general class of random processes that are sub-Gaussian and Hölder continuous.
This is a more general assumption than the Gaussian process assumption in the sense that it covers a larger class of random processes including Gaussian processes, Hölder continuous bounded processes.
In particular, as we mentioned in the first point of the response *To all reviewers*, when the range of the random process is a finite set, then the Hölder continuity condition can be dropped.
This relaxation includes the case of image classification tasks, where the input image is flattened as a sequence and the range for each pixel value is a finite set, e.g., $\{0,1,\ldots,255\}$.
**In line with the comment, we provide more discussions of the new assumption in Section 4.1 and Appendix C, and modified the proof of Theorem 1 accordingly.**

Reviewer comment (Theory):
The reviewer suggests trying to establish exactly how the initialization and regularization terms affect the learned model. They understand that L2 regularization reduces parameter magnitude, but they question what the regularizer in (9) encourages in learned models, and how regularized models differ from regular S4 models. This analysis could inspire better SSM structures.

Author response:
* The regularization method in the paper is different from weight decay which tries to minimize the magnitude of the parameters, but to minimize the correlation between the statistics of the input sequence and the SSM basis function.
As we emphasize in the introduction part of the paper, regular S4 models do not regularize the hidden state matrix $A$ because its imaginary part controls the oscillating frequencies of the SSM basis function (Gu 2022b).
Instead of penalizing the weight norm, our regularization method regularizes the magnitude of the SSM output value, which is determined by the correlation between the input sequence statistics and the SSM basis function. 
* For the initialization scheme, our proposed method is built on the original HiPPO based initialization, which is a *data-independent* initialization method.
However, when the temporal dependency of the input sequence varies, the SSM output value may vary a lot (see the dashed lines in Figure 2).
With our proposed initialization scheme, the SSM is more robust in the sense that the scale of the SSM output value does not depend on the variations of the temporal structures.
**We have added the above discussions in Section 4.2.**

Reviewer comment (Experiments):
The absence of Path-X in the paper is unfortunate, as the reviewer understands Path-X to be the only truly challenging LRA benchmark. While they are willing to overlook this for now, they encourage completing the benchmarks.

Author response:
* **In line with the comment, we complete the PathX task for both directional S4-Legs and directional S4D-Legs models in Section 4.4 and Appendix A.2.**
The results for the PathX task show that our proposed optimization designs improve the generalization performance for both directional models.