On Feature Learning in Structured State Space Models

We sincerely appreciate your thoughtful review and positive assessment of our work. Your feedback has been invaluable in helping us improve the paper. We have addressed your concerns regarding improved presentation and organization of the proofs. We hope that our revisions will merit your consideration for an increased score.

Improved presentation. To enhance the visualization of the Mamba forward pass, we have created a new figure (please see Figure 2 in the PDF attached to the global response). We will integrate this figure into the main paper and expand Section 3.1 to provide additional background on SSMs for enhanced readability. We will also introduce all notation and results used such as the definition of the spectral norm, Kolmogorov’s SLLN and Lindenberg-Feller CLT, in the appendix.

Modularization of the proofs. Thank you for your suggestion. We concur that modularizing the proofs will certainly enhance the paper's accessibility. However, analyzing signal propagation in the Mamba architecture is inherently more complex than in MLPs, necessitating some analysis of specific modules. Nevertheless, we came up with a strategy to modularize the analysis by noting that the Mamba layer can be mechanistically viewed in the forward pass as a sequence of 3 components/sub-layers:

1. Selection
2. Discretization
3. Linear recurrence

Our visualization of the Mamba layer is based on this perspective (in the PDF attached to the global response). To our knowledge, all SSM layers comprise some or all of these sub-layers, allowing us to modularize both forward and backward signal propagation analysis through each sub-layer. It's worth noting that different discretization rules (e.g., ZOH or Euler) would require separate analyses. Beyond this, while SSMs primarily differ in their initialization of the transition matrix A, our analysis remains applicable to any initialization chosen according to HiPPO theory.

We believe these changes will significantly enhance the clarity and accessibility of our work, and we look forward to incorporating them in our revised manuscript.

We would like to express our gratitude for your thorough and insightful review of our manuscript, and useful feedback for improving our work.

Order of limits and practical applicability. In SSM models, $N_u$ (input dimension for SSM component) typically increases much faster than $N_x$ (latent states dimension for SSM component) during scaling (see Table 10 in the original Mamba paper [1]). This makes the limit where $N_u$ approaches infinity before $N_x$ more practically relevant. Furthermore, fortunately, these limits commute which allows us to extend our results to the proportional limit setting.

Additional Experiments. We have conducted additional experiments to further validate our theory. We have summarized them in the global response and the results are provided in the PDF attached to the global response. Addressing your concerns, we now also present results on a randomly sampled subset of the Fineweb dataset, suggesting similar conclusions to the wikitext-103 results in Figure 2 (main paper). Separately tuning the learning rates in SSM layers and non-SSM layers shows the benefit of muP-SSM over spectral scaling more clearly: While muP-SSM continues to improve with scale, spectral scaling no longer monotonically improves beyond a certain width threshold. Even without separate learning rates, muP-SSM markedly improves training stability at larger learning rates over spectral scaling. SP consistently performs worse than muP-SSM in terms of stability and generalization. This indicates that muP-SSM can indeed improve performance at larger scale.

Implications for other SSM variants. Mamba is among the most complex SSM architectures, containing all components present in other recent SSM variants. By providing a scaling analysis for Mamba analogous analyses for S4 or LRUs follow as corollaries. We will clarify implications for other SSM variants in the revision (see our answer to Reviewer iUa6).

Implications of enabling feature learning and computational considerations. We see no negative implications in enabling feature learning in every layer. If feature learning is undesired in a specific layer, training can be explicitly disabled, saving computation instead of letting the updates to implicitly vanish with scale. muP-SSM therefore provides the flexibility to achieve feature learning in SSM layers when desired. muP-SSM introduces no additional computational complexity, as it's merely a different parameterization of the weights.

Have you explored whether μP SSM enables training of wider SSMs than previously possible? While we lack the computational resources to experiment at such large scales, there's evidence suggesting muP-SSM could enable training of wider SSMs. Even at relatively small scales, muP-SSM shows markedly improved training stability at larger learning rates compared to spectral scaling. According to [3], instabilities in small-scale models at large learning rates can predict general training instabilities at larger scales. This evidence provides some support for the possibility of training wider SSMs with muP-SSM, though direct large-scale experiments are needed for confirmation. We hope that the SSM community would conduct experiments at larger scales to further investigate this potential.

Intuitive explanation of why SSMs violate the TP ansatz. There are two primary reasons for why SSMs such as Mamba cannot be represented via TP:

1. Structured transition matrix (A). Typically, the state transition matrix A is highly structured and chosen according to (or loosely based on) Hippo theory [2]. One example is the diagonal matrix with the $i$th diagonal entry being $i+1$. The TP framework crucially relies on matrices with i.i.d entries (such as i.i.d Gaussians) and cannot represent such structured matrices.
2. Selection mechanism. Activations and weights play different roles under TP. The selection mechanism first computes activations (linear transformations of input parameterized via some weights) and uses them as weights in the linear recurrence (see the pdf attached to the global response for a visualization). This underlies a second source of incompatibility of SSMs with the TP framework.

