Mamba Modulation: On the Length Generalization of Mamba Models

First, we thank the reviewer for acknowledging the relevant and vibrant topic of our work, that it is written well and structured and the presented method is simple and clear. We also acknowledge the various points of concern and questions they have raised, which we hope will be resolved with the following response.

W1. Results in the paper are trivial and the significance seems very limited.

A1: We understand these concerns, but highlight that our method offers significant/consistent improvements over the original pre-trained models.

We posit that generalization to longer inputs remains a challenge for Mamba models. Our proposed scaling method demonstrates improvements under such conditions, offering evidence to support our theorems. This insight is also novel, both explaining why existing methods fall short in extrapolating and providing a principled criterion for designing more resilient architectures.

W2. The paper does not discuss known results on other architectures or distinguishes the results and the novelty brought forth in the paper.

A2: Mamba models present a promising alternative to the well-established Transformer architecture. However, despite their rising popularity, they remain comparatively under-explored, particularly when it comes to understanding their limitations.

We address this gap by uncovering the fundamental causes behind Mamba’s failure to generalize across sequence lengths, which previously has not been explained. Our analysis exposes key structural differences from self-attention mechanisms that contribute to this behaviour, providing valuable insight that deepens the field’s understanding and paves the way for more robust model design.

W3. Related work is very lacking - missing works on extrapolation and length generalization in different setups.

A: We will add and discuss the mentioned references regarding length generalization in transformers and RNNs in the final version of our paper as follows.

The literature on transformers highlights positional encoding and attention biases as critical for length extrapolation: Press et al. introduced ALiBi (Attention with Linear Biases), demonstrating that linear distance-based attention biases enable effective extrapolation beyond training lengths. Zhang et al. used synthetic reasoning tasks to analyze transformers' generalization capabilities, revealing how attention patterns and pretraining affect robustness to longer sequences. Emami et al. analyzed the implicit bias of linear RNNs in sequence modeling tasks. Cohen-Karlik et al. investigated gradient descent dynamics for temporal extrapolation and showed how overparameterized RNNs learn low-dimensional state spaces.

While Mamba's efficiency is well-established, its extrapolation behavior requires systematic comparison to transformer approaches and integration of insights from RNN theoretical analyses.

1. Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation. Press et al. ICLR 2022.
2. Unveiling transformers with lego: a synthetic reasoning task. Zhang et al. Arxiv 2022.
3. Implicit Bias of Linear RNNs. Emami et al. ICML 2021.
4. On the implicit bias of gradient descent for temporal extrapolation. Cohen-Karlik et al. AISTATS 2022.
5. Learning Low Dimensional State Spaces with Overparameterized Recurrent Neural Nets. Cohen-Karlik et al. ICLR 2023.

Q1. Can you please provide grounds for [sentence in lines 30-33]?

A: We acknowledge some potential confusion that may have stemmed from a poor choice of phrasing.

The primary motivation of Mamba is with regards to the ability to have models address the computational inefficiency of models on long sequences, which is shown by results in the original paper. However, Mamba does claim extrapolation abilities, namely on synthetic tasks such as selective copying and induction heads, but the same experiments have not been tested on language data, which is why we believe this claim prompts further investigation.

Q2. What is s in line 151? it isn't defined before it is used.

A: ‘s’ simply appears after its definition: ‘scalar values’, refers to the scaling factors used by Azizi et al. to multiply with $\Delta_t$.

Q3. Figure 1 isn't clear from its description. What does it mean that certain layers have fast decay while others decay slowly?

A: Figure 1 visualizes the spectrum of each layer to motivate our theoretical analysis. It reveals a consistent pattern across layers: the presence of both large (near 1) and small (near 0) eigenvalues. This is not incidental; it underpins Section 4.2, where the convergence behaviour is shown to depend critically on the spectral distribution. If the largest eigenvalues of different states in a layer are concentrated close to zero, then the information contained in the states of this layer decays fast. Conversely, layers with slower spectral decay (i.e., eigenvalues closer to 1) preserve state information for longer durations.

Q4. Lemma 4.1. is trivial and does not require a dedicated proof.

A: This Lemma ensures that we do not need to impose any explicit assumption on the norm of the input. By presenting it as a lemma, we establish a clear and self-contained result that supports downstream arguments without relying on unstated constraints.

Q5. In Theorem 4.2, Corollary 4.3 & Corollary 4.4 - why is the assumption on the eigenvalue distribution different per statement?

