MAMBA STATE-SPACE MODELS ARE LYAPUNOV-STABLE LEARNERS

Authors: John Timothy Halloran, Manbir S Gulati, Paul F Roysdon
Primary area: foundation or frontier models, including LLMs
Keywords: Mamba, Mamba SSMs, SSMs, LLMs, PEFT, LoRA, Transformers

Abstract:
Mamba state-space models (SSMs) were recently shown to outperform state-of-the-art (SOTA) Transformer large language models (LLMs) across various tasks.  Despite subsequent widespread adaptation, little work has focused on Mamba LLMs' amenability for fine-tuning frameworks ubiquitously used for Transformer-based LLMs, e.g., mixed-precision fine-tuning (MPFT) and parameter-efficient fine-tuning (PEFT).  For the former, it currently remains an open question whether Mamba's recurrent dynamics are robust to small input changes, such as those encountered during MPFT.  Using dynamical systems theory (in particular, Lyapunov exponents), we answer this question in the affirmative.  We empirically validate this result through several experiments, showing that Mamba SSMs are significantly more stable to changes introduced by mixed-precision than comparable Transformers, even when both MPFT and PEFT are combined.  For PEFT, we show how targeting specific memory buffers in Mamba's customized CUDA kernels for low-rank adaptation regularizes SSM parameters, thus providing both parameter efficient learning and computational savings. Finally, with both MPFT and PEFT enabled, we explore the impact of instruction tuning Mamba SSMs for in-context learning (ICL) on natural language tasks.  While pretrained Mamba and Mamba-2 models only achieve 38% and 82% (respectively) of the ICL improvements of comparable Transformer-based LLMs, we show that instruction tuning allows Mamba models to narrow this gap to 81% and Mamba-2 models to skyrocket over this gap to 132%.

Decision: reject
Reviewer scores: 5, 6, 3

Reviews and discussion:
<thread>
creation date:
1730650417797

thread_type:
official_review

<thread_object>
creator: reviewer_jrwp

<content>
summary:
This paper studies Mamba, a state-space model (SSM) that has previously been introduced as an efficient alternative to Transformers. The authors analyze Mamba's stability in mixed-precision and parameter-efficient fine-tuning by applying Lyapunov exponents, theoretically and empirically confirming its resilience against input perturbations. Experiments show that Mamba maintains stability better than comparable Transformers under mixed precision and achieves notable memory and computation efficiency. Additionally, instruction tuning enhances Mamba's in-context learning, suggesting that SSMs can approach or surpass Transformer models in certain few-shot tasks.

strengths:
This work provides promising evidence that fine-tuning Mamba on instruction-tuning data can improve its in-context learning abilities.

weaknesses:
- Line 51:  For some training framework it does not keep master weight in fp32 which problematic for Mamba. But Mamba was still trained with mixed precision using AMP.

- Theory 1 seems shallow: Isn't Lemma 1 trivial to derive? \( dx_{t+1}/dx_t \) directly represents the decay rate, which must be \(\leq 1\) to prevent exponential growth in cumulative products. This insight is obvious for any recurrent model, making Theorem 1’s assertion about exponential decay rather redundant.

- Theory 2 seems unnecessary: It’s clear that concatenating multiple weight matrices with LoRA applied to the larger matrix creates dependencies. However, there's no practical advantage demonstrated here, nor is there an empirical basis for this approach.

- Line 274: Why use MMLU as a proxy metric? For models at this scale, MMLU scores are generally around random (i.e., ~25%), leading to limited interpretability. Additionally, only the differential score is provided, without the original MMLU results, which weakens MMLU’s reliability as a metric in this context.

- Instruction tuning experiments: Pythia feels like a weak baseline. And is it plausible to claim Mamba has an “emergent ability” (line 439) for in-context learning after instruction tuning, given SSMs' inherent limitations, as noted in [1,2]? Figure 3 also shows a performance drop for Mamba2 as shot count rises, likely due to state-space models’ well-known difficulties with long-sequence processing and in-context learning.

[1] https://arxiv.org/abs/2402.18510

[2] https://arxiv.org/abs/2402.18668

questions:
See weakness.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730681454096

thread_type:
official_review

<thread_object>
creator: reviewer_twtf

