Thank you very much for your precious feedback.

Structure regarding the Experiments section

We agree that having a separate “Experiments” section is a bit strange. We will consider renaming Section 5 (“Experiments”) to “Forget and Recall Thresholds” to have a better parallel between this section and Sections 3 and 4.

Typo in line 277.

Thank you very much for this catch. “The second plot of” should be removed.

Possible Explanation for https://arxiv.org/pdf/2410.04727

Yes! Our findings can be viewed as a possible explanation for the observation in that paper. In fact, this work started exactly because we observed the same phenomenon as that paper.

If there are any other aspects of our paper that you believe could be improved, please let us know.

Thank you very much for your precious review. We would like to make the following response.

Viable and Practical Solutions to the Tradeoff Between Short-Context and Long-Context Performance.

There is indeed an inherent trade-off between short-context performance and long-context generalization. This is a fundamental research question that is outside the scope of our work. We believe further improvements to the architecture and training approach can mitigate the issue. Interestingly, a similar (though not exactly the same) trade-off has also been observed with decoder-only Transformer models as well [1].

Pretraining beyond the identified Forget Threshold is computationally impractical.

To pretrain with the Forget Threshold identified in our paper is indeed a great challenge. Some techniques that we think are promising include recomputation, gradient truncation, sequence parallelism, gradient offloading, zero-order optimization, etc. Moreover, architectural improvements may also alleviate the need for very long contexts, such as designing update rules that are better at learning robust forgetting.

Sliding Window Mamba is Not Compelling.

The main purpose of the proposed sliding window method is to show that inducing forgetting can mitigate length generalization failure. Nonetheless, we would argue that this method has some merits compared to sliding window attention (SWA).

• Firstly, it can be applied to any Mamba models (or other models whose recurrent state can be reformulated as a weighted sum) without additional training, even though the model was pretrained without a sliding window. In contrast, attention-based models that were not trained with SWA either require additional training or require modification to attention patterns and position embeddings to turn into a SWA model.
• Secondly, the recurrent state may be considerably smaller than the KV cache in SWA, which means the model may have a lower memory cost.

Extensively comparing these methods is outside the scope of our work.

The Lack of a Baseline: Context Length Extension by Post-Training

You are raising a very important point and we agree that such a baseline is very important. However, we found that directly training on longer sequences require large amounts of data for the model to adapt to the next context length. This can be partly explained by Figure 3. It shows that, after sufficient training, the majority of SSM heads in Mamba-2 have a memory retention strength ($\alpha_t$) whose cumulative products are very close to 0, resulting in vanishing gradient. This causes the model to struggle to learn to model long-distance dependencies. How to improve the data-efficiency in long-context extension for Mamba-2 remains an important future research direction.

References

[1] Buckman and Gelada. 2024. https://manifestai.com/articles/compute-optimal-context-size/

Thank you very much for the highly detailed review. We will try to address your remarks one by one. Bold text indicates "reason to reject"/question mentioned by you, and the text following each bold text is the corresponsing response.

Delta Rule.

The link you provided is not accessible for us. Can you please share an accessible link?

Nonetheless, as you mentioned, improving Mamba-2 is valuable since it has been applied to several LLMs and has important practical values.

Contradicting statements.

We respectfully disagree with your claim that our paper contains contradicting statements. In response to the example you provided, we believe you have misunderstood the context of those two statements. Statement (1) describes the scenario when the training context lengths are too short for the state size (i.e., $T_\text{train} < T_\text{forget}$), which causes overfitting. Statement (2) describes the scenario when we increase the training context length beyond the Forget Threshold (i.e., $T_\text{train} \ge T_\text{forget}$), in which case overfitting does not occur. These two statements describe different scenarios, which do not occur at the same time. Thus, they are not contradicting statements.

In response to your remark that “There are also statements which require verification and missing details”, can you kindly provide examples? We would be very glad to provide additional details/verification measures.

Real applications of the phenomena.

The findings are meant to provide insights for designing better Mamba-based models for long-context modeling. Thus, the application value of this paper lies in the takeaways, some important ones are:

• To train a Mamba-2 with robust length generalization, one should use training lengths that grow linearly with the state size.
• Length generalization failure can be caused by the lack of forgetting, instead of the lack of memory retention, as claimed in [1].

Mamba-2’s learning process should not depend on the training context length.

