Stuffed Mamba: Oversized States Lead to the Inability to Forget
Authors: Yingfa Chen
Xinrong Zhang
Shengding Hu
Xu Han
Zhiyuan Liu
Maosong Sun
Venue: COLM 2025 (2025)
Keywords: state space models
mamba
long-context modeling
linear attention
TL;DR: We discover and explain why Mamba-based models cannot robustly forget past information, which hurts length generalization abilities.

Abstract: Recent advancements in recurrent architectures, such as Mamba and RWKV, have showcased strong language capabilities. Unlike transformer-based models, these architectures encode all contextual information into a fixed-size state, leading to great inference efficiency. However, this approach can cause information interference, where different token data conflicts, resulting in performance degradation and incoherent outputs beyond a certain context length. To prevent this, most RNNs incorporate mechanisms designed to "forget" earlier tokens. In this paper, we reveal that Mamba-based models struggle to effectively forget earlier tokens even with built-in forgetting mechanisms. We demonstrate that this issue stems from training on contexts that are too short for the state size, enabling the model to perform well without needing to learn how to forget. Then, we show that the minimum training length required for the model to learn forgetting scales linearly with the state size, and the maximum context length for accurate retrieval of a 5-digit passkey scales exponentially with the state size, indicating that the model retains some information beyond the point where forgetting begins. These findings highlight a critical limitation in current RNN architectures and provide valuable insights for improving long-context modeling. Our work suggests that future RNN designs must account for the interplay between state size, training length, and forgetting mechanisms to achieve robust performance in long-context tasks.

Published as a conference paper at COLM 2025
Stuffed Mamba: Oversized States Lead to the Inability to Forget
Yingfa Chen, Xinrong Zhang, Shengding Hu, Xu Han", Zhiyuan Liu' & Maosong Sun Department of Computer Science and Technology, Tsinghua University, Beijing, China yingfa-c24@main tsinghua. www.iaaedd.ccom enu.cn, (han-xu, lisziygt edu.cn
Abstract
Recent advancements in recurrent architectures, such as Mamba and RWKV, have showcased strong language capabilities. Unlike transformer-based models, these architectures encode all contextual information into fixed- size state, leading great inference efficiency. However, this approach can cause information interference, where different token data conflicts, resulting in performance degradation and incoherent outputs beyond a certain context length. To prevent this, most RNNs incorporate mechanisms designed to "forget" earlier tokens. In this paper, we reveal that Mamba- based models struggle to effectively forget earlier tokens even with built-in forgetting mechanisms. We demonstrate that this issue stems from training on contexts that are too short for the state size, enabling the model to perform well without needing to learn how to forget. Then, we show that the minimum training length required for the model to learn forgetting scales linearly with the state size, and the maximum context length for accurate retrieval of 5-digit passkey scales exponentially with the state size, indicating that the model retains some information beyond the point where forgetting begins. These findings highlight a critical limitation in current RNN architectures and provide valuable insights for improving long-context modeling. Our work suggests that future RNN designs must account for the interplay between state size, training length, and forgetting mechanisms to achieve robust performance in long-context tasks.
1 Introduction
Transformer-based large language models (LLMs) (Achiam et al., 2023; Dubey et al., 2024) have shown impressive capabilities in processing very long sequences (Gemini Team et al., 2024; MiniMax et al., 2025). However, these models incorporate self-attention (Vaswani et al., 2017) whose complexity scales quadratically with sequence length, making long-context processing costly. In contrast, recurrent neural networks (RNNs) (Bengio et al., 1994) have fixed-size contextual memory. Thus, their per-token time and space complexities are constant and they are much more efficient for long sequences. Despite this advantage, their effectiveness in modeling long contexts remains underexplored. Most recent state-of-the-an (SOTA) RNNs, such as Mamba-1 and Mamba-2 (Gu & Dao, 2023; Dao & Gu, 2024), GLA (Yang et al., 2024a), and RWKV (Peng et al., 2024a) are trained on context lengths below 10K, and existing works have shown that their performance degrades sharply when the context length exceeds the model's training length' (Ben-Kise Kish et al., 2024; Zhang et al., 2024a; Waleffe et al., 2024).
In this paper, we analyze the factors that cause the inability of Mamba-based models to handle contexts longer than the training length. By inspecting memory retention strength and modifying the forgetting mechanism, we discover that the performance drop is caused
"Corresponding Authors.
'Throughout this paper, "training length" refers to the context length used during training.
1

---