<content>
summary:
- The authors work to demonstrate that the Mamba SSM is robust to small input perturbations during training, meaning that Mamba is compatible with mixed-precision fine-tuning (MPFT) and parameter-efficient fine-tuning (PEFT). The authors include both theoretical and empirical evidence for this finding.
- Theory results are in section 3, using dynamical systems theory (Lyapunov exponents) to demonstrate that Mamba’s recurrent dynamics are robust to small input changes
- Section 5’s empirical results contain the following experiments (this list defines the numbers which I’ll use to refer to experiments later):
  1. Main results: divergence in accuracy between FP32 and FP16 or BF16 inference, showing that the accuracy divergence for Mamba is roughly similar to that for pythia and OLMo
  2. Non-divergent mamba fine-tuning: divergence in MMLU accuracy compared to full-precision for FP16 or BF16 models fine-tuned fully, or using LoRA on all or a subset of linear layers. Shows that divergence for Mamba is generally less than similar-size transformers.
  3. Hardware throughput and memory-utilization improvements: Average tokens-per-second and maximum memory-per-token for mamba of varying sizes, using FP32 full fine-tuning, or FP16/BF16 LoRA on all or a subset of linear layers. Shows that any form of PEFT is faster and less memory-intensive than the FP32 full fine-tuning
  4. Instruction tuning’s impact on Mamba ICL for natural language tasks: authors claim that instruction tuning narrows the ICL gap between Mamba and Pythia
  5. ICL as an emergent ability of Mamba SSMs: compares ICL performance for Mambas and several transformer models of varying sizes

strengths:
- The paper addresses a very important problem that, until now, has helped to prevent the wider use of SSMs: training resources for an SSM have been greater than a similarly-sized transformer, because the transformer has been known to be compatible with efficient fine-tuning techniques. Addressing this disparity will help more people to consider working with SSMs, as their computational requirements will suddenly be within reach
- The combination of both theory and empirical results is quite persuasive and is convincing that Mamba is indeed compatible with MPFT/PEFT
- Theoretical results are simple, elegant and important. The proofs in the appendix are easy to follow and look solid.
- Care has been put into the experimental setup, for instance everything being trained using the same batch size in the main results

weaknesses:
- My most major critique is that there aren’t any error bars or sense of the spread/randomness in the empirical results, especially for experiment 2 / figure 1, experiment 4 / figure 3 for the fine-tuned models and experiment 5 / figure 4 for the fine-tuned models.
-  My second biggest critique is that each experiment is only performed with one dataset (MMLU, fine-tuning with Alpaca for fine-tuning experiments). Including more than one dataset or task would strengthen the claims you’re drawing from the empirical evaluations. Experiment 5 includes more datasets, but the results for all datasets are averaged together and readers can’t consider their performances separately
- The graphs are nice in figures 1-4, but it would also be helpful to include the actual numbers in tables in the appendix
- The paper seems to have a lot of extra contextualization in it, taking up room which could be devoted to more detailed empirical analysis. In particular, section 2 background and section 5 related work seem to just both be disjoint related work sections. Usually I would consider the “background” section to be merely preliminary information and definition of symbols that are necessary to understand the rest of the paper, like a statement of the SSM equations and formal definition of Lyapunov exponent. Then, related work would be a longer section detailing, well, related work. Just as an example for comparison, you can see this differentiation in roles between background and related work in this paper’s sections 2 and 3 here: https://arxiv.org/pdf/2409.00717 Contextualization/related work is important to include, but I would consider prioritizing and moving at least some of it to appendix to include more empirical results
- Some of the number references seem to be wrong:
  - Line 345, “Figure 5” seems to actually be referring to Figure 1
  - Line 379, “Figure 5” seems to refer to Figure 2 (did you maybe put your \label{} command for the figures in the wrong spot inside your figures’ code blocks?)
  - Appendix theorem 3 corresponds to main paper theorem 2.

Overall, I'm setting my rating to 6 due to the significance of the theoretical results. However, I would appreciate more detail to the empirical evaluations as I mention above.

questions:
Why did you choose to work with different sets of transformer models in the main results versus the other experiments? For instance, you give results for OLMo on the first experiment (Table 1) but not second, and results for OpenELM in all experiments but the first one.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730689505415

thread_type:
official_review

<thread_object>
creator: reviewer_pctm

<content>
summary:
This paper aims to address a noticeable gap in research concerning the amenability of the Mamba model to popular fine-tuning frameworks such as PEFT and MPFT. Utilizing theoretical insights from dynamical systems, the paper demonstrates that minor input variations within the SSM layer of either model do not lead to outputs that deviate exponentially. This theoretical assertion is further validated by empirical experiments.

strengths:
1. The paper is well-written and easy to follow.
2. The paper provides a theoretical analysis to support experimental performance.
3. Comprehensive experiments compare the Mamba model with Transformer-based models across different scenarios.

weaknesses:
Though the theoretical analysis of this paper is really solid, since most researchers nowadays use MPFT and PEFT for Mamba, the contribution of this paper may be limited. It may be better if some weaknesses about MPFT and PEFT for Mamba are figured out and then optimized accordingly.

questions:
1. The introduction could be improved by adding a summary of contributions to provide clearer insight into the paper's impact and contribution.
2. Is it appropriate to set the same learning rate for different models (e.g. Mamba and Pythia) when comparing training stability? Whether it is possible that a better-tuned Transformer model could demonstrate enhanced stability and performance, suggesting that the comparison might benefit from adjustments in learning rate settings. Conducting an ablation study or sensitivity analysis on the learning rates for different models maybe can help.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732433373699

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Updated results and manuscript