We respectfully disagree. The learning dynamics of Mamba-2 depend on the training context length because the output hidden representations at any layer and time step (i.e., $y_t$) depend on the input hidden representations in all previous time steps (i.e., $[u_1, \cdots, u_t]$). Thus, the training context length will affect the loss function, which influences the learned neural weights. In other words, the learned model weights (i.e., $A$, $W_x$, $W_\Delta$, $W_B$, etc. in Appendix A) are influenced by the training context length. Additionally, the time-varying variables (i.e., $\alpha_t$, $\overline B_t$ and $x_t$ in Eq. 2) are produced by these model weights, so they are also influenced by the training context length. Hence, the update rule and forgetting mechanism in Mamba-2 are influenced by the context length.

Mamba-2 has a stable exponential decay (converging to a constant value when variables are fixed).

The proof is as follows:

Assuming fixed variables, we have $\alpha_i = \alpha, \overline B_i = \overline B$ and $x_i = x$ for all $i=1,2,\cdots$, thus:

$$h_t = \overline B x + \alpha \overline B x + \alpha ^ 2 \overline B x + \cdots + \alpha^{t-1}\overline Bx=\left(\sum_{i=0}^{t-1} \alpha^i \right)\overline B x \underset{t\rightarrow \infty}{=} \frac{1}{1-\alpha}\overline Bx$$

QED.

This proof will be added to the revised paper.

Mamba’s state-space equations’ influence on forgetting.

Can you kindly clarify why Mamba’s state-space equations would undermine the observation that context length affects the forgetting mechanism? As mentioned above, the learned weights are influenced by the context length; the variables in the update rule ($\alpha_t, \overline B, x_t$ in Eq. 2) are influenced by the learned weights; and the actual forgetting mechanism of a trained Mamba is influenced by the update rule. Hence, we conclude that the forgetting mechanism in a trained Mamba model is influenced by the context length.

Why cannot $\overline B_i$ be mutually orthogonal?

This is because $\overline B_i$ has $N=128$ dimensions. Thus, when there exists more than $N$ of them ($[\overline B_1, \overline B_2, \cdots, \overline B_T] \text{ where } T > N$), they cannot be mutually orthogonal. This is a basic result in linear algebra. This concurs with our statement in Line 148 because the context length can be greater than $N$.

Remark: We have mistakenly left out the \overline in Line 148. This will be fixed in the future.

References

[1] Ben-Kish. 2024. DeciMamba: Exploring the Length Extrapolation Potential of Mamba

By state size, the authors are referring to the internal Mamba-2 state dimension, correct? E.g., d_state \in {64, 128} for official Mamba-2 checkpoints?

As defined in Section 4.2 and Line 272, the state size is $N_S$, which is the dimensionality of all states (i.e., $h_t$) at any time step. There are $H$ of them in each layer, where $H$ is the number of SSM heads. In this paper, we use M parameters as the measurement unit for $N_S$.

"This is because the amount of information in the context does not increase with its length" <- This contradicts many other statements in the paper

This does not contradict other statements in the paper because this statement is describing the prompts used in the passkey retrieval task, and not the other kinds of contexts. The concrete prompt in passkey retrieval is provided in Appendix B.1. In this prompt, a 5-digit passkey is located inside a large context of repeated texts. When increasing the context length, we just increase the number of repetitions, and there is no additional content. Hence, the amount of information is essentially independent of the context length. This prompt is adopted from the original authors of the passkey retrieval task [2].

References

[2] Mohtashami et al. 2023. Random-Access Infinite Context Length for Transformers.

Why are the retention strengths of earlier tokens always smaller than those of more recent tokens?

The retention strength of the $i$-th token at time step $t$ is defined as $\alpha_{i:t}=\prod_{j=i+1}^t \alpha_j$, since $\alpha_i\in(0, 1)$, we have:

$$\alpha_{i:t}=\alpha_i \cdot \alpha_{i+1:t} < \alpha_{i+1:t}$$

for any $t$ and $i$. Thus, the retention strengths are monotonically increasing. QED.

Remark: In Line 113, we defined “memory retention strength” as $\alpha_t$, which might be misleading due to the similar name, thus, we will change this term in the future.

We will add this proof in the revised paper.

Why are 8K tokens not enough to overload the memory?

