From generalization analysis to optimization designs for state space models

We thank the reviewer for the valuable comments and suggestions. In the following, we address the 4 weaknesses from Weakness 1 to Weakness 4 and the 5 questions from Q1 to Q5.

Weakness 1. My biggest concern is the assumption that the inputs in Theorem 1 are sampled from a Gaussian process.

• 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.

Weakness 2. 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.

• 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.

Weakness 3. 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.

• 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.

We thank the reviewer for the valuable comments and suggestions. In the following, we address questions from Q1 to Q3.

Q1. In Theorem 1 , the authors claim that the SSM generalization is characterized by the temporal dependencies of the sequential data. More details on how does the dependency of the sequential data affect the generalization error should be included. Moreover, in order to achieve small generalization error, the mean and variance of the GP should remain a small level. While these two key parameters rely on the GP assumption, independent of data. This seems inconsistent with data-dependent generalization error bounds, as claimed in the paper.

• Our generalization bound is data-dependent in the sense that it depends on the statistics of the input sequence and we can only calculate the bound after observing the input sequence. In particular, we can see that the dominant term is given by

For the SSM basis function $\rho_\theta(s) = C e^{As}B$, it is exponentially decay with respect to $s$ because the eigenvalues of $A$ are negative. Therefore, we do not require the mean and variance to be uniformly small in $[0, T]$ to get a small generalization gap. We also relax the GP assumption to a more general class of random processes that are sub-Gaussian and Hölder continuous, and we refer the reviewer to the response to the comment 5 from Reviewer 2971 for more details. We have added the above discussions to make the generalization bound more clear after Theorem 1.

Q2. In speak of enhancing the robustness of SSMs on different temporal dependencies, the authors take $1/\sqrt{\tau(\theta)}$ as a rescaling factor for initialization. Any theoretical guarantees (e.g. variance analysis) on the robustness comparing with the HiPPO framework?

• We would like to emphasize that the rescaling method is built on the HiPPO based initialization, and it is applied before training. The HiPPO based initialization method is independent of the input sequence while we rescale the initialization based on the input sequence. This rescaling method improves the robustness of the initial output value scales of SSMs across different temporal dependencies of the input sequence. As Proposition 1 shows, the scale of the SSM output values 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.2.

Q3. The main techniques adopted in the proof are sub-exponential property of r.v. and Borell-TIS inequality, how did they yield to temporal dependency generalization bounds since both of them are temporal independent.

• 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.

Q3. Can you clarify how the rescaling in Line 7 of Algorithm 1 works: (a) across epochs and (b) extends to multiple layers. R.e. (a): is the value of $\tilde{C}$ rescaled at the beginning of every epoch? I.e. it is being “reinitialized” by rescaling its previous values. R.e. (b): does rescaling $\tilde{C}$ at each layer disrupt the action of other layers? Or is there a different method for rescaling between layers?

• The rescaling method is only applied for SSMs before training. We calculate the complexity measure layer by layer and then rescale the matrix $C$ from the first layer to the last layer sequentially. For example, for a $k$-layer SSM with the initial matrix $C_1,\ldots,C_k$ at each layer, we first calculate the complexity measure $\tau_1(\theta)$ for the first layer and rescale $C_1$ by $C_1/\sqrt{\tau_1(\theta)}$. Then we calculate the complexity measure $\tau_2(\theta)$ for the second layer for the updated input sequence of layer 2 and rescale $C_2$ by $C_2/\sqrt{\tau_2(\theta)}$. We repeat this process until the last layer. We have added the above discussions in Section 4.2.

Q4. The experiment in Figure 2, is it really Gaussian white noise? Or is it more like Brownian motion?

• It is Gaussian white noise. The autocovariance function for Brownian motion is given by $K(s,t) = \min(s,t)$ rather than $\frac{1}{|b| \sqrt{\pi}} e^{-((s-t)/b)^2}$.

Q5. Does the training loss in Figure 2 (right) include the regularization term? I believe it should actually be labeled as “Training set MSE”.

• Figure 2 (right) shows the training loss when training with the initialization scheme but without the regularization method. We have highlighted it in Section 4.2.

Q6. An appealing benefit of S4 is the zero-shot transfer to other sampling frequencies. However, this might change the scale of the time-dependencies. Can the authors clarify whether there are drawbacks to this method with respect to zero-shot transfer?

• This is an interesting question. In fact, our generalization measure is zero shot transerable because the input sequence statistics $\sqrt{K(s,s)}, \mu(s)$ in our generalization bound both have the same scaling as the input sequence $x(s)$, i.e., by changing sequence $x(s)$ to $x(k s)$ for $k>0$, the standard deviation and the mean of $x(ks)$ is changed to $\sqrt{K(ks,ks)}, \mu(ks)$, respectively. Then the generalization measure is invariant after changing the sampling frequency if the SSM output value is also invariant. Therefore, the zero shot transferability for SSMs is automatically inherited by our generalization bound. In line with the question, we have added a thorough discussion at the end of Section 4.1.

Q7. The theory is presented for linear SSMs, but in practice, multi-layer S4 models are interleaved with position-wise nonlinearities. I cannot see any discussion of how these nonlinearities (and the parameters in these nonlinearities, e.g. GLU) interact with the regularization of the parameters in the SSM, or, how the warping effect of the nonlinearity affects/interacts with the theoretical results.

• Since we have relaxed the original assumption of Gaussian process to a more general class of random processes that are sub-Gaussian and Hölder continuous, the theory can be extended if the nonlinearities do not affect the Hölder condition and the sub-Gaussian property. In fact, the assumption is to ensure that the input sequence is bounded (with high probability) in a given time interval. Therefore, intuitively, if the nonlinearity function has a controllable growth rate, then the boundness of the input sequence with the nonlinearity is still satisfied and our theory works if one considers a single layer SSM. For example, for Lipschitz nonlinearities, we list some examples in Appendix C to show that the Hölder condition and the sub-Gaussian property are still satisfied. We have added the above discussions in Section 5.

We thank the reviewer for the valuable comments and suggestions. In the following, we address the comments for the 7 comments and questions from Q1 to Q7.

Comment 1. There is quite a lot of work here to just get better generalization as shown on a near-pathological synthetic example, and marginal improvement on LRA (see Q.1. as well).

Q1. Can the authors clarify whether, in Table 2, w/o (8, 9) corresponds to the original S4 model? The numbers are slightly lower than in the original paper, and I am trying to clarify whether these numbers are like-for-like within the table, and how comparable to S4 they are.

• 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).

Comment 2. I would have also liked to have seen more evaluation of the initialization across different sizes of models, sensitivity to hyperparameters etc. Experimental repeats are also important to ensure that the results are reliable.

• 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.

Comment 3. The additional time complexity is also theoretically commensurate with the original S4 model, but I would like to see a concrete comparison of the runtimes to confirm this.

• 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.

Comment 4. I would also like to see a more thorough comparison to, e.g., the initialization and metrics suggested by Orvieto et al. [2023], or evaluation of whether this initialization/regularization scheme can be applied to methods adjacent to S4 (e.g. S4D, S5, Liquid-S4, MEGA).

• 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.