A: The assumptions in Corollary 4.3 and Corollary 4.4 are specific instantiations of the more general condition stated in Theorem 4.2. We provided detailed derivations in Appendix A.3, showing how each corollary follows from the general theorem under simplified eigenvalue distributions.

In particular, the convergence behaviour for the Mamba state is obtained by taking the limit $\lambda_{\min}⁡ \to 0$. For Mamba2, we consider the regime where $\lambda_{\max} \to \lambda_{\min}$​. These are chosen to reflect the parameterizations used in each model. We will clarify this by presenting the assumptions in formal mathematical notation and revise the manuscript accordingly.

Limitations: The empirical benefit of scaling instead of isn't clear. [...] It would be good to discuss [...] the tradeoffs and benefits of the different scaling methods.

A: In future revisions, we will compare the tradeoffs of different scaling strategies and clarify their benefits, especially in comparison to MambaExtend. Corollaries 4.3/4.4 suggest that both eigenvalue adjustment and discretization step scaling can influence the convergence rate. While performance can vary depending on the model and sequence length, our results indicate that the scaling method becomes increasingly important as sequence length grows. We posit that scaling $A$ can be a dominant factor for extremely long sequences, where numerical stability and convergence speed become critical.

Paper Formatting Concerns: The paper has multiple flaws in the form of unclear sentences and other formatting issues, below are a few:

A: Thank you for the detailed feedback, we will fix these typos in the final version.

We appreciate that the author acknowledges our work in terms of the evidence we provide for our claims are convincing, that our paper is solid and the actionable nature of our ideas. We also appreciate the various questions and suggestions they make, which we hope are resolved through the rebuttal that follows.

C1: I think the paper could have been written in a much more succinct manner. Do you think this long paper format of NeurIPS makes most sense?

A1: We understand the concern regarding the length and format, and we would like to take this opportunity to explain why we believe the current 9-page version is appropriate and necessary to fully communicate the significance of our contribution.

Shorter page limits may reduce persuasiveness. There is a fine line between communicating an idea succinctly/effectively, while also providing sufficient justification through analysis and experiments. Thus, it is difficult to conclude regarding a shorter page limit, as we wish to maintain the same level of detail and persuasiveness.

A heavier appendix may reduce readability. We already carefully offload proofs and more ancillary experimental and implementation details to appendix sections, allowing the main text to stay focused and readable. Condensing more main content could risk shifting essential arguments into the appendix, impacting overall readability and accessibility.

Additional analysis provided. We appreciate your feedback and would like to add more analysis to enhance our work. We can provide the norm analysis of the ssm_state to enhance section 4.1 by conducting more experiments to record the trend of the ssm_state as the input sequence length increases as we replied to reviewer up1U. We randomly sampled data from ProofPile and PG19, and measured the maximum difference between the largest and smallest ssm_state norms across all layers of the Mamba2-1.3B model. The results, presented in the following table, clearly demonstrate a divergent trend in the ssm_state norm as sequence length increases. This empirical observation is consistent with our theoretical prediction and further validates our analysis.

We feel the depth of analysis offered by the current format aligns well with the standards and expectations of the NeurIPS long paper. At the same time, we are also happy to discuss and willing to make changes based on this interaction period.

C2: Why does section 4.1 exist, what’s the point of that analysis?

A2: Section 4.1 motivates the theoretical developments in Section 4.2 by examining the numerical structure of the Mamba transition matrix.

Section 4.1 provides essential empirical motivation for the theory in Section 4.2. The heatmap in Figure 1 reveals a consistent spectral pattern across layers: the coexistence of both large (near 1) and small (near 0) eigenvalues. This observation is crucial, as it sets the foundation for the theoretical analysis that follows—namely, the identification of two distinct dynamical regimes that govern the behavior of Mamba states.

Without this empirical grounding, the theory would lack connection to real model behavior. As mentioned above, we will add the state norm analysis to further ground the following theoretical section 4.2. This will make the connection between empirical trends and theoretical insights even more concrete in the final version of the paper.

C3: Although 6.1 is very convincing, just maybe add a remark about why there could still be cases like GovReport Mamba-1.4B where the perplexity of scaling Delta is better. Acknowledge it or just share intuition.

A3: One way of interpreting this behavior is that at shorter lengths (ex. <= 16k), simply being able to scale either factor is sufficient, as the differences in perplexity are marginal. This may be because at these lengths, the compounded decay may approximate rather well; however, as the lengths keep increasing, this eventually leads to a significant difference between them, which may be the case of why at shorter lengths both methods appear to work rather well.