Here, we are claiming that overloading the memory would cause the model to learn to produce smaller memory decay terms ($\alpha_t$), as argued in Section 3.1. Moreover, we are stating that the 8K context length is insufficient for the model sizes of the official Mamba-2 checkpoints (i.e, those we investigated in Section 3). In later experiments (Sections 5.1 and 5.2), we showed that the models of such sizes require greater than 8K to learn forgetting. Hence, our experiments support this claim.

Typos in Equation 8

Yes, it is a typo, the $w$ in Line 176 should be replaced with $r$.

The usage of “minimal” versus “negligible”.

Nice catch, we will change it to negligible.

"Similar observation can be found with any model size."

Yes, this is empirically verified, and we are reporting the results for one model size to save space.

"This indicates that the model converges toward more retention and less forgetting" <- What learning schedule was used for pre-training, and how does it compare to the original Mamba-2 recipe?

The training details are reported in Appendix F.1. We will add a reference to this sentence to clarify this. The training recipe is quite similar to the original Mamba-2 recipe, with some slight variations that we believe do not affect our conclusions. Extensive ablation studies on the influence of the training recipe is interesting, but out of the scope of this work.

Figure 8 is empirical evidence, more models seem warranted to support the underlying hypothesis.

We will add the result of Mamba-2 130M and 780M to the Appendix. This phenomenon is observed across different model sizes in our experiments as long as the training context length falls below the Forget Threshold (defined in Section 5).

"To determine whether the model has learned robust forgetting, we feed prompts with 1M tokens to the model and check if the model’s loss exceeds 2× the maximum loss within Ttrain tokens at any point." <- 2x seems very arbitrary

The choice of 2x was determined such that it can reliably capture severe length generalization failure while ensuring that fluctuations are not mistakenly regarded as failures. While we acknowledge that this value is an arbitrary choice, the conclusions are not very sensitive to the actual value (using 1.5x would give exactly the same conclusions). This is because at such huge context lengths, models that are unable to forget earlier tokens have very large loss values, so models with robust length generalization have a stable loss.

"The fact that the amount of information in the tokens exceeds the state’s capacity," <- the experiments are not convincing that this happening

This sentence describes the conclusion from the previous section (Section 4.2), which is empirically validated in Sections 5.1 and 5.2.

"To save cost, we continue pre-training from three official checkpoints of Mamba-2 (130M, 370M, and 780M). They were pre-trained with 8K sequences. The other model configurations (36M, 47M, and 85M) are trained from scratch." <- Where are the training details? As previously stated, if details are in the appendix, please forward reference.

As mentioned in Section 5 (Line 253), the training details are given in Appendix F.

"We can see that for each model size, there is a training length threshold, beyond which the model has much better length extrapolation, which supports our arguments discussed in Section 4.2." <- But why? Is this an emergent ability? I.e., if the x-axis on Figure (10) a-b were further extended, would the 16k and 64k trained models continue to have stable loss?

Yes, this is an emergent behavior when we scale up the training context length, and the reason for this emergence is discussed in Sections 4 and 4.2. We have empirically verified that the models in Figure 10 (a and b) have stable loss up to 10M tokens.

"The rightmost data point in the plot corresponds to Mamba-2 370M" <- What are the other data points?

The other data points correspond to the smaller model sizes listed in the “Models” paragraph in Section 5, which are 36M, 47M, 85M, and 130M.

Thank you very much for your response.

Influence of Context Length on Mamba-2's Update Rule

When we say that Mamba-2's forgetting mechanism is influenced by the training context length, we mean that the data-dependent variables (specifically, the memory decay multiplier $\alpha_t$) are influenced by the training context length. This is true since the value of the neural weights that generate $\alpha_t$ are dependent on the training context length.

As you have suggested, we will make this more explicit in the introduction to avoid confusion.

Disambiguating Statements

We agree that we can disambiguate the paragraph you mentioned further by explicitly stating that the two statements describe two different scenarios. Thus, we will revise this paragraph in the updated paper.

Greater Placement within Existing Works

Regarding RWKV-7, which is a concurrent work, we would like to emphasize that our analysis is limited to recurrent architectures whose update rule can be written as an element-wise affine operator (Mamba-2, Mamba-1, RWKV-6, HGRN-2, etc.). RWKV-7 and DeltaNet are out of the scope of our paper.

LongMamba

Your claim regarding LongMamba is false. The LongMamba [1] paper referred to in the response to Reviewer rF9s25 is not the same as the one appearing in the paper [2] (which is a GitHub repository without any paper).