comment:
We thank all reviewers for their feedback and time reviewing our paper.  We have updated the manuscript to reflect reviewer comments, including the following:
- Ablation results demonstrating the regularization implications of Theorem 3 and effects on improved performance generalization.
- Expanded divergence results, including additional fine-tuning datasets and benchmarks.  Agreeing with Reviewer's tWTf intuition, these extensive results further empirically confirm the paper's main theoretical results showing Mamba SSM's are stable learners, while also empirically demonstrating Transformer-based LLMs can be highly unstable when training using MPFT and PEFT.
- Writing changes: summary of major contributions was added to the introduction, shifting of the background section (which previously contained extra contextualization) to make way for more detailed empirical analysis in the main text, and reference fixes.  We thank the reviewers for these suggestions, we believe such changes have greatly improved the structure of the submission.
- Standard deviations to accompany mean divergence results, and raw performance numbers to accompany ICL results
- A study of the stability with respect to instruct tuning hyperparameters (Appendix H), which demonstrates Mamba models are far more stable to variations in such hyperparameters than Transformer LLMs.

Below, we also address general concerns from reviews.  We look forward to any additional feedback, please let us know if any points need further clarification or discussion.
# Theoretical results are shallow: Lemma 1 is trivial
We respectfully disagree.  The stability of recurrent neural models is not a given, with prior works promoting adjustment of the training schedule to combat unstable dynamics of such recurrent models [1, 2].  By theoretically showing Mamba SSMs are guaranteed to to be stable (Theorems 1 and 2), Mamba LLMs may directly use standard Transformer-based training frameworks (e.g., MPFT and PEFT) without adjusting for potential instabilities, as demonstrated at length in the paper.  

Furthermore, calculating the lypanuv exponents of a recurrent model is nontrivial, i.e., we refer the reviewer to the proof of Lemma 1, which requires careful analysis of Mamba's state-space equations, variable activations, and variable domains.  Finally, we note that other reviews have highlighted the strength of the paper's theoretical results.