On one note, we can observe that on Mamba (first generation) models, scaling $\Delta$ usually works slightly better on shorter lengths; we hypothesize that this might be due to how the $A$ matrix is constructed as a diagonal matrix rather than a scalar-times-identity like in Mamba2, which may make scaling $A$ with a single layer-specific scaling factor more appropriate in that scenario.

One specific outlier is on Mamba-1.4b with GovReport, at a sequence length of 32k; scaling $\Delta$ for this example works significantly better than $A$ but it is the only case. We believe that this is more of a point of acknowledgement given the lack of other settings where this occurs (either with 1.4B models, first-generation Mamba models or just with other models on GovReport) and will mention this accordingly.

C4: It seems like initially controlling for $\Delta$ starts out as being better, and then controlling for $A$ becomes better. Is there any scope of controlling for both, or is that stupid?

A4: There are manners in which this is possible; for example, if one is to introduce scaling parameters for each factor and train them in the same manner as we do, then by all accounts it is possible to control them simultaneously.

Characteristics of zeroth-order optimization: One reason why this isn’t done explicitly is the use of zeroth-order optimization, which is known to become more unstable as the number of parameters that need to be tuned increases. Alternatively, one way would be to directly fine-tune everything through gradient descent, but this is a much more time- and resource-consuming process compared to the proposed design.

Multi-stage optimization: Another method would be to conduct a type of multi-stage optimization, where perhaps one subset of scaling factors is tuned, followed by another. However, there are some known potential issues with such types of methods in deep learning research that make this somewhat difficult.

C5: MambaExtend is by Azizi et al, that wasn’t clear from the main paper when it was mentioned, maybe add the citation in the section as well, instead of it just being in the Appendix.

A5: We thank the reviewer for raising this detail; we did mention Azizi et al. at multiple times within the preceding sections but did not explicitly mention MambaExtend; we will make this very clear by mentioning them together.

We thank the reviewer for the time and effort they have placed towards providing a thorough evaluation of our work, particularly their comments about the problem being important, the experiment being solid, and both the ideas and results being insightful. We are also appreciative of the weaknesses that have been raised and hope the following additional details/results will clarify these points and resolve these existing doubts about work.

C1: There is no direct experimental validation of the state norm behavior throughout the paper, which raises questions about the validity of the theoretical analysis.

A1: We can fully address this concern by conducting additional experiments that track the behavior of the ssm_state as the input sequence length increases.

Specifically, we randomly sampled data from ProofPile and PG19, and measured the maximum difference between the largest and smallest ssm_state norms across all layers of the Mamba2-1.3B model. The results, presented in the following table, clearly demonstrate a divergent trend in the ssm_statenorm as sequence length increases. This empirical observation is consistent with our theoretical prediction and further validates our analysis.

We appreciate your suggestion and will incorporate this analysis into the final version of the paper, as we agree it provides a valuable supplement that reinforces the strength of our contributions.

C2: Prior works, such as DeciMamba, have also proposed explanations for the failure of Mamba models in handling length extrapolation. This paper does not sufficiently argue that its explanation is superior to those provided in prior work, nor does it point out any flaws or inaccuracies in the existing explanations. This weakens the persuasiveness of the theoretical analysis.

A2: Thanks for your feedback. While prior works, including DeciMamba, LongMamba, and MambaExtend, provided useful explanations for Mamba’s failure in long-sequence generalization, we believe our theoretical analysis offers a more comprehensive and principled understanding.

Specifically, our contributions extend existing explanations in three key ways:

• New insight – divergent states due to large eigenvalues: Our analysis reveals that large eigenvalues can lead to the tendency of exploding state norm, introducing instability that is not captured by existing explanations. This insight identifies an additional failure mode beyond vanishing memory and provides a more complete theoretical account of why Mamba models struggle with length extrapolation.
• Consistency with prior insights: Our framework also recovers and formalizes the core idea shared by previous works: that small eigenvalues, when combined with accumulated discretization steps $\sum\Delta_{t}$, cause the decay term $\exp(-A \sum \Delta_t)$ to vanish, resulting in fast memory loss. In this sense, our theory subsumes and strengthens existing intuitions about the vanishing states.
• Unified and explanatory framework: Through Corollaries 4.3 and 4.4, we show how modifying the discretization step $\Delta_{t}$ helps modulate convergence rates. However, we further clarify that such adjustments do not address the root cause for extremely long inputs, which lies in the spectral properties of the state transition matrix A. Our framework provides a principled explanation for why these heuristics sometimes succeed—and when they may fall short.

