Paper: StableSSM: Alleviating the Curse of Memory in State-space Models through Stable Reparameterization (ICLR-2024)

Reviewer comment (Theory):
The analysis on the curse of memeory is limited to the simple state space model approximating linear decay. Reparameterization technique may not be applicable for sequential models with more complex structures.

Author response:
The reparameterization technique can also be applied to sequential models with more complex structures in the following sense: 
    In multi-layer state-space models, we implement the analysis outlined in Theorem 3.4 on a layer-by-layer basis. 
    When examining the gradient of recurrent weights, denoted as $\lambda=f(w)$, the same bounds on gradient norm are applicable even for more complex structures. 
        $$G_{f}(w) := \left | \frac{\partial \textrm{Loss}}{\partial w} \right | \leq C_{\mathbf{H}, \widehat{\mathbf{H}}_m} \frac{|f'(w)|}{f(w)^2}.$$
    The coefficient $C$ is a target-dependent and model-dependent coefficient, yet it is **independent of the parameterization scheme $f$**.
    Therefore our Theorem 3.4 (in Page 6) is also effective for more complex structures. 
    The reparameterization only change the recurrent weights' gradients, the gradients of other weights will not be influenced. 

    We include further numerical experiments (see Appendix D, page 22-24) on the image classifications and associative recalls to demonstrate the effectiveness of stable reparameterizations.
    It can be seen the stable reparameterizations guarantee a good performance while the case without reparameterization fail to achieve convergence.
    As we maintain the same initializations of the weights, the only difference across these models are the parameterization schemes of the recurrent weights $\Lambda$.

Reviewer comment (Theory):
It is not clear that why the "best" reparameterization should satisfy that the gradient is Lipshitz to the weight.

Author response:
We propose the "best" reparameterization to further improve the optimization stability. 
    The **intuitive** explanation is that a function with "gradient is Lipschitz to the weight" is more similar to the quadratic functions.
    Therefore this property is important for the optimization. 

    A theoretical explanation will be given as follow: 
    For non-convex optimizations, the Lipschitz-smooth property is a commonly used assumption to achieve the convergence of the models [1]. 
    Without Lipschitz-smooth property, function $f(x) = \sqrt{x}$ cannot converge to the optimal solution $x^*=0$ using fixed positive learning rate $\eta$. 

    The criterion "gradient is Lipschitz to the weight" is a **necessary condition** for the stability in the following sense (We include the detailed proof in Appendix F, page 25):
- Consider functions $f(x) = x^4$, the gradient function $g(x)=4x^3$ does not have a global Lipschitz coefficient for all input values $x$. Therefore for any positive fixed learning rate $\eta$, there exists an initial point $x_0$ such that the convergence cannot be achieved via the gradient descent step. 
    $$x_{k+1} = x_k - \eta g(x).$$ 
    Usually a smaller learning rate and a fine-tuned learning rate scheduler is necessary for the  convergence at that $x_0$. 
- Consider functions $f(x) = x^2$, the gradient function $g(x) = 2x$ is associated with a Lipschitz constant $L=2$. Then the same gradient descent step converges for any $\eta \leq \frac{1}{L} = \frac{1}{2}$ and any initial point $x_0$.
- As can be seen in the above two examples, the criterion **"gradient is Lipschitz to the weight" is associated with the convergence under large learning rate**. 
    As the use of larger learning rate is usually associated with faster convergence [2], smaller generalization errors [3], we believe the Lipschitz criterion is a suitable stability criterion for the measure of optimization stability. 
- The gradient-to-weight ratio evaluated in Figure 4(a) is a numerical verification of our Theorem 3.4. 
    The gradients scale of stable reparameterizations are milder therefore less susceptible to the well-known issue of exploding or vanishing gradients [4,5].

Reviewer comment (Evaluation):
As the numerical examples mainly show the results on gradient to weight ratio, how can you justify that this quantity is related with better training?

Author response:
Kindly refer to the response provided for point 2 of the weaknesses, particularly the last paragraph, for further clarification on this question.

Reviewer comment (Theory):
Can this analysis be extended to more complex models such as multi-layer RNN or transformer-based RNN?

Author response:
As we mentioned in reply to point 1 of the weaknesses, our analysis can be extended in the following sense: 
    When we evaluate the gradients of the recurrent weights in complex multi-layer models, the state-space model layers also satisfy the same gradient bounds from the Theorem 3.4.

Reviewer comment (Experiments):
My main concern with this paper is the disconnect between theory and experimental results: the authors train few small-scale Hyena-SSM model on wikitext, and then attempt to explain how the ranking in training loss corresponds to a ranking with the proposed gradient-norm scale. There are no attempts to perform multiple runs, try in different applications, or verify whether this hypothesis holds with different hyperparameters. In fact, Appendix D shows the rankings change completely by tweaking the learning rate.

Author response:
Here are the connections between theory and experimental results:

1. We study the stability of state-space models from the parameterization viewpoint. 
2. For **approximation** theory, we point out the importance of stable reparameterization in learning long-term memory. 
    We prove that stable reparameterization is **necessary** (Theorem 3.1, Page 5) and **sufficient** (Theorem 3.3, Page 6) for the long-term memory learning. 
    The necessity is proven in the sense that without reparameterization, the state-space models can only stably approximate targets with exponential memory decay. 
    The sufficiency is proven in the sense that with stable reprameterizations, the linear RNN can stably approximate the linear functionals with **any** memory decay. 
    **We identify a family of stable reparameterizations** (Definition 3.2) that can achieve the stable approximation of long-term memory targets: 
    For some continuous function $g: [0, \infty) \to [0, \infty), g(0)=0$,
    $$\sup_w \left[ |f(w)| \sup_{|\tilde{w} - w| \leq \beta} \int_0^\infty \left | e^{f(\tilde{w}) t}  - e^{f(w) t} \right | dt \right ] \leq g(\beta).$$
    This family includes a large set of stable reparameterizations. 
    In Appendix B.4 Remark B.3 (Page 17), we show the commonly used parameterizations such as the exponential and softplus parameterization are all stable parameterizations. 
3. Building upon the above theoretical results, it's important to note that there exist alternative reparameterization methods, which, though less explored, are equally capable of achieving a stable approximation of long-term memory. 
    This leads us to a pertinent inquiry: What constitutes the **optimization criterion** for selecting the most stable reparameterization approach in practical applications?
    The Lipschitz smoothness of the gradient is commonly adopted as an assumption for non-convex optimization, therefore we propose the Lipschitz continuity of gradient in Equation 11 (Page 6) as an optimization criterion for good stability. 
    Further details regarding the motivation for this optimization criterion are comprehensively outlined in Appendix F, on page 23.
    The gradient norm bounds and the Lipschitz continuity criterion allow us to compare different stable reparameterizations and find the **"best" in the criterion of Lipschitz continuity**. 

    (**Wikitext as verification**)
    The training of Hyena-S5 model with 120M parameters on Wikitext is used as a verification for the approximation and optimization properties in Theorem 3.4. 
    We show the initial gradient-over-weight scale in Figure 4(a) and corresponding training loss in Figure 4(b).

    (**Numerical verification of "best" parameterization**)
    The experiments on Wikitext only present the performance of stable reparameterizations over unstable reparameterizations. 
    We further verify the optimization stability of "best" reparameterization in the following extreme setting. 
    We construct a large scale language model with 3B parameters and train with larger learning rate (lr=0.01).
    As can be seen in the following table, the only convergent model is the model with "best" reparameterization. 
    We emphasize that the **only difference between these models are the parameterization schemes** for recurrent weights. 
    Therefore the best reparameterization is the most **stable** parameterization. 
    (We repeat the experiments with different seeds for three times. The detail of the experiment is given in Appendix D.2, page 24)

| Reparameterizations | Convergent / total experiments |
|------------|-----|
| **"Best"** | 3/3 |
| Exp        | 0/3 |
| Softplus   | 0/3 |
| Direct     | 0/3 |

Reviewer comment (Experiments):
The experiments rely on exploring different options for the implicit parametrization \( f(w) \) of the eigenvalues of a real-valued SSM. Have you considered using small hypernetworks, or alternative parametrizations? Why parametrize only the poles of the SSM implicitly, and not also the residues?

Author response:
Thank you for the interesting proposal. 
    The approximation and optimization results of our work have close connections with the suggestion in the following sense:
    - First, the result in Equation 9 gives a sufficient condition for stable approximation. 
        This condition identifies a family of "feasible hypernetworks" for learning long-term memory. 
        Therefore the parameterization with the constraint from Equation (9) gives a feasible set for further optimization. 
    - Second, the Lipschitz criterion (Equation 11) is taken as the specific target for optimization stability in this paper. 
        Our "best" parameterization is proposed based on the optimization stability criterion - "Lipschitz continuity criterion". 
        We solve the ordinary differential equation and achieves a low-freedom parameterization: $f(w) = \frac{1}{aw^2+b}$ with only two degress of freedom. 
        This is consistent with the optimization over hypernetwork idea that your suggested. 
    - Future promising direction includes using hypernetworks and finding suitable reparameterization with gradient-based optimization. 
        With constraints from Equation 9 and new objectives, we can still optimize the hypernetwork, which is out of scope of current work. 

    The reason we only parameterize the poles of the SSM implicitly is based the empirical evidence that the output of the model for large time is mostly sensitive to the value of poles [1,2,3,4]. and the gradient of the state-space model parameters
    The state-space models accumulated the hidden state with linear system $h_{k+1} = \Lambda h_k + U x_k$. 
    **When the sequence length $T$ is large, $h_T = \sum_{i=0}^{T-1} \Lambda^{T-k} U x_k$, even a minor difference of $\Lambda$ will result in a significant change of hidden state $h_T$.**
    
    The selection of parameters for other layers is equally important, though they are less prone to stability issues.

Reviewer comment (Evaluation):
Have you considered the effect of the entire architecture block, composed of gating and the SSM, on the choice of parametrization? Could that inform a better metric that better correlated with performance in practice?

Author response:
The entire architecture block is usually complex with different stack of normalization, residual connection and dropouts. 

    Our analysis can be generalized to the multi-layer state-space models of complex architectures: 
    Consider the gradient of the recurrent weights in a complex architecture, we still have the same gradient norm bound as characterized by the Theorem 3.4. 
    In particular, the gradient norm bound only depends on the target function, current model function, and the parameterization-dependent scale $\frac{|f'(w)|}{f(w)^2}$. 
    The detailed discussion is presented in the Q1 of reply-to-all.

Reviewer comment (Evaluation):
Can you provide some downstream evaluation of the language model, or at the very least validation loss?

Author response:
This paper primarily presents a theoretical exploration of parameterization properties.
    We provide the validation loss and perplexity with the standard error for the wikitext103 language modelling. 

    | Reparameterizations | Train ppl | Train loss | Test ppl | Test loss |
    |-------------|-----------|------------|----------|-----------|
    | Best        | 17.182    | 2.844      | 20.811   | 3.035     |
    | Exp(S5)     | 15.721    | 2.755      | 20.218   | 3.007     |
    | Softplus    |**14.570** | **2.679**  |**20.136**| **3.003** |
    | Direct      | 18.916    | 2.940      | 28.167   | 3.338     |
    
    Due to the small size of wikitext dataset and limitation of computing resources, we do not carry out the downstream evaluation of the language models. 
    This evaluation will be part of future numerical evaluation directions.

Reviewer comment (Theory):
Eq (2) does not correspond to practical instantiation of SSM-based models, which have linear readout (\( c^T h_t \)), then a pointwise shallow MLP (either with gating or without).

Author response:
The Equation 2 corresponds to the practical instantiation of SSM-based models in the following sense: 
    Any practical instantiation of SSM-based models can be implemented as a stack of Eq (2). 
    The pointwise shallow MLP can be realized with state-space model layer with layer-wise nonlinearity by setting recurrent weights $W$ to be 0. 
    $$h_{k,l+2} = U_1 \sigma (U_2 h_{k, l} + b_2) + b_1.$$
    We give a graphical demonstration in Appendix E, page 24. 
    Therefore, the Equation 2 is given in a simplified form without sacrificing the expressiveness of the multi-layer state-space models. 
    The results established from Theorem 4.1, 4.3 and 4.4 also hold for the practical versions.