Other Questions 1

Thank you for these questions. The number of training tokens is 50 times the model size, and we only train for 1 epoch (no training data repetition). We will add these details to the paper.

Other Questions 2

Regarding $\alpha_t$ values for $t \gt 8K$, each $\alpha_t$ is a function of the $t$ input hidden state $u_t$ (please see Eq. 4). Mamba-2 does not have positional encodings, so this is well defined, and we do not make any modifications to the model.

Once again, thank you very much for your precious time. If there are more concerns, we are very happy to address them through further discussions.

References

[1] Ye et al. 2025. LongMamba: Enhancing Mamba's Long Context Capabilities via Training-Free Receptive Field Enlargement

[2] https://github.com/jzhang38/LongMamba

Thank you very much for your time in reviewing.

Novelty of this paper and comparison to existing works.

We first emphasize that LongMamba [1] is a concurrent work because it was published after COLM’s submission deadline. Moreover, our work differs from existing works involving length generalization of Mamba-based models in the following aspects:

• Relationship between memory decay and length generalization: In both DeciMamba and LongMamba, it is argued that too much memory decay (i.e., forgetting too much) is the cause of length generalization failure. In contrast, we show that the failure is caused by too little memory decay (i.e., the inability to forget).
• Length generalization ability: Existing work in length generalization methods, such as DeciMamba [2] and LongMamba [1], can only extrapolate up to 40-50K tokens.
• Relationship between state size and memorable context length: Existing works (e.g., [4]) that investigate the relationship between state size and memory are limited to toy tasks, and we are the first to establish this relationship with language modeling training.

Figures 9 and 11 have too few data points.

We agree that more data points would make the conclusions more convincing, thus, we have conducted additional experiments (with model sizes of 15M, 25M, 55M, and 100M). The new results (combined with previous results) are reported below, and they show that our established relationships/conclusions still hold. These results will be added to the revised paper.

The result of near-perfect retrieval on a 256K context length.

This result is presented in Figure 9 (the rightmost point). “Max Recall Len.” (i.e., $T_\text{recall}$, defined in Section 4.3) refers to the context length at which the model has more than 95% accuracy in passkey retrieval, and the rightmost data point in Figure 9 corresponds to the 370M Mamba-2 model with near-perfect accuracy at 256K context lengths.

Additionally, you mentioned “RoPE extension applied to transformer”. We are unsure about what the concern is. Would you kindly clarify?

More linear attention models beyond Mamba/Mamba2 are highly likely to be included in the discussion.

We are not sure what your concern is. We have presented results involving the length generalization failure of RWKV-6 and HGRN-2 in Appendix H. Our analysis about the importance of forgetting is not specific to Mamba, so it should be generalizable to other kinds of linear attention models. Empirical validation with all linear attention variants is impractical, so we have focused on Mamba (which is the most widely used linear attention variant at the moment).

References:

[1] Ye et al. 2025. LongMamba: Enhancing Mamba's Long Context Capabilities via Training-Free Receptive Field Enlargement

[2] Ben-Kish et al. 2024. DeciMamba: Exploring the Length Extrapolation Potential of Mamba

[3] Yang et al. 2024. Gated Linear Attention Transformers with Hardware-Efficient Training

[4] Arora et al. 2024. Simple linear attention language models balance the recall-throughput tradeoff.

Thank you very much for the response.

Our response to the two concerns is as follows:

1. We are not making a conclusion that contradicts the conclusions from DeciMamba. Instead, we are providing another explanation for the length extrapolation failure. In fact, both too much and too little memory decay can lead to length generalization failure. However, in DeciMamba, they only considered the scenario where there is too much decay. Regarding token pruning, our analysis in Section 3.2.2 is already an analysis of token pruning on the "forgetting" effect. Sliding window is equivalent to pruning all tokens outside of the last $r$ tokens ($r$ is the window size). The conclusion is that token pruning induces forgetting, which mitigates the issue caused by the inability to forget.
2. Our claim that "Mamba-2 outperforms similarly sized transformer models" is based on publicly released transformer-based models. To the best of our knowledge, there are no publicly released transformer-based models of 370M parameters with such a strong performance on the passkey retrieval task. We are not claiming that transformer-based models cannot achieve this performance, but this is an impressive performance considering the great inference efficiency of Mamba-2.

Thank you for your time, if there are any more concerns, please let us know.