In DeciMamba, the claims are quite complementary to our analysis. For example, they also make the claim about the cumulative product of the per-timestep transitions converging to 0; they do not provide any additional intuition as to why this may be the case. Our analysis, meanwhile, delves deeper into this phenomenon. While we don’t provide a subjective quantitative measure as they do, our inspection of the eigenvalue spectrum is meant as a way of highlighting that our claims have grounding and serves to motivate the rest of our analysis in Section 4. We hope this clarification addresses your concern and strengthens the persuasiveness of our theoretical contribution.

Thank you for your valuable feedback. We greatly appreciate your recognition of our framework's efficiency, good reproducibility and competitive results. We nevertheless appreciate the additional points that have been mentioned as limiting factors and hope that the following response will to address these follows:

C0: This scaling of A is training-free (using pre-trained weights directly) without updating model weights (but it actually tunes A values through calibration).

A0: To clarify, our method introduces a lightweight set of scaling factors applied to the pre-trained model; the number of scaling factors is smaller than the number of A. While it is possible to tune the entire A, our method and experiments avoid doing so. The scaling set functions as a plug-and-play module that can be activated only when the input sequence length exceeds the model's original maximum context length during inference.

C1: The contributions are limited : scaling a different parameter in the same training-free calibration framework has only subtle differences from MambaExtend.

A1: We would like to highlight our contribution and distinguish our work from previous studies.

Novel Theoretical Analysis. We are the first to provide a theoretical explanation for Mamba’s failure to generalize to long sequences, identifying the root cause as the divergent growth of the state norm due to spectral properties of the state transition matrix. This insight is not only novel but also foundational, as it enables a deeper understanding of the underlying dynamics rather than relying on heuristic modifications.

Strong Interpretability. While we recognize the simplicity and practicality of the scaling methods in MambaExtend, our approach is theoretically motivated and analytically grounded. In particular, our analysis explains why certain empirical scaling factors (e.g., those greater than 1) observed in MambaExtend emerge despite appearing counterintuitive under their original motivation.

C2: The theoretical analysis assumes uniform, diagonal spectra, limiting real-world applicability.

A2: We could fully clarify your concerns regarding the uniform and diagonal assumption.

1. The premise that our theoretical analysis assumes diagonal spectra is a point of misunderstanding. In fact, the diagonality is not a simplifying assumption—it is a structural characteristic of Mamba and Mamba-2. Their state transition matrices are inherently diagonal, which means the spectrum and the matrices of the form $\text{diag}(\exp(-A))$ naturally share identical entries. We don’t impose any additional assumptions on them.

2. Our choice to model the spectrum using a uniform distribution is empirically grounded. As shown in Section 4.1, the eigenvalues of pretrained Mamba models consistently span the range from 0 to 1 across all layers. This observed behavior supports our decision to use a uniform distribution as a practical and realistic approximation.

3. The divergent convergence behavior is irrespective of spectral assumptions. We appreciate your rigorous comment and accordingly extend our analysis to accommodate a more generalized case without imposing any strict assumptions on the distribution of $\lambda$. Notably, the asymptotic behavior of the system remains unchanged.

Elaboration on Theorem 4.2: Expected State Norm under General Spectral Distribution

Assume the transition matrix $\mathbf{\Lambda} = \mathrm{diag}(\lambda_1, \dots, \lambda_d)$ has diagonal entries $\lambda_i \in (0,1)$ drawn i.i.d. from a probability density function $p(\lambda)$ supported on $[\lambda_{\min}, \lambda_{\max}] \subset (0,1)$.

Suppose the system evolves as: $\mathbf{h_t} = \mathbf{\Lambda} \mathbf{h}_{t-1} + \mathbf{B} \mathbf{x}_t,$ where $\mathbf{x}_t \sim \mathcal{N}(0, I)$ and each row of $\mathbf{B}$ is independently sampled as $\mathbf{b} \sim \mathcal{N}\left(0, \frac{1}{\sqrt{d}} \mathbf{I}\right)$.

Then, as $t \to \infty$, the expected squared norm of the hidden state converges to $\boxed{ \mathbb{E}[\|\mathbf{h_\infty}\|^2] = \mathbb{E}[\|\mathbf{B} \mathbf{x}\|^2] \cdot \int_{\lambda_{\min}}^{\lambda_{\max}} \frac{p(\lambda)}{1 - \lambda^2} \, d\lambda. }$

Asymptotic Analysis of $\mathbb{E}[\|\mathbf{h}_\infty\|^2]$