Published as a conference paper at COLM 2025
130M
700M
100 75 50 25
Ane. Depit 1943
130 T. 25
Ane. Depti
50 25
Ann. D420
Ane. Depen
-05
120 7
16K
Content Lon.
Context Lon.
Context Lor.
Context Lon.
Token position
(a) 130M
(b) 370M
(c) 780M
(d) 1.3B
Figure 1: The LM loss of Mamba-2 as a function of token position. The training length is 8K.
Figure 2:
The accuracy of Mamba-2 2 on the passkey retrieval
task. "Ans. Depth" refers to the passkey position divided by the context length.
by the inability to forget earlier tokens, although Mamba has a built-in forgetting mechanism. Insufficient forgetting leads to interference between multiple token representations, causing faulty memory recall and, ultimately, performance degradation in longer contexts. We provide two lines of evidence for this discovery: (1) The first token's retention strength is very high throughout its training context window. (2) Artificially inducing forgetting via interventions of the state's update rule can mitigate this performance degradation.
We hypothesize that the inability to learn effective forgetting mechanism is due to state overparameterization-- where the model's state is excessively large, allowing it to minimize language modeling loss without much forgetting. Two key pieces of evidence support this hypothesis. (1) Initially, the model demonstrates robust forgetting, retaining only the last k tökens and forgetting earlier ones. However, as training progresses, the model's ability to forget diminishes while its recall of contextual information improves, resembling overfitting. This suggests that the model increasingly attempts to retain all available information within the context. (2) We observe that forgetting occurs only when the training context length exceeds the state's s capacity to retain all information, forcing the model to forget less relevant details. Notably, larger states require longer training context lengths to effectively learn and implement forgetting.
We next investigate the minimum training length required for Mamba-2 to effectively learn forgetting and the maximum context length which the model can recall information. First, by varying model sizes and training lengths, observe that the training length threshold scales linearly with the state size, confirming that forgetting only occurs when the training length exceeds the model's state capacity. Second, while this threshold represents the point at which the contextual information exceeds the state's capacity, we demonstrate that the model can still recall tokens beyond this context window. Evaluation on passkey retrieval (Mohtashami & Jaggi, 2 simple retrieval task-shows that the maximum context length with perfect retrieval accuracy scales exponentially with state size. Notably, with continued pre-training, Mamba-2 with 370M parameters achieves near-perfect retrieval on 256K context length, outperforming similarly sized transformer models. These findings suggest that current training lengths for RNN models may be suboptimal and underscore the potential of RNN-based architectures for modeling long-context sequences.
This paper is structured as follows. Section 2 describes the Mamba-2 architecture and provides evaluation results showing its inability generalize beyond its training length. Section 3 provides arguments for the importance of forgetting and provides evidence showing that the model has failed to learn a robust forgetting mechanism. Section 4 presents high-level explanation for why Mamba-2 has failed to learn how to forget. Finally, in Section 5, experiments are conducted to verify the claims and provide important conclusions for training long-context context recurrent models.
The main findings of this paper can be summarized as follows:
The inability to forget. We discover that Mamba-2, and RWKV-6 (some SOTA RNNs) do not know how to robustly forget earlier information to avoid memory overload. This causes performance degradation for contexts longer than the training length.
2

---

Published as a conference paper at COLM 2025
State overparameterization. We provide overparameterization a plausible explanation for the inability to forget and provide empirical evidence for this hypothesis.
Minimum training length of forgetting. In alignment with the state overparameterization hypothesis, we empirically discover that for any state size, there exists a training length threshold where the Mamba-2 learns forgetting if and only if the training length is above that threshold. We also find that this relationship is linear.
2 Preliminaries
In this section, we first describe the Mamba-2 architecture and corresponding notations. Then, we evaluate Mamba-2 on language modeling and passkey retrieval with context lengths exceeding their training length to illustrate the consequence of the inability to forget.
Most experiments in this study focus on Mamba-2 (Dao & Gu, 2024) because has shown strong capabilities on several tasks and has publicly available checkpoints of multiple sizes, allowing us to explore the relationship between state sizes and length limits. Moreover, is more widely studied than other RNNs, making it easier to use existing works as reference.
2.1
Mamba-2
The Mamba-2 architecture consists of L layers, each consisting of H heads computed in parallel. The layer's output is the sum of the heads' outputs. Let E R E R denote the input and output vectors of the layer at time step. The computation at f time step for each head can be formulated follows:
y/+ = C,h₃ E R IxP
(Query rule)
(1)
h; h1-1 0t
+ +Ba
E RxP
(Update rule)
(2)
Decay Insertion
where C E R E x
RIXP, 0a R are functions of ,,, d, N, P are hyper- parameters, denoting the hidden dimensionality, state dimension, and head dimension, respectively, and h, is the 1-th recurrent state. Eq 1 and 2 are called the "query rule" and "update rule" because they determine how memory is queried from the recurrent state, and how the state is updated.
The other variables are parameterized as follows:
B, B,A, Nx1
(3)
&; exp( A exp E R
(4)
A; Softplus(u; WA bx) ER
(5)
where B, Nx1 is a function of ", and A,b E R, WA E Rdxl are trainable model parameters. Appendix A presents more details on the model. Notably, Mamba-2's update rule is similar to many existing RNNs (Peng et al., 2024a; Sun et al., 2023; Yang et al., 2024a). Thus, some conclusions/ /insights may apply to other architectures. We leave such exhaustive ablation studies for future work.
Importantly, h, is the contextual memory representation that stores information from all tokens up to f. &1 E (0,1) is the memory decay multiplier that controls the strength of forgetting. Past information is completely forgotten when &; 0 and completely retained with & + 1. In this paper, we refer to &; as the memory retention strength.
2.2 Length Generalization Failure of Mamba-2
Language Modeling Figure 1 shows the language modeling loss Mamba-2 as a function of token position. The result shows that Mamba-2 models suffer great performance degra- dation when the context length is much longer than their training lengths. Furthermore, we find that larger models have worse length generalization abilities.
3

---

Published as a conference paper at COLM 2025
Passkey Retrieval Evaluation Language modeling may not reflect downstream capabili- ties, thus, we also evaluate Mamba-2 recall ability on the passkey retrieval task (Mohtashami & Jaggi, 2023; Zhang et al., 2024a). It is widely-used simple synthetic task where model is prompted to recall 5-digit passkey from a lengthy context.
The passkey retrieval result is reported in Figure 2. We find that Mamba-2 (except for the smaller 130M checkpoint) has near-perfect retrieval accuracy within 8K tokens, but poor or even zero accuracy on sequences longer than 16K, regardless of model sizes. This behavior is unexpected because the update rule (Bq. 2) has a stable exponential memory decay (it converges to constant value the variables are fixed). Therefore, we expect RNNs of such form to have good retrieval accuracy on the last k tokens, and tokens earlier than that are forgotten. However, when the context too long, Mamba-2 fails to even recall very recent tokens. This implies that the limitation is not the inability to retain memory about the answer, but the inability forget irrelevant past tokens.
More experimental details and the results of some other recurrent architectures can be found in Appendix B and D.
The Inability to Forget
In this section, we argue that Mamba-2's length generalization failure can be attributed to its inability to forget contextual information. We first provide arguments for the impor- tance of forgetting. Then, we present empirical evidence to verify that the model has not learned a robust forgetting mechanism. Furthermore, we show how the inability to forget is manifested in the statistics of the recurrent state.
3.1
The Importance of Forgetting
The fact that the model fails to retrieve tokens at any position when the context length exceeds certain threshold has one critical implication: the existence of earlier tokens damages the model's ability to recall both earlier and more recent tokens. This is because the model's state ht at time step can be formulated as a weighted sum of past information:
(1i,,)
h
Ags B
agr
(0.1)
(6)
When we try to retrieve the memory inserted at time step s, would query the state with C Bs, which returns:
y/ G Σ axx;:x; (C,Baa)
Σ ift
Nii C,
Retrieval error
When all B, are mutually orthogonal, querying the memory with G B3 returns a scaled version of xş. However, the context length increases, B, cannot be mutually orthogonal, in which case multiple memory entries interfere and cause retrieval errors. small error may not affect the final output because subsequent calculations may tolerate these errors. However, if there are many preceding tokens, then this retrieval error can be too large. To optimize retrieval accuracy, the model needs to produce small enough decay a, to diminish the inference from earlier tokens, which can be viewed as forgetting them.
3.2 Evidence for the Inability to Forget
Here, we provide empirical evidence that confirms that Mamba-2 has failed learn how to forget past information.
3.2.1
Evidence 1: High Retention of the First Token
Based on Eq. 6, we can view aj as the memory strength of the i-th token at i time step. The retention strengths of earlier tokens are always smaller than those of more recent tokens.
4

---

Published as a conference paper at COLM 2025
1.0
HD
Original
Mean
Variance
1,000
111
LongNamba
Siidin Win
150
HD
Cumunitiv decay
H2
H2
0.5
0.0
100
-03
0.995
0.2
H6
0.1
47
5000
10000 15000
12K 20K 28K
10K 20K 30K Toxee position I7)
10K. 24K 30K
Posisin H1
Token position (1)
Token postion (3)
Figure 3: The reten- tion strength of the first token (41:)) over time. Each curve represents head.
Figure 4: LM loss of Mamba-2 370M at differ- ent positions when induc- ing forgetting (see Section 3.2.2).
Figure 5: The mean and variance of the first 8 heads in layer 38 of Mamba-2 370M. exhibits a clear explosion when f is greater than the training length.
We find that some heads have very strong inclination toward retaining all information within the training length. As an example, Figure shows the cumulative decay of the first token in the first eight heads of the 38th layer, and three of the heads have memory retention strength over 0.9972 at (=8K. Similar observations can be found in other heads and in other layers as well. This implies that the model has not learned to forget information (by producing a smaller ay), but it still has decent language modeling capabilities because the information of 8K tokens is typically not enough to overload the memory.
3.2.2 Evidence 2: Inducing Forgetting Can Improve Length Generalization
Here, we demonstrate that artificially inducing more forgetting without training can improve performance by reducing past memory interference.
Reduced Memory Retention and Insertion (RRI) This method assumes that a, and B control the memory retention and insertion strength, respectively. We scale them with multiplier smaller than The actual multipliers used are 0.9999 for aj and 0.75 for B,, which is chosen by validation using average loss pre-training data with 32K context length.
Sliding Window We can utilize the fact that the state it can be written as weighted sum (Eq. 6) to simulate a sliding window mechanism without re-processing from the start of
the window at every step. Let w E NN denote the window size and h," E NRP denote the hidden state when applying the model on the last tokens at time step t. We can then compute n!TT) exactly as the difference between two states:
11
Σ
ai BiX
Σ
Kis Bix - at rils
a Ba2-BBx = h1 a, 111341-
(8)
i--++1
During streaming generation, we only have to maintain (4, 1,, &s r) and advance each of them in parallel. However, directly computing ata- may suffer from instability due to floating-point imprecision. Therefore, we maintain A,-r - - A instead, and re- compute a; rf exp (-A; rit exp(A)) at every step, which incurs minimal computational cost. This method can be applied to all RNNs that can be written as weighted sum.
Result From Figure 4, one see that the original model has the worst length generaliza- tion abilities. LambMamba (a length extrapolation method) (Zhang, 2023) and our methods for inducing forgetting can alleviate this generalization failure, although LongMamba and RRI compromise short-context performance due to weaker memory insertion. This further confirms the fact that memory over-retention the culprit of this degradation.
Thhis is a cumulative product, so the decay at each time step is even closer to We also have to cache the last token IDs, but their size is negligible compared to 14-1 and h

---

Published as a conference paper at COLM 2025
10a
= 20K
Mean
Var
10a
0.02
107
0.00
0.02
-0.02
10a
0-
20000
20000
25000
100
Value
(a) 4 over time.
(b) over time. (c) time. Each curve repre- Each curve repre- Each curve repre- sents head. sents a head.
Figure 6: The distri- bution of the channels in an exploding state (head 2 in layer 38) at two different time steps.
sents a channel.
Figure 7: The value of various components in the update rule (A,, Bi, and xi) on some heads with large retention values in the 38th layer in Mamba-2 370M. The red dotted line indicates the training length.
3.3 The Manifestation of Over-Retention
We also examine how the inability to forget is manifested in the state's values. Since the recurrent state's dimensionality does not change over time, the sharp change of behavior during length generalization must be result of a change in the state's ' distribution. For reproducibility and better visualization, we use the "newlines" prompt (string with only "An") and inspect the statistics of the recurrent states every head Mamba-2 370M and find that the mean and variance of some heads change sharply when the context length exceeds the training length. One example is shown in Figure 5. Appendix reports the statistics of every layer. The state at f 20K of one head with exploding variance is shown in Figure 6. From it, we discover that this variance explosion can be largely attributed to few outlier channels while most channels are relatively stable.
4 State Overparameterization
Here, we present high-level explanation for why Mamba-2 over-retains memories: state overparameterization. The state is excessively large for the training length, allowing the model to achieve strong language modeling performance without learning how forget when the state is overloaded with memories.
To support this hypothesis, we present two pieces of evidence: (1) Mamba-2 starts with the ability to forget, but slowly loses this ability as the amount of training data increases, which coincides with behaviors of overfitting, and (2) for any state size, there training context length threshold Tooret such that Mamba-2 learns to forget if and only if Ttrain > Trorgtt where Trrain denotes the training context length.
4.1 Evidence 1: More Training => Less Forgetting
We pre-train Mamba-2 370M from scratch with Train = 512 using the RedPajama (Computer, 2023) corpus and evaluate the intermediate checkpoints on passkey retrieval, reported in Figure 8. It shows that the model's retrieval accuracy for contexts longer than the training length slowly decreases we increase the amount of training data. Meanwhile, the model's accuracy for contexts shorter than the training length increases. This indicates that the model converges toward more retention and less forgetting. Since language modeling loss is only computed for tokens within the training length, this behavior is induced by minimizing loss. This reduced forgetting leads to conflicts between token representations, impairing memory recall accuracy when an excessive number of tokens are inserted.
4Similar observation be found with any model size.

---

Published as a conference paper at COLM 2025
Ans. Depth (%)
73 50 23
Ans. Depth (%)
25
0
(%)
25
Max Recall Len (K)
400
b. - +
1.355
b=4.756
50
c= -0.442
100 15
75
200
10
(a) 10B tokens.
(b) 20B tokens.
(c) 40B tokens.
State size (M parameters)
Figure 8: Passkey retrieval results intermediate check- points during the pre-training of Mamba-2 370M on 512 sequence length. Generalization failure only occurs in the model beyond certain amount of training data.
Figure 9; The maximum con- text length of accurate passkey retrieval (i.e., Tecall in Section 4.3) as function of state size.
This behavior be viewed kind of overfitting because the state's distribution with short contexts is not varied enough for the model to generalize to the distribution in longer contexts. In other words, the state has too many parameters for the given training length.
4.2 Evidence 2: Forgetting Learned + Sufficient Training Length
We empirically confirm that larger recurrent states require longer training contexts for the model to learn to forget. This because the model will only learn to forget when the amount of contextual information exceeds the state capacity. This hypothesis implies the following law:
Let Nş and Ttrain denote the recurrent state size and training length, spectively, there exists threshold Tooget (Nş) such that the model learns to forget if and only Ttrain > forget-
We empirically validate this by sweeping different training lengths for different state sizes, and checking whether the model has successfully learned how to forget. Concretely, we train multiple Mamba-2 with different state sizes and training lengths to find the relationship between Torget and Ng. To determine whether the model has learned robust forgetting, we feed prompts with 1M tokens to the model and check the model's loss exceeds 2x the maximum loss within Tirain tokens at any point. The loss is averaged over 128 prompts. The result is reported in Section 5.1.
4.3 Maximum Recall Context Length
The fact that the amount of information in Torget tokens exceeds the state's capacity, does not necessarily imply that the model fails to recall information beyond the last Troget tokens, especially when there clear distinction between the target information and other contextual information. Therefore, we also search for the maximum context length from which the model can accurately perform passkey retrieval. We refer this context length as the model's maximum recall context length, denoted with Trecall- Similar to the previous section, we train with different lengths for different state sizes and identify the maximum context length where the model has an accuracy over 95% as Trecall- In this task, the noisy context repetitive, thus, the amount contextual information is largely independent of the context length. Therefore, ideally, the recall threshold should grow roughly exponentially with the state 5
If we train Mamba-2 on passkey retrieval data, the model can theoretically handle infinitely long contexts. Here, the model is only trained with the next token prediction objective, which means the model will not ignore the irrelevant context, and the ability to retain information for extended time emerges from language modeling.
7

---

Published as a conference paper at COLM 2025
Training Len. (
60
a
4K
b= -4.469
4.0
3.5
40
3.6
LOSS
16K
32K
3.2
3.0
2.8
0K
8K
16K
24K
32K
0K
32K
64K
96K
128K
position
State (M parameters)
(a) Mamba-2 130M
(b) Mamba-2 370M
Figure 11: Minimum train- ing length at which the model learns robust forgetting (i.e., Torget in Section 4.2).
Figure 10: LM loss at each token position for different training lengths. Evaluated on RedPajama.
Experiments
We briefly describe the data and model configurations used to identify the forget threshold furget and maximum recall context length recall- Due to limited space, more comprehensive experimental details are reported in Appendix F.
Data We start from RedPajama-V2 V2 (Computer, 2023), an open dataset with 30T tokens from the Internet, and perform deduplication to ensure data quality and discard documents that are too short.
Models We experiment with six model sizes to find the relationship between state capacity and size. For each of them, we perform an extensive search with training lengths up to 256K tokens. To save cost, 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.
5.1 The Existence Forget Threshold
In Figure 10, we plot the language modeling perplexity as a function of token position for Mamba-2 130M and 370M with different training lengths. We can see that for each model size, there is training length threshold, beyond which the model has much better length extrapolation, which supports our arguments discussed in Section 4.2.
5.2 Forget Threshold Function of the State Size
Figure 11 shows the minimum training length needed for Mamba-2 to learn forgetting, The rightmost data point in the plot corresponds to Mamba-2 370M. We have confirmed that the 780M model (with state size of 19.3M) also has poor length generalization at training lengths below 128K, but do not have enough resources to train the model beyond this length. The results establish linear relationship Toorett 5.172 Ns 4.469 between the length Tirain forget at which the model can learn robust forgetting and the state size Ns. The R value is over 0.999. This indicates that to train Mamba-2 with robust length generalization, one should use training lengths that grow linearly with the state size.
5.3 Maximum Recall Context Length as Function of the State Size
The second plot of Figure shows the recall threshold of Mamba-2 passkey. The maximum contexts length in which Mamba-2 can accurately retrieve 5-digit passkeys is exponential concerning the state size, the function is Treaall 4.756 (1.365 No - 1) - 00744, with an R² value over 0.999. This is because the amount of information in the context does not increase with its length. In other words, storing constant amount of information while the

---

Published as a conference paper at COLM 2025
number of combinations of the state grows exponentially with the number of elements. The result is very promising because, to the best of our knowledge, no previous models with less than 1B model parameters have near-perfect accuracy at this length in this task.
Related Works
RNN-Based Language Models This paper focuses Mamba-2, a recurrent architecture that can be viewed as a variant of gated linear attention (Yang et al., 2024a). Many recently proposed RNNs can also be viewed as GLA variants. These include the RWKV series (Peng et al., 2023; 2024a), the Mamba series (Gu & Dao, 2023; Dao & Gu, 2024), GLA (Yang et al., 2024a), and many more (Zhang et al., 2024b; Yang et al., 2024b; De et al., 2024; Arora al., 2024b; Orvieto et al., 2023; Sun et al., 2023). Our methods may apply to these architectures as well. Some recent/ concurrent RNNs such as Gated DeltaNet (Yang al., 2025), RWKV-7 (Peng et al., 2025), xLSTM (Beck et al., 2024), and Titans (Behrouz et al., 2024) have gone beyond gating-based memory decay mechanism and are out of the scope of this paper.
Length Generalization Most SOTA language models in the last few years have been based on the transformer (Vaswani et al., 2017) architecture. These models, when using certain variants of position encoding, can process arbitrarily long sequences. However, they exhibit severe performance drops on tokens beyond the training length (Zhao et al., 2024). To alleviate this shortcoming, many works have focused on modifying positional encoding (Peng et al., 2024b; Zhu et al., 2024; Ding et al., 2024; Jin et al., 2024), some achieving training-free length generalization certain extents (Zhang al., 2025).
Length Generalization of Mamba Some prior works investigated the performance of Mamba as function of context length (Park et al., 2024; Wen et al., 2025). Jelassi et al. (2024) empirically showed a sharp performance drop beyond the training length for Mamba on copying task and also showed that Mamba struggles to copy from context unless its state size grows linearly with the context length. Arora et al. (2024a) discussed the associative recall abilities of transformer and some RNNs. Wang et al. (2025) is most related to our work. They discussed the issue of smottsioo introduced by the memory decay term. In contrast, our paper explores a setting where recency may be preferred, but interference from earlier tokens damages the recall accuracy of recent tokens.
Some concurrent works have explored extending Mamba's 's context length by controlling the discretization term (A, in Eq. 2) (Ben-Kish et al., 2024), such as dividing it by a constant to make it smaller (Zhang, 2023). This makes the memory decay factor (a; in Eq. 4) closer to 1, which makes the state retain more contextual information. However, it also unnecessarily diminishes the inserted information on all tokens. Similar to the above works, this study explores the cause of Mamba-2's inability to generalize beyond its training context length and provides valuable insights into training Mamba-2 models that generalize better.
7 Conclusion
This paper demonstrates that while the Mamba architecture includes memory decay mechanism, it fails to effectively learn forgetting in practice. As result, when the context exceeds the training length, the model produces incoherent outputs. This issue arises from training with contexts that are too short relative to the state size. Empirical results show that robust forgetting is only learned when the training context length surpasses certain threshold, which increases linearly with the state size. Notably, the model is still capable of recalling some contextual information beyond this threshold. These findings offer valuable insights into the causes and consequences of the model's inability to forget, highlighting key limitations of the Mamba architecture. Nevertheless, the insights gained from this study provide promising foundation for improving Mamba's performance in long-context modeling, paving the way for more effective applications in tasks requiring extended context lengths.

---

Published as a conference paper at COLM 2025
Acknowledgements
This work is supported by the high-quality development project of MIIT and grant from the Guoqiang Institute, Tsinghua University.
References
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. GPT-4 technical report. arXiv preprint ar arXio:2303.08774 2023.
Simran Arora, Sabri Eyuboglu, Aman Timalsina, Isys Johnson, Michael Poli, James Zou, Atri Rudra, and Christopher Ré. Zoology: Measuring and improving recall in efficient language models. In The Twelfth International Conference on Learning Representations, 2024a.
Simran Arora, Sabri Eyuboglu, Michael Zhang, Aman Timalsina, Silas Alberti, Dylan Zinsley, James Zou, Atri Rudra, and Christopher Re. Simple linear attention language models balance the recall-throughput tradeoff. In Workshop on Efficient Systems for Foundation Models II @ ICML2024, 2024b.
Maximilian Beck, Korbinian Pöppel, Markus Spanring, Andreas Auer, Oleksandra Prud- nikova, Michael Kopp, Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. XLSTM: Extended long short-term memory. In Amir Globersons, Lester Mackey, Danielle Belgrave, Angela Fan, Ulrich Paquet, Jakub M. Tomczak, and Cheng Zhang (eds.), The Thirty-eighth Annual Conference on Neural Information Processing Systemis, 2024.
Ali Behrouz, Peilin Zhong, and Vahab Mirrokni. Titans: Learning to memorize at test time, 2024.
Assaf Ben-Kish, Itamar Zimerman, Shady Abu-Hussein, Nadav Cohen, Amir Globerson, Lior Wolf, and Raja Giryes. DeciMamba: Exploring the Length Extrapolation Potential of Mamba, 2024.
Yoshua Bengio, Patrice Simard, and Palo Frasconi. Learning long-term dependencies with gradient descent is difficult. IEEE transactions on neural networks, 5(2):157-166, 1994.
Jacob Buckman and Carles Gelada. Compute-optimal Context Size, 2024.
Together Computer. RedPajama: an open dataset for training large language models, 2023.
URL https://githul wtwccooo//eeeeeee potpereeecetete/eeaaaaaa...tt RReaaajama--aaaaaa Data.
Tri Dao and Albert Gu. Transformers are SSMs: Generalized models and efficient algorithms through structured state space duality. Xiv preprint ar XrXii:2405.2100 2024.
Soham De, Samuel Smith, Anushan Fernando, Aleksandar Botev, George-Cristian Muraru, Albert Gu, Ruba Haroun, Leonard Berrada, Yutian Chen, Srivatsan Srinivasan, Guillaume Desjardins, Arnaud Doucet, David Budden, Yee Whye Teh, Razvan Pascanu, Nando de Freitas, and Caglar Gulcehre. Griffin: Mixing gated linear recurrences with local attention for efficient language models. Xrv, abs/ / 02//24402..12222 2024.
Yiran Ding, Li Lyna Zhang, Chengruidong Zhang, Yuanyuan Xu, Ning Shang, Jiahang Xu, Fan Yang, and Mao Yang. LongroPE: Extending LLM context window beyond million tokens. In Forty-first International Conference on Machine Learning, 2024.
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, et al. The llama herd of models. ar arxiv preprint or XrXi:2247.217883 2024.
Stefan Elfwing, Eiji Uchibe, and Kenji Doya. weighed- linear units for neural network function approximation reinforcement learning. Neural Networks, 107:3-11, 2018. doi: NO.00
10

---

Published as a conference paper at COLM 2025
Gemini Team, Petko Georgiev, Ving Ian Lei, Ryan Burnell, Libin Bai, Anmol Gulati, Garrett Tanzer, Damien Vincent, Zhufeng Pan, Shibo Wang, et al. Gemini 1.5: Unlocking multi- modal understanding across millions of tokens of context. arXiv preprint arXiv:2403.05530, 2024.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. ArXiv, abs/2312.00752. 2023.
Shengding Hu, Yuge Tu, Xu Han, Chaoqun He, Ganqu Cui, Xiang Long, Zhi Zheng, Yewei Fang, Yuxiang Huang, Weilin Zhao, et al. MiniCPM: Unveiling the potential of small language models with scalable training strategies. arXiv preprint arXiv:2444.06395, 2024.
Samy Jelassi, David Brandfonbrener, Sham M. Kakade, and Eran Malach. Repeat After me: Transformers are better than state space models at copying. In Forty-first International Conference on Machine Learning, 2024.
Hongye Jin, Xiaotian Han, Jingfeng Yang, Zhimeng Jiang, Zirui Liu, Chia-Yual Chang, Huiyuan Chen, and Xia Hu. LLM maybe longLM: Selfextend LLM context window without tuning. In Forty-first International Conference on Machine Learning, 2024.
MiniMax, Aonian Li, Bangwei Gong, Bo Yang, Boji Shan, Chang Liu, Cheng Zhu, Chunhao Zhang, et al. MiniMax-01. Scaling Foundation Models with Lightning Attention, 2025.
Amirkeivan Mohtashami and Martin Jaggi. Landmark attention: Random-access infinite context length for transformers. In Workshop on Efficient Systems for Foundation Models@ ICML2023, 2021.
Antonio Orvieto, Samuel L Smith, Albert Gu, Anushan Fernando, Caglar Gulcehre, Razvan Pascanu, and Soham De. Resurrecting recurrent neural networks for long sequences, 2023.
Jongho Park, Jaeseung Park, Zheyang Xiong, Nayoung Lee, Jaewoong Cho, Samet Oymak, Kangwook Lee, and Dimitris Papailiopoulos. Can mamba learn how to learn? A compar- ative study on in-context learning tasks. In Forty-first International Conference on Machine Learning, ICML 2024, Vienna, Austria, July 21-27, 2024. OpenReview.net, 2024.
Bo Peng, Eric Alcaide, Quentin G. Anthony, Alon Albalak, Samuel Arcadinho, Stella Bider- man, Huanqi Cao, Xin Cheng, Michael Chung, Matteo Grella, G Kranthikiran, Xuming He, et al. RWKV: Reinventing RNNs for the Transformer Era. In Conference on Empirical Methods in Natural Language Processing, 2023.
Bo Peng, Daniel Goldstein, Quentin Anthony, Alon Albalak, Eric Alcaide, Stella Biderman, Eugene Cheah, Xingjian Du, Teddy Ferdinan, Haowen Hou, Przemyslaw Kazienko, Kranthi Kiran Jan Kocoñ, et al. Eagle and Finch: Rwkv with Matrix-Valuee States and Dynamic Recurrence. Xiv, abs/ 2404.05892 2024a.
Bo Peng, Ruichong Zhang, Daniel Goldstein, Eric Alcaide, Haowen Hou, Janna Lu, William Merrill, Guangyu Song, Kaifeng Tan, Saiteja Utpala, Nathan Wilce, Johan S. Wind, Tianyi Wu, Daniel Wuttke, and Christian Zhou-Zheng. Rwkv-7 " goose" with expressive dynamic state evolution, 2025.
Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole. YaRN: Efficient context window extension of large language models. In The Twelfth International Conference on Learning Representations, 2024b.
Zhen Qin, Songlin Yang, Weixuan Sun, Xuyang Shen, Dong Li, Weigao Sun, and Yiran Zhong. HGRN2: Gated Linear RNNs with State Expansion. In First Conference on Language Modeling, volume /bbp//24000794 2024.
Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue, Jianyong Wang, and Furu Wei. Retentive Network: A Successor to Transformer for Large Language Models. arXiv, abs/2307.086/) 2023.
11

---

Published as a conference paper at COLM 2025
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances neural informa- fion processing systems, 30, 2017.
Roger Waleffe, Wonmin Byeon, Duncan Riach, Brandon Norick, Vijay Korthikanti, Tri Dao, Albert Gu, Ali Hatamizadeh, Sudhakar Singh, Deepak Narayanan, Garvit Kulshreshtha, Vartika Singh, Jared Casper, Jan Kautz, Mohammad Shoeybi, and Bryan Catanzaro. An Empirical Study of Mamba-based Language Models, 2024.
Peihao Wang, Ruisi Cai, Yuehao Wang, Jiajun Zhu, Pragya Srivastava, Zhangyang Wang, and Pan Li. Understanding and Mitigating Bottlenecks of State Space Models through the Lens of Recency and Over-smoothing. In The Thirteenth International Conference Learning Representations, 2025.
Kaiyue Wen, Xingyu Dang, and Kaifeng Lyu. RNNs are not transformers (yet): The key bottleneck on in-context retrieval. In The Thirleenth International Conference on Learning Representations, 2025.
Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim. Gated linear at- tention transformers with hardware-efficient training. In Forty-first International Conference on Machine Learning, ICML 2024, 2024a.
Songlin Yang, Bailin Wang, Yu Zhang, Yikang Shen, and Yoon Kim. Parallelizing linear trans- formers with the delta rule over sequence length. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024b.
Songlin Yang, Jan Kautz, and Ali Hatamizadeh. Gated delta networks: Improving mamba2 with delta rule. In The Thirteenth International Conference on Learning Representations, 2025.
Biao Zhang and Rico Sennrich. Root mean square layer normalization. In The Thirty-second Annual Conference on Neural Information Processing Systems, 2019.
Peiyuan Zhang. Longmamba, 2023. URL
Xinrong Zhang, Yingfa Chen, Shengding Hu, Zihang Xu, Junhao Chen, Moo Hao, Xu Han, Zhen Thai, Shuo Wang, Zhiyuan Liu, and Maosong Sun. coBench: Extending long context evaluation beyond 100K tokens. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics, 2024a.
Xinrong Zhang, Shengding Hu, Weilin Zhao, Huadong Wang, Xu Han, Chaoqun He, Guoyang Zeng, Zhiyuan Liu, and Maosong Sun. Optimal Rop' extension via Bayesian Optimization for training-fre free length generalization. Al Open, 2025.
Yu Zhang, Songlin Yang, Rui-Jie Zhu, Yue Zhang, Leyang Cui, Yiqiao Wang, Bolun Wang, Freda Shi, Bailin Wang, Wei Bi, Peng Zhou, and Guohong Fu. Gated slot attention for efficient linear-time sequence modeling. In NeurlPS, 2024b.
Liang Zhao, Xiachong Feng, Xiaocheng Feng, Weihong Zhong, Dongliang Xu, Qing Yang, Hongtao Liu, Bing Qin, and Ting Liu. Length extrapolation of transformers: survey from the perspective of positional encoding. In Findings of the Association for Computational Linguistics: EMNLP, 2024.
Dawei Zhu, Nan Yang, Liang Wang, Yifan Song, Wenhao Wu, Furu Wei, and Sujian Li. PoSE: Efficient context window extension of LLMs via positional skip-wise training. In The Twelfth International Conference on Learning Representations, 2024.
12

---

Published as a conference paper at COLM 2025
A Mamba-2 Architecture
For completeness, we give a more detailed formulation of the Mamba-2 architecture here, A
although we recommend the readers refer the original paper (Dao & Gu, 2024) or detailed blog post by the authors The model accepts sequence of T token IDs as input
I= [1, iT] E N', where i, E f1,2, ... V), denotes the vocabulary size. It performs next token prediction by predicting the probability distribution over the vocabulary at each time step, denoted as P R The model can be formulated as follows.
U(0) Embed (T) E RdxT
((() = Mamba (1) NNor (uu--
E RdXT
Norm [ R VXT
P Embedout
where L denotes the number of layers, I {1, L} denotes the layer index, U!!) E R represents the input of the I-th layer, U represents the input of the first layer. Mamba (.) denotes the I-th Mamba layer, Embedin! (-) and Embedaut(.) denote the input and output embedding layers, and Norm(.) denotes RMS normalization (Zhang & Sennrich, 2019). denotes the number of dimensions of each token embedding. Similar to many other models, Mamba-2 ties the weight of the input and output embedding layers.
Each Mamba layer consists of H "heads" that are computed in parallel. The result of which is summed together. Notably, the notations here are slightly different from Section 2 because simplified the notations in the main content save space. The t-th token {1, T}) in a head is computed as follows:
J/: W,Norm (o, + Wgaaeelt) E Rdxx
(9)
0, C,h+ + +D RIP
(10)
h; h exp( exp A)) + A, E NxP
(11)
G ((Conv(Wcui)) E RRNN
(12)
B1 c(Conv(WBu)) E RR
(13)
X1 o(Conv(W..u)) T ERIx
(14)
A; Softplus(WAtt) ba) E
(15)
ut E I denotes the t-th input representation. In other words, for the I-th layer, we
have U = 1 u7 (1) u(++) = [, ,y.",...., ".."...., and u, 1) y!" Conv(.) denotes
channel-wise one-dimensional convolutional layer with kernel size 4, and o(.) denotes the SiLU activation function (Elfwing et al., 2018). The result of a matrix multiplied by scalar is the matrix with each element multiplied by that scalar.
Wgate W₂ E RdXP pixx,,W W. E RPxd We, WB E RdxN WA E Rdx1, E IR are trainable parameters of the layer, and P, N are hyperparameters. The authors call P the head dimension and N the state size. In practice, the weights of Wg, Wc are shared among different heads.
A.1 State Size
The authors of Mamba-2 always set P = 64, N = 128, and H = 2d/P. Thus, the state size of each Mamba-2 layer is HPN 2dN 256d. In transformer-based models, when using multi-headed attention, usually, the product of the number of heads H and head dimension Pequals the hidden dimension Therefore, the KV cache of a transformer-based model is 27d, which means that when using the same hidden dimension, the state of a Mamba-2 layer is equal in size to KV cache of 128 tokens. Additionally, the number of Mamba-2
Phttps://triie me/bloog/2224//amaa22
model/
13

---

Published as a conference paper at COLM 2025
State size parameters
Mamba-2
Mamba-1
30
RWKV-5/6
10
0
1000
2000
3000
Model size (M parameters)
Figure 12: The relationship between state size and model size of various RNN models in this paper.
layers in a Mamba-2 model usually twice the number of attention layers in a similar-sized Transformer model. Thus, the state size of Mamba-2 model is equal in size to KV cache of a vanilla Transformer model.
Compared to many other recurrent models (e.g., the RWKV series (Peng et al., 2023; 2024a), GLA (Yang et al., 2024a), and RetNet (Sun et al., 2023)), Mamba-2 does not have a state-less feed-forward network and has considerably more heads in each layer, making the state size much larger than other recurrent models. Compared to Mamba-1 (Gu & Dao, 2023), Mamba-1 uses N 16, which means that the state size in Mamba-2 8 times larger than the state in Mamba-1 model of roughly the model parameter count. Figure 12 shows the relationship between state size and model size of the RNN models in this study.
A.2 Short Convolution
The Conv(-) function in Mamba-2 is a one-dimensional convolutional layer applied to each channel separately. For i-th channel, it can be formulated as follows.
Vr
Wis jj R, 1
& denotes the kernel size (set to 4 by default). i denotes the channel index, n, denotes the number of channels. yr, IR denotes the i-th channel of the output vector at 1-th time step. x1, represents the i-th channel of the input vector 1-th time step. Wis E IR denotes the j-th value in the convolutional kernel for channel i.
This model component accepts the last 4 token embeddings at the input. Therefore, it also has state that contains information about the context, which we refer to as the convolutional state. To be concrete, due to information propagation through the layers, the short convolutional layer a function of the last 4L tokens. For the 370M model size, this length is 4 x 48 192. Therefore, we reasonably assume that this component contains much less contextual information relative to the recurrent state ht Thus, we have largely ignored this state in various discussions in this paper. However, we have also reported the distribution of the input to this short convolutional layer over time in Figure 17, for reference. As we can see, the convolutional state relatively stable over time (compared to the recurrent state).
B Passkey Retrieval Inference Parameters
Throughout the whole paper, greedy decoding, not just for reproducibility, but also because our preliminary results show that other decoding parameters give noticeably worse performance passkey retrieval.
14

---

Published as a conference paper at COLM 2025
We use 32-bit floating point precision for both model parameters and activations during inference, to ensure that precision errors do not introduce noise to the result. We have conducted some preliminary evaluations with BF16 and FP16 and found that there are noticeable differences with using FP16, but computing some activations, especially the A and a, with BF16 introduces error around 1c-3. However, the explosion of channels in the states is consistently observed despite this precision error.
B.1 Passkey Retrieval Prompt
The prompt that we use for the passkey retrieval task is as follows, using 34847 as the passkey for example, which is adapted from existing works (Zhang et al., 2024a). We also evaluate with slight variations to the template in preliminary experiments but do not observe considerable differences in the results.
There is important info hidden inside a lot of irrelevant text. Find it and memorize it.
The grass is green. The sky is blue. The sun is yellow. Here we go. There back again.
The grass is green. The sky is blue. The sun is yellow. Here we go. There and back again. The passkey is 34847. Remember it. 34847 is the passkey. The grass is green. The sky is blue. The is yellow. Here go. There and back again.
The grass is green. The sky is blue. The sun is yellow. Here we go. There and back again.
What is the passkey? The passkey is
We sweep different context lengths T (1K,2K, ..., 25.K),,,,,,,,, and for each length T, generate π prompts with evenly distributed needle positions, ie,, the i-th needle (i e ... n - 1)) of sample is inserted position 7' i/n 1, from the beginning.
C Mamba-2 with Modified A, on passkey retrieval
Ben-Kish et al. (2024) and GitHub user jzhang28 propose to improve Mamba's length generalization by reducing the value of A, Ben-Kish et al. (2024) propose a heuristic method for identifying which head to modify and how to modify Á, However, their method requires task-dependent tweaking, so we not consider comparing against it. jzhang28 propose to simply multiply A by a constant (they used 0.5). We apply this method and sweep different A, for the best passkey retrieval performance, but got near-zero accuracy across all passkey positions and context lengths.
D Passkey Retrieval Evaluation with Other Architectures
Here, we also evaluate RWKV-5, RWKV-6, and Mamba-1 (some popular and strong RNNs) on the passkey retrieval task. The result is reported in Figure 13, 14, and 15. We can see that length generalization failure is observed in Mamba-1, but it is less severe for RWKV-5 and RWKV-6. We hypothesize that this difference result of architectural differences and that the state size is smaller in RWKV-5 and RWKV-6 (see Figure 12.
hhttps://www gi thub. com/ http:/jhhnnnnnnognaaaage
15

---

Published as a conference paper at COLM 2025
Ans. Depth (%)
1K
1K
Depth
Accuracy
2K
0.75
2K
80
0.5
4K
0.50
4K
8K
8K
100
100
16k
16k
20
10
10a
10*
Context
Input Length
Context Len.
Input Length
(a) 1.5B
(b) 1.5B
(c) 3B
(d) 3B
Figure 13: The performance of RWKV-5 official checkpoints on the passkey retrieval task. Each curve in (b) and (d) represents the accuracy of retrieving the needle when it is within the last r tokens, with E {1K, 2K, 8KJ.
Depth (%6
50
1K
0.50
2K
100
0.25
16k
10
15
Context Len.
Input Length (K)
Figure 14: RWKV-6 1.6B result on the passkey retrieval task. The left plot shows the retrieval accuracy of the needle when appears in the last r {1K, 2K, AK,8K) tokens.
Passkey
Passkey
Passkey
Relatiee Anownr Postlon 4/5
Relalive Anrewer Position 156)
Relative Anowar Positon
24
Context Length (K takens)
Context Length (K tokenss
Context Length (K tokens))
(a) 790M
(b) 1.4B
(c) 2.8B
Figure 15: The performance of Mamba-1 official checkpoints on the Passkey task. We can clear exhibition of the inability to forget, similar to Mamba-2.
E Pre-Trained Checkpoints
The pre-trained checkpoints used in our experiments are given in Table 1.
More Experimental Details
ensure that the data contains as much long-term structure as possible, we filter out sequences with less than 4K tokens. Buckman & Gelada (2024) have shown that this critical for training effective long-context models. Although we train on sequences longer than 4K tokens, we do not use higher length threshold because the above threshold already removes about 67.6% of the data the original corpus. To train on longer sequences, simply concatenate sequences and delimit them with special EOS FEnd-of-Sequence) token. During evaluation, we sample documents longer than 16K tokens and concatenate them if they not long enough.
Data Processing
16

---

Published as a conference paper at COLM 2025
Model
Checkpoint URLs
co/Kk// rwkv-5-world-all-phh
https://hugginngffce.oo/
RWKV-5
RWKV-6
Mamba-1
30m
Mamba-2
Table 1: The pre-trained checkpoints used in our experiments.
Truncated Backpropagation Through Time In the vanilla Mamba-2, the states are initial- ized to zeros for each data sample. Instead, we initialize the states as the final state of the previous sequence. This is equivalent to concatenating multiple sequences, but stopping the backpropagation of gradients at certain intervals. This technique has been shown to help extend the context length of RNNs (Yang et al., 2024a) and alleviate the memory cost of caching activations for computing gradients. We employ this technique to make the distribution of the state's initial value (e.g,, the state before processing the first token ha) more varied. Based Yang et al. (2024a) and our preliminary tests, concatenate 12 sequences with this technique by default.
E1
More Hyperparameters
We use the WSD LR scheduler (Hu et al., 2024) with 10% decay steps. This scheduler is chosen because it is competitive with the commonly used cosine scheduler while allowing simple resumption from intermediate checkpoints, saving large amounts of computational resources. We report the result of the best checkpoint selection by validation on passkey retrieval.
We perform a hyperparameter search on learning rates, sweeping {le - 5, 2e 5, 5, 4,2e 4,5e 4, le 3), selecting the best performing one by validation on passkey retrieval Regarding the WSD scheduler, it warms up linearly for 1000 steps and decays linearly with 50K steps. This setup is inspired by the authors of WSD (Hu et al., 2024).
Other hyperparameters are kept as similar to the original papers for Mamba-2 as possible. That means we use 0.5M tokens per batch because we found this to give more stable results for continual pre-training instead of the 1M batch size from the original paper. Training is done mainly in BF16, with some activations in FP32 (in the same manner as the official implementation). The optimizer is AdamW, with 0.1 weight decay. Moreover, we use 1.0 gradient clipping.
All experiments are run on A800 80G, some are run with multiple nodes, and others with multiple GPUs on single node.
Model Configurations
For the models smaller than the 130M official checkpoint, we pre-train from scratch using the configurations reported in Table 2. We try to follow the same depth-to-width ratio found
WWiile the loss of many checkpoints was highly similar, their performance in passkey retrieval can differ lot.
17

---

Published as a conference paper at COLM 2025
Model size
State size
# Layers
Hidden dim.
# heads
Official checkpoints
780M
19.3M
48
1536
48
370M
12.9M
48
1024
32
130M
4.8M
24
768
24
Our checkpoints trained from scratch
84.6M
2.4M
12
768
24
47. 47.0M
1.6M
12
512
16
36.4M
0.8M
6
512
16
Table 2: The configurations of the models used in finding the passkey retrieval memory capacity as function of the state size.
in the official checkpoints, although the ratio is not entirely consistent in those checkpoints. Hyperparameters not mentioned are kept the same the 130M checkpoint.
G State Statistics over Context Length
Here, we provide more detailed result the inspection of state distribution over time.
Figure 16 shows the distribution hidden state h, the recurrent mechanism described in Eq. 2. Additionally, Br, Ct, and x; in Mamba-2 are generated with short channel-wisg convolutional layer with a kernel size of
B = (Conv (ut WBl)
C = (Conv [u, Wcl)
X; = oCConv Wx]
where o is the SiLU activation function. This function is also stateful because it operates on the last 4 tokens, therefore, also collect the statistics of this convolutional state and report them in Figure 17. As we can see, the convolutional states are much more stable compared to the recurrent states. This because only the last 4 tokens contribute this state which avoids the explosion as a result of cumulative sum.
H Length Generalization of Other Architectures
We additionally evaluate HGRN-2 (Qin et al., 2024) and RWKV-6 (Peng et al., 2024a) on the "newlines" prompt (string with only " "\n"") and find that they also exhibit severe performance degradation on the "newlines" prompt. The phenomenon is less severe in RWKV-6, which concurs with our argument that with longer training length, the model will learn to more robust forgetting mechanism, thus avoiding memory overload. Perhaps surprisingly, the increase in perplexity happens considerably before the context length reaches the training length for both models. We hypothesize that this is result of the training distribution, and that by continual training on data with more long-distane dependencies can alleviate this degradation.
I The "newlines" Prompt
In this paper, collect the statistics of the state computed on "newlines" prompt, prompt where every token is the newline token n")n"), as shown below.
However, we again emphasize that similar state distribution and model behavior are observed on prompts extracted from the pre-training corpus, the passkey retrieval task, or
18

---

Published as a conference paper at COLM 2025
Mean
Variance
Mean
Variance
0.015
3.002
0.6
Layer
1.5
Layer 8
Layer
Layer 9
0.001
0.4
Layer
Layer
0.005
1.0
Layer
Layer
Layer
Layer
Layer
Layer
Layer 7
-0.005
Layer 15
-2.001
10000
20000
10000
20000
10000
20
10000
20000
Token position
Tiken positon
Token position
Taiken postion
(a) Layer 0-7
(b) Layer 8-15
Mean
Varlance
Meann
Variance
0.003
0.002
Layer 16
Layer
0.001
Layer 17
0.002
0.3
Layer
Layer 25
0.000
Layer
Laye:
0.000
Lager
Layer
Layer
-0 002
Layer 23
0.002
Layer 31
10000
20000
10000
20000
10000
20000
10000 Thown position
20000
There position
Token position
Toben position
(c) Layer 16-23
(d) Layer 23-31
Mean
Varlance
0.006
0.004
0.002
1.5
1.0
Layer 34
Layer 32
Layer 33
0.00100
0.004
0.003
Laysr
Layer
layer
0.000
Layer 28
Layer
Layer
0.00350
Laya
Layar 45
Layer
Layer 39
-0.002
0.00000
10000
20000
10000
20000
10000
20000
30000
20000
Thken position
Token positton
Taken position
Tiken postion
(e) Layer 32-39
( Layer 40-48
Figure 16: The mean and variance of the hidden state of each layer of Mamba-2 370M, computed the "newlines" prompt (string with only "\n").
other randomly generated sequences. We have chosen the "newlines" prompt because the samples from the pre-training corpus are too short, and this prompt produces the most consistent and smooth layer statistics.
19

---

Published as a conference paper at COLM 2025
Mean
Variance
Mean
Variance
0.02
0.04
0.01
Layer 6
Layer 8
6.00
Layer
Layer1 1
Layer
Layer 9
-0.02
Layor
Layer
Layer
-0.02
Layer
2
Layer
Layer
Layer7 7
-0.04
Layer
0
10000
20000
10000
20000
10000
20000
10000 Token postion
20000
Toxion position
Tokan pos tio
Tikon position
(a) Layer 0-7
(b) Layer 8-15
Mean
Variance
Mean
Variance
Layer 16
Layer 24
Layer
0.02
Layer
0.04
3.5
0.00
0.8
Layer
Layer
Layer
0.6
Layer
Layer 22
Layer 30
-0.02
Layer 23
0.4
Layer 31
0.5
-0.04
0
30000
20000
10000
20000
10000
20000
0
10000
20000
18kan positiio
Token position
laken position
Taken posktin
(c) Layer 16-23
(d) Layer 23-31
Mean
Varlance
Mean
Verlance
0.03
0.04
1.5
Layer 32
Layer 40
Layer 33
Layer 41
0.00
0.6
1.0
Layer
Layer
Layer
Layer
-0.04
Layer 28
Layer 39
-0.02
00.03
Layer
Layer 46
10000
20000
10000
20000
10000
20000
10000 Thkon position
20000
Tiken position
Tokon position
Taken position
(e) Layer 32-39
(f) Layer 40-48
Figure 17: The mean and variance of the convolutional states (the representation of the last four tokens) of each layer in Mamba-2 370M, computed on the "newlines" prompt. We can see that the mean and variance are visibly more stable than the recurrent state.
10
1.6B
60
10
Perplexity
Loss
10
3
20
10-4
0
0K 10K 20K 30K
16K
Taken position (t)
Token position (t)
(a) Loss of RWKV-6 series on the "newlines" prompt as function of time.
(b) The perplexity of HGRN-2 1.3B on the "newlines" prompt as function of time.
20


Reviews:
[{"review_id": "5NBi9d3SVg", "reviewer": "Reviewer_rF9s", "summary": "This work investigates the forgetting behavior of Mamba-based models and shows that they struggle to effectively forget earlier tokens, even with built-in forgetting mechanisms. It demonstrates that the minimum training length required for the model to learn forgetting scales linearly with the state size, while the maximum context length for accurate retrieval scales exponentially with the state size. These observations may provide insights for the community to better understand Mamba's behavior.", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "1. The paper is clearly written and easy to follow.\n\n2. The finding regarding the relationship between minimum training length and state size is interesting and novel.\n\n3. While not entirely new, the overall observations could enhance the community’s understanding of Mamba and other linear attention models.", "reasons_to_reject": "1. The concept and causes of forgetting have been analyzed in previous works, e.g., DeciMamba and LongMamba. This work seems to provide a similar analysis of memory decay factors and then uses retrieval performance as a proxy to quantify forgetting. If the authors aim to emphasize the forgetting aspect, a deeper or lower-level analysis of the forgetting mechanism is expected, in addition to the empirical study similar to previous works.\n\n2. When fitting the scaling laws in Figures 9 and 11, too few data points are used, making the conclusions less convincing.\n\n3. The findings of this work are not entirely new. For example, the strong memory retention of early tokens, the relationship between state size and memorable context length, and the difficulty of extrapolating beyond the training length have been discussed in previous works like LongMamba and DeciMamba.\n\n4. The authors mention that *\"Mamba-2 with 370M parameters achieves near-perfect retrieval on a 256K context length, outperforming similarly sized transformer models,\"* but I did not find this set of experiments. For this comparison, is ROPE extension applied to transformers to ensure the performance gap is not due to improper handling of other design factors?\n\n5. More linear attention models beyond Mamba/Mamba2 are highly likely to be included in the discussion.", "questions_to_authors": "My questions have been included in the weakness section. I'm willing to adjust my scores if my concerns are properly addressed.", "extra_scores": {}}, {"review_id": "Css1q4tfIc", "reviewer": "Reviewer_pm8h", "summary": "The paper offers empirical evidence seeking to confirm conditions under which Mamba-2 models \"learn to forget.\"  The paper thus considers several retrieval-based experiments, aiming to demonstrate that there is a certain point during training that Mamba-2 models \"learn to forget\" and essentially overfit, and that if the training context length is sufficiently larger than Mamba-2's SSM dimension, the model retains the ability to forget.  Several experiments, model sizes, training context lengths, and SSM state sizes are presented to support the underlying hypotheses.", "questions": "", "limitations": "", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "Mamba-2 models have garnered significant interest over the past two years.  Quantifying exactly how and when they retain (or lose) the ability to forget information has the potential to impact the performance of these models for both training and downstream application performance.", "reasons_to_reject": "The underlying hypotheses and experiments need more refinement.  In particular, a discussion of the underlying additive state problem of RNNs like Mamba is warranted, as this is a known issue that has been addressed in other work (see [the Delta Rule section of the following](https://openreview.net/pdf?id=ayB1PACN5j) for more background).  However, exactly quantifying this for Mamba-2 models could still be a valuable contribution.  Yet, the experiments require further work; at times they lie at odds with one another.  In particular, the following two statements:\n > (1) Initially, the model demonstrates robust forgetting, retaining only the last k\n> tokens and forgetting earlier ones.  However, as training progresses, the model’s ability to\n>  forget diminishes while its recall of contextual information improves, resembling overfitting.\n>  This suggests that the model increasingly attempts to retain all available information within\n> the context. (2) We observe that forgetting occurs only when the training context length\n> exceeds the state’s capacity to retain all information, forcing the model to forget less relevant\n> details. Notably, larger states require longer training context lengths to effectively learn and\n> implement forgetting.\n\nClearly, the two statements contradict one another; if (2) is true, (1) does not get to the point of overfitting.  Similarly, given (1), what if the state is especially small in the non-overfitting region (thus invalidating statement 2).  The paper has many such contradictory statements/hypotheses which require ironing out.  There are also statement which require verification and missing details.  Finally, while the underlying phenomena is interesting, what demonstratable does this have on real applications?  The last question should be addressed at length.", "questions_to_authors": "- \"We observe that forgetting occurs only when the training context length exceeds the state’s capacity to retain all information\" <- Why would forgetting be correlated with context length?  Mamba's main innovation are the time-varying parameters, per time-step, so it effectively learns each time position during training.  Particularly since it is a first order process, learning should depend on the number of training steps, not on the training context length.\n\n- \"has a stable exponential memory decay (it converges to a constant value if the variables are fixed)\" <- Citation or please prove this statement\n\n- \"Therefore, we expect RNNs of such form to have a good retrieval accuracy on the last k tokens, and tokens earlier than that are forgotten\" <- Given Mamba's state-space equations, the argument that forgetting is dependent on length still doesn't make sense\n\n\"However, as the context length increases, Bs cannot be mutually orthogonal\" <- Why?  Please explain.  If this is a mathematical argument, please try to prove the assertion.\n\n\"The retention strengths of earlier tokens are always smaller than those of more recent tokens.\" <- With your set up, you can prove this claim.  Please prove this claim.\n\n\"This implies that the model has not learned to forget information (by producing a smaller αj), but it still has decent language modeling capabilities because the information of 8K tokens is typically not enough to overload the memory\" <- Why \"8K tokens is typically not enough to overload the memory?\" There is a break in logic between the experiment at the conclusions\n\nIn Equation 8, shouldn't r be w?\n\n\"which incurs minimal computational cost.\" <- This is not minimal, please change to neglible\n\n\"Similar observation can be found with any model size.\" <- Was this verified by the authors?\n\n\"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?  If in the appendix, please forward reference.\n\nFigure 8 is empirical evidence, more models seem warranted to support the underlying hypothesis.\n\n\"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\n\n\"The fact that the amount of information in the tokens exceeds the state’s capacity,\" <- the experiments are not convincing that this happening\n\n\" 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.\n\n\"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?\n\n\"The rightmost data point in the plot corresponds to Mamba-2 370M\" <- What are the other data points?\n\nBy 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?\n\n\"This is because the amount of information in the context does not increase with its length\" <- This contradicts many other statements in the paper", "extra_scores": {}}, {"review_id": "Toq3PZRNvQ", "reviewer": "Reviewer_DDLr", "summary": "This paper investigates the surprising failure of Mamba-based RNNs to extrapolate to longer sequences, despite their recurrent design. Contrary to the common belief that RNNs forget too much historical information, the authors argue that Mamba struggles because it forgets too little—retaining too much information due to oversized recurrent states. The paper provides compelling empirical evidence for this claim, showing that models trained on short contexts fail to learn effective forgetting mechanisms, leading to memory interference and degradation in performance beyond the training length. Notably, the authors introduce useful diagnostic concepts such as the Forget Threshold (minimum training length required to induce forgetting) and Maximum Recall Context Length (longest context length with reliable retrieval), offering valuable insights into the interplay between state size, training context, and memory retention.\n\nI recommend a clear acceptance of this work.", "questions": "", "limitations": "", "rating": 8, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "- The paper proposes a very interesting hypothesis—that RNNs fail to generalize to long contexts not because they forget too much, but because they retain too much due to overparameterized states. This sharply contradicts the common belief that RNNs are prone to forget historical information\n\n- The empirical evidence supporting this claim is thorough and convincing, including well-designed diagnostics (e.g., forget thresholds, recall lengths) and clear demonstrations of the failure modes and mitigation strategies.", "reasons_to_reject": "While the empirical observations and analyses are insightful, the paper lacks a viable and practical solution to the length generalization problem. There appears to be an inherent trade-off between short-context performance and long-context generalization: pretraining on longer sequences tends to increase perplexity at short-context positions, and enforcing forgetting during training degrades short-context modeling. Moreover, pretraining beyond the identified Forget Threshold is computationally impractical at the 700M scale and beyond.\n\nThe proposed sliding window workaround preserves short-context performance but effectively discards all historical information. This undermines one of the key motivations for using RNNs in the first place—their ability to retain long-range dependencies in a compact manner—rendering the approach less compelling than simply using sliding-window attention.\n\nFinally, the paper lacks an important baseline: post-training on longer sequences from an existing Mamba checkpoint for few tokens. This approach—similar to RoPE extension in Transformers—could potentially enable length generalization with minimal additional cost. Without evaluating such a baseline, it remains unclear whether the observed failure to generalize is a fundamental limitation of RNNs or merely a training artifact.", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "wUkNbEj7O9", "reviewer": "Reviewer_cVoT", "summary": "The paper analyses in depth the problem of Mamba where the performance drops drastically beyond context length. The paper made the following contributions in the understanding and improvement of this issue:\n- In section 3, the paper exhibits various evidence that the model does not learn to forget and this behaviour explains the rapid decrease when the context exceeds the training context length\n- In section 4, the paper examines the root cause of this memory retention and expose the relationship between training length and number of parameters characterizing the states.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "The paper made important observations for RNN/Mamba type of architecture for future improvements. The observation at least guides for state matrix design and the training process (train for larger context length) to improve Mamba type performance. \n\nThe paper has significant experiments to support its claim. To argue that the performance drop is because of the inability to forget requires significant amount of experiments and I argue that the paper performs well. For example, 3.2.1 shows that some tokens are never forgot with certain heads and 3.2.2 shows that simply re-introducing forgetting improves the performance. In general, I find the paper's claim insightful and are backed up with solid experiments.", "reasons_to_reject": "It seems that the paper can be presented more consistently in structure. Section 3 and Section 4 both are insights that are grounded with experiments. So it seems a bit strange to have a separate experiments section that only talk about threshold, maybe a better naming or organization can help readers.\n\nMinorly, line 277, there seems some error as Figure 9 does not have a second figure.", "questions_to_authors": "It seems that the paper provides possible explanations to Mamba/RWKV problems pointed out at https://arxiv.org/abs/2410.04727 where authors observe that at large length, the next token prediction performance falls below Mamba/RWKV that does not have access to the context.", "extra_scores": {}}]

Author rebuttal:

Meta-review: Maybe accept, if there is room.

This analysis paper makes and defends the interesting claim that Mamba suffers at long context length for a specific pair of reasons: because it not good at forgetting at those lengths, because of the limited context lengths at which it is trained.

There are split opinions by reviewers on this paper.   Supportive reviewers were enthusiastic and found the paper's proposal and development of a "Forget Threshold" and "Maximum Recall Context Length" (and the associated experimental data) to be valuable conceptual contributions for understanding the limitations of the Mamba architecture. AC agrees these seem like valuable contributions.

However, the main unresolved objection by negative reviewers is that the paper is insufficiently grounded in current work, in particular that both RWKV-7 and DeltaNet have advocated a "delta rule" that improves long context behavior at the cost of additional computational complexity, and that the simultaneous discoveries by LongMamba are not discussed.  The authors do not share these concerns because they deem those methods' (more computationally complex) non-elementwise update rule as "out of scope" for their paper, and also point out that LongMamba was not published at time of submission.

Both reviewer's questions and the author's point of view are reasonable.  I support accepting the paper if there is room.

If there is room to accept this paper, I urge the authors to add a more in-depth discussion of the delta rule, even if the main point is to help readers understand how it differs, and why it is out of scope for the authors' analysis.  And despite the timing of the simultaneous submission, I would ask the authors to consider adding a discussion of LongMamba and whether the authors' findings are consistent or not with the findings of that other concurrent work.  If the needed discussions are lengthy, it would be appropriate for supplemental material.
Decision: Accept