# No practical demonstration of Theorem 3
We have added new experiments (Section 4.2 of the updated manuscript) demonstrating the implications of Theorem 3 (i.e., parameter tying of time-varying parameters via LoRA's target layer).  The new results demonstrate that targeting the weight parameter described by Theorem 3 near uniformly results in better performance than only targeting other layers, while also showing that **only targeting this weight parameter outperforms targeting all linear layers using LoRA on a majority of tasks**.  The latter thus demonstrates the regularization implications of Theorem 3, further discussed at length in the updated submission.

# Recent Mamba MPFT works/MPFT and PEFT weaknesses for Mamba models
We thank the reviewer for raising this interesting point and their suggestion to find Mamba weaknesses for MPFT and PEFT.  As noted in the review, recent works [3, 4] have begun fine-tuning Mamba LLMs using mixed-precision.  We believe such decisions can only be strengthened when backed by strong theoretical results (such as the stability theorems featured in the submission).  

We also note that **the paper highlights a weakness of the most widely used LLM architecture, i.e., that Transformer-based LLMs can be highly unstable when using some of the most common fine-tuning frameworks (e.g., MPFT and PEFT)**.  As a potential solution to such instability, the paper demonstrates MPFT and PEFT may be used with Mamba models without this weakness, and fine-tuned Mamba LLMs can demonstrate comparable emergent abilities (i.e., instruction tuned ICL) compared to Transformer-based LLMs.

# References
[1] Pascanu, R. "On the difficulty of training recurrent neural networks." arXiv preprint arXiv:1211.5063 (2013).

[2] Mikhaeil, Jonas, Zahra Monfared, and Daniel Durstewitz. "On the difficulty of learning chaotic dynamics with RNNs." Advances in Neural Information Processing Systems 35 (2022): 11297-11312.

[3] Wang, Junxiong, et al. "The Mamba in the Llama: Distilling and Accelerating Hybrid Models." The Thirty-eighth Annual Conference on Neural Information Processing Systems.

[4] Waleffe, Roger, et al. "An Empirical Study of Mamba-based Language Models." arXiv preprint arXiv:2406.07887 (2024).
</content>
</thread_object>

</thread>
<thread>
creation date:
1733259013631

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Comment from the authors

comment:
Dear all,

The authors would like to thank Reviewers Pctm and tWTf for their constructive criticisms.  Based on this feedback, we have added a large number of experiments extensively validating Theorems 1 and 2 (spanning both different instruction tuning datasets and benchmarks), further highlighted the instability of Transformer-based LLMs for MPFT+PEFT in contrast to Mamba-1/2 LLMs, added hyperparameter search experiments demonstrating that Mamba LLMs are significantly more stable than Transformer LLMs for this paradigm, streamlined the writing of the paper, and added ablation experiments clearly demonstrating the generalization impact of Theorem 3.  Overall, these mark significant improvements to the paper, and we thank the reviewers again for their positive discussions.

To Reviewer jrwp: the reviewer has revealed their identity to the authors (without the authors' consent), requested additional experiments then admitted they never intended to positively adjust their review, repeatedly ignored author responses to their criticisms, and repeatedly stated incorrect claims which are not applicable to our results.  We thank the reviewer for stating Theorem 3 had no practical demonstration, and note these additional experiments were described in this [comment](https://openreview.net/forum?id=i9RTCC6whL&noteId=CsLrDeXS3Q).  We address your remaining criticisms below:

> Let’s be direct. I could simply choose not to respond, rate this submission poorly, and likely see it rejected.

Yes, thank you for not doing this.

> First, the stability analysis builds on well-known facts about RNN models, specifically that their decay terms are bounded by 1... they essentially restate the intuitive property that the exponential decay in linear RNNs...

This statement is incorrect.  The reviewer has stated they base their stability claims on [1].  [1] proves stability of linear *time-invariant* RNNs.  For such RNNs, [1] uses exponentiation of the time-invariant $A$ matrix during unrolling  (i.e., $\sum_j^{k-1} A^j$) to subsequently derive their stability results (Equation 4 of [[1]](https://arxiv.org/pdf/2303.06349).  In contrast, when the linear RNN is time-varying as in Mamba, i.e., $A_j$, unrolling does not maintain this exponentiation (instead, we have $\sum_{j=0}^{k-1} \prod_{i=j}^{k-1} A_i$, as described in Section 3.1 of [[2]](https://arxiv.org/pdf/2405.21060)).  The decay terms you mention are derived through eigenvalue analysis owing to the linear time-invariance of the RNNs in [1] (i.e., $\sum_j^{k-1} A^j = \sum_j^{k-1} P \Lambda^j P^{-1}$, where $A = P \Lambda P^{-1}$).  **This property does not hold for linear time varying Mamba SSMs, thus these decay terms do not exist for Mamba SSMs, and a stronger notion of stability is required, i.e., Lyapunov stability** (a cornerstone of stability analysis for linear dynamic systems [3,4]).  **Thus, our theoretical results are non-trivial and novel.**

> Second, I am not convinced that numerical stability is a significant issue to address here. 

Thank you, we agree.  Indeed, this is the central thesis of the paper: **Mamba SSMs are naturally stable to numerical deviations.  Strong theory supports this, and extensive empirical evidence demonstrates this for MPFT/PEFT, in sharp contrast to Transformer LLMs.**

> To my knowledge, only matrix ( A ) needs to be stored in fp32, while all other linear projection parameters can safely use bf16—a method I’ve applied in practice without issues. I believe Mamba's official repository suggests is for the same reasons

This is incorrect, **the official Mamba repo states:**
> Our models were trained using PyTorch AMP for mixed precision. AMP keeps model parameters in float32 and casts to half precision when necessary... **SSMs are sensitive to their recurrent dynamics. If you are experiencing instabilities, as a first step please try a framework storing parameters in fp32**

Indeed, the last two sentences motivated us to theoretically characterize exactly how sensitive Mamba SSMs are to their recurrent dynamics.  Finally, the authors cannot verify the following: "a method I’ve applied in practice without issues."  Please provide citable, peer-reviewed evidence as the basis for criticisms.  

# References
[1] Orvieto, Antonio, et al. "Resurrecting recurrent neural networks for long sequences." International Conference on Machine Learning. PMLR, 2023.  Arxiv link: https://arxiv.org/pdf/2303.06349

[2] Dao, Tri, and Albert Gu. "Transformers are SSMs: Generalized models and efficient algorithms through structured state space duality." arXiv preprint arXiv:2405.21060 (2024).

[3] Mikhaeil, Jonas, Zahra Monfared, and Daniel Durstewitz. "On the difficulty of learning chaotic dynamics with RNNs." Advances in Neural Information Processing Systems 35 (2022): 11297-11312.

[4] Yang, Lujie, et al. "Lyapunov-stable Neural Control for State and Output Feedback: A Novel Formulation." Forty-first International Conference on Machine Learning.
</content>
</thread_object>

</thread>


Metareview:
metareview:
Reviewer find the paper to have some interesting hypothesis on stability of Mamba models but expressed strong concerns over the choice of datasets and strength of theoretical results. Authors addressed some of these concerns by adding more experimental results, however reviewers are not fully convinced about the strength of the results. Overall I think this paper requires another round of rewriting to clearly present the hypothesis and the evidence. Hence I recommend rejection.