Define $I(\lambda_{\min}, \lambda_{\max}) := \int_{\lambda_{\min}}^{\lambda_{\max}} \frac{p(\lambda)}{1 - \lambda^2} \, d\lambda.$

We analyze the behavior of $I(\lambda_{\min}, \lambda_{\max})$ as $\lambda_{\max}$ approaches the edges of the domain.

1. Behavior as $\lambda_{\max} \to 1^{-}$:

Near $\lambda = 1$: , we have $1 - \lambda^2 \approx 2 (1 - \lambda)$ via Taylor expansion.

If $p(\lambda)$ is continuous at $\lambda=1$ and $p(1) > 0$, then for $\lambda_{\max}$ close to 1, $I(\lambda_{\min}, \lambda_{\max}) \approx \frac{p(1)}{2} \int_{\lambda_{\min}}^{\lambda_{\max}} \frac{1}{1 - \lambda} \, d\lambda = \frac{p(1)}{2} \log\left(\frac{1 - \lambda_{\min}}{1 - \lambda_{\max}}\right).$

As $\lambda_{\max} \to 1^{-}$, $\log(1 - \lambda_{\max}) \to -\infty$, so $\boxed{I(\lambda_{\min}, \lambda_{\max}) \sim -\frac{p(1)}{2} \log(1 - \lambda_{\max}) \to +\infty.}$

The expected state norm diverges logarithmically as the largest eigenvalue approaches 1.

2. Behavior as $\lambda_{\max} \to 0^{+}$

When $\lambda_{\max}$ is close to 0, $\frac{1}{1 - \lambda^2} \approx 1,$ since $\lambda^2$ is very small.

If $p(\lambda)$ is continuous near zero with $p(0) > 0$, then $I(\lambda_{\min}, \lambda_{\max}) \approx \int_{\lambda_{\min}}^{\lambda_{\max}} p(\lambda) \, d\lambda \approx p(0) \cdot \lambda_{\max}.$ Hence, $\boxed{I(\lambda_{\min}, \lambda_{\max}) \to 0 \quad \text{linearly as } \lambda_{\max} \to 0^{+}}$.

The expected state norm vanishes linearly with $\lambda_{\max}$ approaching zero. The expanded analysis still reveals the inherent tendencies of Mamba’s state dynamics toward explosion or vanishing, irrespective of spectral assumptions.

C3: It is not clear how section 4.2 directly leads to the motivation of controlling A only.

A3: The eigenvalues of the transition matrix is $diag(\exp(-A))$ , thus we can control the spectrum distribution by applying a scaling factor to control the magnitude of A accordingly.

C4: Using one scalar per layer coarsely addresses spectral outliers.

A4: For fair comparison, we follow the default setting of MambaExtend for both language modeling and NIAH experiments. The experiments on the language modeling task are applied with per-layer-wise scaling. We could perform state-wise scaling accordingly; however, this will lead to more computation compared to MambaExtend methods, as it will introduce more trainable parameters.

On one note, the NIAH results do use state-wise scaling for each layer for MambaExtend and ours, with the end results similar to what we observe on other datasets. We contend that this serves as additional evidence to support the proposed method of scaling $A$ as opposed to $\Delta$ is valid.

C5: Does this calibration method actually perform better than fine-tuning only A or t values (with other parameters fixed) directly to long context length? What happens if combining both t and A scaling?

A5: We note that the goal of using zeroth order optimization method here is intended to make calibrating the model significantly simpler and faster, as it only requires forward passes and no explicit gradient. Scaling both $A$ and $\Delta_t$ simultaneously can possibly lead to better results, but at the cost of greater complexity and potential instability as zeroth-order optimization often relies on tuning a small set of parameters.

Furthermore, $\Delta_t$ is input-dependent, not a model parameter, and therefore cannot be directly fine-tuned. We manage to add fine-tuning $A$ directly and calibrate full $A$ for language modeling on PG19 with Mamba models. Enable more tuning freedom with additional scaling factors can improve the performance accordingly.

The following is a full fine-tuning of only $A$ for the NIAH task, which we observe does not significantly deviate from directly scaling $A$ but remains nearly exactly comparable to our method of simply adding scaling parameters. Here, X indicates the model correctly solved all examples.

Mamba-1.4B

We use the same evaluation criteria as described within the manuscript for all additional experiments.

Q: Figure 2 legend color should be corrected.

A: Thank you for this note; the legend was intended to show that the changing $A$ scales were represented by different colors while the scaling of $\Delta_t$ used different line styles. We can understand where the confusion came about and will fix this issue.