[1] Gu, Albert, and Tri Dao. "Mamba: Linear-time sequence modeling with selective state spaces." arXiv preprint arXiv:2312.00752 (2023).

[2] Gu, A., Dao, T., Ermon, S., Rudra, A., & Ré, C. "Hippo: Recurrent memory with optimal polynomial projections." NeurIPS 2020.

[3] Wortsman, M., Liu, P. J., Xiao, L., Everett, K., Alemi, A., Adlam, B., ... & Kornblith, S. "Small-scale proxies for large-scale transformer training instabilities." arXiv:2309.14322 (2023).

Thank you for your constructive feedback and positive evaluation of our work. We appreciate your insights and we have carefully addressed your major concerns in the following response. We will rectify minor issues like typos in our revision. We hope that our revisions warrant your consideration for an increased score.

Additional Experiments. We have conducted additional experiments to further validate our theory. We have summarized them in the global response and the results are provided in the attached PDF. Addressing your concerns, we conduct experiments on a randomly sampled subset of the Fineweb dataset which suggest similar conclusions as those on wikitext-103 in Figure 2 (main paper). muP-SSM outperforms SP in terms of test loss, training stability, monotonic improvement of generalization with scale, and transferability of optimal learning rate across scales. Observations in comparison to spectral scaling are more nuanced. muP-SSM markedly improves training stability at large learning rates over spectral scaling, and slightly improves generalization performance. First tuning the learning rate in non-SSM layers and subsequently tuning a separate learning rate in SSM layers separates muP-SSM and spectral scaling more clearly: While muP-SSM continues to monotonically improve with scale, spectral scaling no longer improves beyond a certain width, suggesting that the reasonable performance of spectral scaling stems from the non-SSM layers. In SSM layers, the spectral scaling approach lacks rigorous theoretical justification.

Clarification on the asymptotically i.i.d assumption. In Section 2.2 & A.1 of our paper, we discuss that activations and their updates in neural networks representable as Tensor Programs (TPs) become asymptotically independent and identically distributed (i.i.d.) with increasing width. This result underlies our i.i.d assumption. Note that, we consider the practical setting where SSM layers are embedded within a network of non-SSM layers (such as MLPs or Layernorms) representable via TP. Accordingly, it's natural to assume that inputs to the SSM layer (e.g., activations from the previous non-SSM layer) are asymptotically i.i.d at each recurrence step. In fact, this i.i.d. assumption isn't strictly necessary. It suffices to assume that inputs are correctly scaled, i.e., $\vert \vert u\vert \vert_2 \in \Theta(\sqrt{N_u})$. This scaling has been demonstrated for correctly scaled network modules representable via TP. Crucially, we do not assume that the sequence of inputs $u_1, u_2, \cdots, u_L$ are i.i.d. Rather, we assume that the coordinates of a single input $u_i$ are asymptotically i.i.d. (or at least correctly scaled) in the sense described above.

Why should N_u and N_x go to infinity? We would like to clarify that our analysis does not require both $N_u$ (input dimension for SSM component) and $N_x$ (latent states dimension for SSM component) to go to infinity. It merely allows both quantities to scale up. To derive the results for when only $N_u$ goes to infinity, one can simply set $N_x \in \theta(1)$. In this simplified scenario where the dimension of the latent states $N_x$ is fixed, heuristic muP spectral scaling yields $\eta_B,\eta_C=\Theta(\frac{1}{N_u})$ while muP-SSM gives $\eta_B,\eta_C=\Theta(\frac{1}{\sqrt{N_u}})$ (see Table 1 in the main paper). Experiments carried out independently by researchers (to be acknowledged in the public version), indeed verify that the correct width-independent scaling aligns with muP-SSM (results as shown in Figure 4 in the attachment). Note, however, that as models are scaled up both $N_u$ as well as $N_x$ may be increased in practice, albeit at very different scales. For example, see Table 10 in the original Mamba paper [1].

In Definition 2.1, why do we require that there exists one input that is correctly scaled? This follows the same logic as the paper that proposed muP [3, Definition H.9], which defines a parameterization to be feature learning iff there exists a training routine and input that result in the correct update scaling. This is a technicality we have to adopt as there exist degenerate combinations of inputs and learning rates that result in smaller scalings. Even under this definition, there exists only one choice of layerwise initialization and learning rate scalings that achieves feature learning.

[1] Gu, Albert, and Tri Dao. "Mamba: Linear-time sequence modeling with selective state spaces." arXiv:2312.00752 (2023).

[2] Wortsman, M., Liu, P. J., Xiao, L., Everett, K., Alemi, A., Adlam, B., ... & Kornblith, S. . Small-scale proxies for large-scale transformer training instabilities. arXiv:2309.14322 (2023).

[3] Yang, Greg, and Edward J. Hu. "Feature learning in infinite-width neural networks." arXiv:2011.14522 (2020).