DeciMamba: Exploring the Length Extrapolation Potential of Mamba

We thank the reviewer for their comments, and provide responses to the raised concerns below.

Concern 1:

---

Most of the related work on length generalization is for transformers. Is there any work on this for SSMs? I'm not sure exactly where this falls in the literature.

---

To the best of our knowledge, our work is the first to investigate length generalization in Mamba. We note, however, that several recent follow-up studies [1,2,3] have also explored this topic and were submitted to ICLR 2025.

Concern 2:

---

How does the DeciMamba architecture compare to Mamba at short context length tasks? Ie, is there a tradeoff between performance on short and long context tasks?

---

The operation of DeciMamba is identical to that of Mamba in short-context tasks. By definition, DeciMamba pools L_base tokens in the decimation layer. If the sequence length is already shorter than L_base​, it will pool all the tokens, leaving the sequence unchanged.

References:
[1] LongMamba: Enhancing Mamba's Long-Context Capabilities via Training-Free Receptive Field Enlargement. Anonymous, Submitted to ICLR 2025.
[2] MambaExtend: A Training-Free Approach to Improve Long Context Extension of Mamba. Anonymous, Submitted to ICLR 2025.
[3] Stuffed Mamba: State Collapse and State Capacity of RNN-Based Long-Context Modeling, Anonymous, Submitted to ICLR 2025.

We appreciate the reviewer’s feedback and provide responses to the raised concerns below.

Concern 1:

---

The experimental results are limited. The authors only combine the method with the Mamba model and verify the effectiveness of their method. It would be more convincing if the method can be validated on a broader range of SSM-based and linear-attention-based models.

---

We respectfully disagree with the assertion that the experimental results are limited, as they include multiple model sizes (130M, 1.4B, and 2.8B) across several datasets (PG-19, 16 datasets in LongBench, Passkey-Retrieval, and SQuAD v2) and tasks (language modeling, retrieval, question answering with or without fine-tuning, and summarization), as well as efficiency benchmarking. All of these findings support our claims.

Regarding other models ("...can be validated on a broader range of SSM-based and linear-attention-based models"), we would like to note that Mamba has recently gained significant attention in the field, receiving over 1,000 citations in less than a year, with more than 100 ICLR submissions focused solely on this architecture. We therefore consider our scope to be sufficiently broad and kindly ask the reviewer to reassess this point. Additionally, please note that similar concerns did not arise in the reviews of follow-up studies [4, 5] (also submitted to ICLR '25) that have explored length generalization in Mamba models.

In terms of application to other models, our method is based on the selective, data-dependent delta_t parameter within the selective state-space layers (S6) to measure token importance. Consequently, it cannot be directly applied to models that lack a per-token selective delta_t parameter (though it could be generalized to any data-dependent recurrent gate). Thus, our method cannot be applied to previous SSMs, such as S4, DSS, GSS, S5, Liquid SSMs, and others, which do not incorporate selective mechanisms. Please note that non-selective layers perform significantly worse than Mamba in NLP tasks, making them less relevant for context extension.

Regarding other linear attention models, our method relies on a data-dependent recurrent gate, making it infeasible for non-recurrent models. While some linear attention models could theoretically accommodate our approach (e.g., Griffin, via the r_t parameter in Eq. 1), this is outside the scope of our current work, as we consider Mamba to be broad enough.

To address the reviewer’s concerns, we will add a new technical subsection titled "Extensions to Other Models," detailing how insights from our work can be generalized to other models, such as Griffin, HGRN, RWKV, and others.

Concern 2:

---

The experimental results are not satisfactory. As shown in Table 1, DeciMamba fails to achieve more than a 10% LongBench score on most datasets. In contrast, models with a context window size of only 4k, as reported in the original LongBench paper, often perform better (e.g., Llama-2-7B-chat has an average score of 31 on English tasks).

---

Like many other context extension methods [1,2], DeciMamba does not aim to enhance the base model's capabilities (e.g., reasoning). Instead, it primarily helps the model maintain performance on longer sequences.

DeciMamba’s context extension capabilities are clearly demonstrated in the LongBench results. To highlight this, for every benchmark with a score higher than 10% (for either Mamba or DeciMamba), we compute the percentage improvement of DeciMamba:

In the 0-4k length range (non-long context), DeciMamba improves Mamba by an average of 22%. In sharp contrast, in the 4-8k range (long context), DeciMamba significantly enhances Mamba with an average improvement of 113%.

In the other tasks (e.g. HotpotQA, Qasper, TREC, etc.) Mamba underperforms even in the 0-4k range. Therefore, we do not expect the context extension method to achieve significant performance in these tasks as well.

Finally, Llama-2-7b was trained on 2 trillion tokens and has 7 billion parameters, whereas Mamba-2.8b was trained on 50 billion tokens and has 2.8 billion parameters. Given that in many cases the gap between 2.8b and 7b models holds emerging properties that are crucial for tasks such as LongBench [3], we find this comparison somewhat unfair.

---

DeciMamba's token dropping strategy can only be applied during the pre-filling stage and cannot benefit the decoding stage's token dropping. However, the authors did not discuss and present experimental results for this limitation in the paper.

---

We find that focusing exclusively on improving the pre-filling phase does not limit the practical impact of our method. To demonstrate this, we introduce additional pooling in the decoding phase and test this approach on the longest generation task in LongBench, GovReport, which requires summarizing long documents with an average length of about 10,000 tokens. We find that adding decoding pooling has a negligible effect on performance.

GovReport: Long Summary Generation With and Without Decoding Decimation

Where LB is LongBench; 0-4k, 4-8k, 8k+ are the three LongBench-e context length groups; and +DD is DeciMamba with additional Decoding Decimation. The additional decimation during decoding was performed by iteratively combining prefill and decoding for each generated chunk, where the chunk size is 50 tokens.

A possible explanation for the limited impact of adding decimation during decoding is that, despite being long (about 500 tokens), the generated summaries are relatively short compared to the number of tokens needed to cause limited ERFs. Empirically, for that to happen, the entire sequence should be at least two to three times longer than the length used to train Mamba.

Note also that, while the generated responses in GovReport are relatively long, for most of the remaining long-context tasks in LongBench - such as Multi-Document QA - the model generates only 10–30 tokens per answer. However, if a new task emerges where the generated sequence exceeds the ERF, decimation during decoding can be applied to that task.

In addition, there is a substantial body of work [1-11] published in top-tier conferences such as ICLR, NeurIPS, and ACL that focuses solely on enhancing the prefill stage of transformers without modifying the decoding process. Long-context handling during prefilling is a crucial area of research with significant applications, enabling, for instance, domain-specific LLM-based chatbots that can process entire books as context, Multi-Document Question Answering, etc.

We will add this experiment and discussion to the final version of the paper.

In light of the above, we kindly ask the reviewer to reconsider the assessment of this concern as a significant limitation.

References:
[1] MInference 1.0: Accelerating Pre-filling for Long-Context LLMs via Dynamic Sparse Attention. Jiang et al. NeurIPS 2024
[2] In-context Autoencoder for Context Compression in a Large Language Model. Ge et al. ICLR 2024.
[3] LongLLMLingua: Accelerating and Enhancing LLMs in Long Context Scenarios via Prompt Compression. Jiang et al. ACL 2024.
[4] Compressing Context to Enhance Inference Efficiency of Large Language Models. Li et al. EMNLP 2023
[5] LLMLingua: Compressing Prompts for Accelerated Inference of Large Language Models. Jiang et al. EMNLP 23
[6] CacheGen: KV Cache Compression and Streaming for Fast Large Language Model Serving. Liu et al. SIGCOMM 2024
[7] Prompt Compression and Contrastive Conditioning for Controllability and Toxicity Reduction in Language Models. Wingate et al. EMNLP 2022
[8] Learning to Compress Prompts with Gist Tokens. Mu et al. NeurIPS 2023
[9] Adapting Language Models to Compress Contexts. Chevalier et al. EMNLP 2023.
[10] Extending Context Window of Large Language Models via Semantic Compression. Fei et al. ACL 2024
[11] In-Context Former: Lightning-fast Compressing Context for Large Language Model. Wang et al. EMNLP 2024. [12] MambaExtend: A Training-Free Approach to Improve Long Context Extension of Mamba. Anonymous, Submitted to ICLR 2025. [13] LongMamba: Enhancing Mamba's Long-Context Capabilities via Training-Free Receptive Field Enlargement. Anonymous, Submitted to ICLR 2025. [14] Stuffed Mamba: State Collapse and State Capacity of RNN-Based Long-Context Modeling, Anonymous, Submitted to ICLR 2025

We thank the reviewer once again for their valuable feedback, and for engaging with our rebuttal and additional experiments.

We would like to respond to the major concerns you have raised.
We emphasize that, in addition to identifying and analyzing the limited Effective Receptive Field (ERF) phenomenon in Mamba, our work also proposes a solution that significantly improves Mamba's length generalization capabilities across a variety of tasks and model sizes (Section 5). We have not seen any reference for these results in your review, despite being 75% of the whole experiments section.
Specifically for LongBench, DeciMamba’s context extension capabilities are clearly demonstrated in the previous response, where we show a 113% improvement on context extension tasks. We kindly disagree that this improvement was ‘easily achieved’.

Furthermore, in the previous comment, you expressed concern that DeciMamba’s design limits the sequential mode of Mamba, and that pooling only during pre-fill is a limitation. We explained why the former cannot be true and provided an experiment for the latter, where we show that even for generation-heavy tasks, pre-fill pooling is a robust and effective solution.

We would also like to point out that the new concerns regarding the method’s sequential processing abilities, as well as concerns about the lack of pooling in the decoding phase, were not part of the initial review. Nevertheless, we have made extensive revisions to address all 5 weaknesses that you have identified. The paper has substantially improved in clarity and experimental support as a result of your valuable input:

1. Referenced our extensive evaluation suite (Section 5), which tested a large variety of model sizes over 20 benchmarks.
2. Fixed citation formats and figure caption locations.
3. Showed that limiting pooling to the pre-fill stage is a justified design choice, even in long generation tasks. This is grounded in empirical experiments and by related works from the literature (Appendix D)
4. Nevertheless, for our method, pooling can be naturally extended to the decoding phase (Appendix D)
5. Highlighted that non-Mamba, ssm-based models are not able to benefit from our method because they do not have a data-dependent, selective $\Delta_t$ parameter.

Should any concerns remain, we would be happy to address them further in the time left. If all concerns have been addressed, we kindly request the reviewer to consider raising their score.

We appreciate the reviewer’s feedback and provide responses to the raised concerns below.

Concern 1:

---

Limited scope of application: The method's restriction to the prefilling phase significantly limits its practical impact, especially considering the increasing demand for both long-text prefilling and generation in modern applications

---

We respectfully disagree with the assessment that focusing exclusively on improving the prefilling phase significantly limits the practical impact of our method. There is a substantial body of work [1-12] published in top-tier conferences such as ICLR, NeurIPS, and ACL that focuses solely on enhancing the prefill stage of transformers without modifying the decoding process. Long-context handling during prefilling is a crucial area of research with significant applications, enabling, for instance, domain-specific LLM-based chatbots that can process entire books as context.

Moreover, adapting our method to the decoding phase is straightforward. There are two simple extensions:

1. Threshold-Based Pooling Strategy
   Instead of using a top-K approach, we can select tokens individually based on whether their \delta_t norms exceed a certain threshold.
2. Combination of Prefill and Decode in a Repetitive Manner for Each Generated Chunk
   During the decoding of long sequences, DeciMamba can be applied using the parallel view of Mamba on both the previous history and the newly generated chunk. This allows us to update the current state with the state from DeciMamba. This approach incurs less than twice the computational resources compared to standard decoding, which is negligible [12], but it can enhance long-context capabilities.

A follow-up study [13], which builds on our findings and is also submitted to ICLR '25, employs the threshold-based pooling strategy. Their empirical results suggest that this approach, which extends DeciMamba to decoding, achieves performance similar to that of DeciMamba.

Concern 2:

---

Potential information loss: The token discarding approach may have unintended consequences in scenarios requiring comprehensive context understanding. This is particularly problematic in tasks like document question-answering, where discarded tokens during prefilling might be crucial for subsequent generation phases

---

This isn’t a bug, it’s a feature :)
The token discarding approach in our method isn’t about losing information; rather, it’s about selectively focusing on the most relevant context. By design, DeciMamba minimally intervenes in Mamba's operation: it discards tokens with low delta_t values, which the model tries to neglect anyway. The issue is that while the model flags these tokens as unimportant, it does not discard them, so when aggregated, they contribute to an early collapse of the state (Figure 2).

In addition, neural processing naturally emphasizes certain features over others. This is particularly true for state-based models like Mamba, which are designed to compress the context into a fixed-size recurrent state. Given that the number of tokens can vary, modern methods employ data-dependent compression that, by definition, filters out less relevant information in order to prioritize essential content.

Please consider that numerous context-compression techniques for transformers [1-11], published at top-tier venues, demonstrate that this is a robust approach rather than a limitation.

Regarding information loss: while the reviewer raises potential concerns with our method, it is important to clarify that Mamba naturally experiences severe information loss when processing sequences that are longer than those it was trained on, due to its limited ERF (see Section 3). In practice, DeciMamba mitigates this by reintroducing information that may have been forgotten by the model, effectively helping to recover the most important parts of the lost context.

Moreover, as demonstrated in Appendix B.1, DeciMamba maintains performance in tasks with complex inference demands, such as multi-turn dialogue.

Lastly, it is essential to note that DeciMamba is applied only to a subset of layers, excluding those at the beginning of the model, allowing information from unselected tokens to propagate through earlier layers via the remaining tokens.

For all of these reasons, we kindly request that you reconsider this point.

Concern 1:

---

While it is promising to see that the threshold-based pooling strategy has been explored in a follow-up study [13], relying on unpublished or concurrently submitted work to address core limitations in the current submission is not ideal.

---

We thank the reviewer for their comment. To address this concern in the best possible way, we introduced additional pooling in the decoding phase and tested this approach on the longest generation task in LongBench, GovReport, which involves summarizing long documents with an average length of 10,000 tokens. We find that adding decoding pooling has a negligible effect on performance.

GovReport: Long Summary Generation With and Without Decoding Decimation

Where LB is LongBench; 0-4k, 4-8k, 8k+ are the three LongBench-e context length groups; and +DD is DeciMamba with additional Decoding Decimation. The additional decimation during decoding was performed by iteratively combining prefill and decoding for each generated chunk, where the chunk size is 50 tokens.

A possible explanation for the limited impact of adding decimation during decoding is that, despite being long (about 500 tokens), the generated summaries are relatively short compared to the number of tokens needed to cause limited ERFs. Empirically, for that to happen, the entire sequence should be at least two to three times longer than the length used to train Mamba.

Note also that, while the generated responses in GovReport are relatively long, for most of the remaining long-context tasks in LongBench - such as Multi-Document QA - the model generates only 10–30 tokens per answer. However, if a new task emerges where the generated sequence exceeds the ERF, decimation during decoding can be applied to that task.

In light of these results, as well as our previous arguments regarding the applicability of pre-fill compression, we kindly ask the reviewer to reconsider classifying this concern as a core limitation.

.

Concern 2:

---

While the rebuttal provides a compelling theoretical rationale for the token discarding approach, additional empirical evidence is necessary to fully address concerns about its practical implications for tasks requiring comprehensive context understanding, such as long-document QA where documents are followed by questions.

---

We refer the reviewer to Section 5 (LongBench) and, specifically, to Table 1, where we extensively evaluate our method's performance across LongBench’s 16 benchmarks. These benchmarks include tasks such as Long-Document QA, Multi-Document QA, Summarization, Long-Context Few-Shot Learning, and others. The results show that DeciMamba consistently outperforms the baseline Mamba model, demonstrating the effectiveness of our approach. In particular, we highlight DeciMamba’s improvement on LongBench-E in the 4-8K range, which underscores the impact of our context extension.

.

Concern 3:

---

Addressing the "Questions" subsection

---

We thank the reviewer for this reminder and greatly apologize for missing these questions. We provide full responses below.

---

Regarding the Mamba Mean Distance results in Fig 4, the findings don't definitively prove that Mamba's long-text capability decreases with length, as this metric wouldn't necessarily scale linearly with sequence length. Are there similar studies for attention-based architectures?

---

We thank the reviewer for their comment. We agree that the Mamba Mean Distance experiment does not demonstrate that Mamba’s long context capabilities decrease with context length (The performance drop is shown in Section 5: Passkey Retrieval, Multi-Document Retrieval/QA, Perplexity, and LongBench).

The purpose of Figure 4 is to show that the Mamba Mean Distance is an adequate metric for measuring context utilization, as discussed in lines 213-215. We also agree that the figure’s title is misleading, so we have replaced it with: ‘Mamba Mean Distance Quantifies Effective Context Utilization’. Thank you for your constructive feedback, which has improved the paper. We hope this resolves the issue.

We thank the reviewer once again for their valuable feedback, and hope that our response has addressed all of the remaining concerns.
We would like to respond to the most major concerns you have raised:

---

Limited scope of application: The method's restriction to the prefilling phase significantly limits its practical impact

---

Potential information loss: The token discarding approach may have unintended consequences in scenarios requiring comprehensive context understanding. This is particularly problematic in tasks like document question-answering,

---

In our last response, we showed that neither of these limits our approach in practice, even in long-generation tasks. This was demonstrated through dedicated experiments (Appendix D and Table 1). We kindly request that you reconsider these points, as we have shown evidence that they are not significant weaknesses.

Finally, we have made extensive revisions to address all 5 remaining concerns you raised. As a result of your valuable input, the paper has improved both in clarity and in experimental support:

1. We showed that limiting pooling to the pre-fill stage is a justified design choice, even in long generation tasks. This is grounded in empirical experiments and by related works from the literature (Appendix D)
2. Nevertheless, for our method, pooling can be naturally extended to the decoding phase (Appendix D)
3. We fixed the title of Figure 4 to make it more accurate.
4. We provided an explanation for the counter-intuitive relationship between distance and performance in Figure 1 (Passkey Retrieval).
5. We justified model size selection in experiments and discussed expected impact on generalizability to very large models.

If all concerns have been addressed, we kindly request the reviewer to consider raising their score. Should any concerns remain, we would be more than happy to address them further in the time remaining.

Minor Point 4:

---

Comparison with Transformers

---

We find that vanilla Transformers of equivalent size (trained on the same dataset with a similar training recipe), have inferior length generalization abilities compared to Mamba. This is evident in all long-context tasks tested:

Passkey Retrieval:

The setting is the same as the one in the paper. All models were trained on sequences of length 2k. For each context length we test performance for 5 different needle locations.

MultiDocument Retrieval:

The setting is the same as the one in the paper. All models were trained on sequences of length 2k (10 documents). For each context length we test performance for 100 different queries.

Zero-Shot Perplexity:

The setting is the same as the one in the paper. All models were pre-trained on sequences of length 2k (we remind that DeciMamba is applied directly without any tuning). Inf replaces any result larger than 100.

For LongBench, the equivalent Transformer model (Pythia-2.8B) repeatedly causes an OOM error on our GPU (A6000, 48GB of RAM).

Minor Point 5:

---

Additional clarifications and edits

---

We greatly appreciate the detailed feedback. We incorporated the suggested edits into the revised manuscript.

If all concerns are met, we kindly ask the reviewer to increment their score. Nevertheless, if any concerns remain unresolved, we would be happy to further address them.

References:
[1] Long Range Arena: A Benchmark for Efficient Transformers, Tay et. al, ICLR ‘21
[2] Never Train from Scratch: Fair Comparison of Long-Sequence Models Requires Data-Driven Priors, Amos et. al, ICLR ‘24
[3] Simple linear attention language models balance the recall-throughput tradeoff, Arora et. al, ICLR ‘24 (WS)

We thank the reviewer for their detailed comments and suggestions and provide responses to the raised concerns below.

Concerns 1+2:

---

The approach looks a bit hacky in that it introduces multiple novel hyperparameters: the decay factor, the maximal length of the sequence after the first decimating layer and the number of layers to decimate. And it does not seem very clear how to make these choices without a gridsearch.

---

Have you tried to do an ablation with respect to the number of layers to be decimated or with respect to the decay rate?

---

Maximal Length Parameter (L_base):
In practice, this is the only parameter we sweep. We typically scan 3 or 4 values that are similar in magnitude to the context length used during training (L_train). For example, if the model is trained on 2000 tokens, we check L_base = 1600, 2000, 2400, 2800. We note that the scores do not deviate significantly during the sweep. For instance, across the 16 LongBench benchmarks, the scores deviate by an average of 10%.
The reason for selecting L_base close to L_train originates from an assumption on the long-context data: short training sequences and long evaluation sequences have similar information content and differ mainly by the amount of noise in the sequence. Under this assumption, there are at most L_train important tokens; hence, it makes sense to pool this number of tokens, regardless of the amount of noise. We find this assumption quite reasonable for many long-context tasks, such as retrieval, multi-document question answering, and next-token prediction. For example, in multi-document retrieval/QA, only one document is relevant to the query, regardless of how many random documents we append to the context. Another example is next-token prediction, which is usually very local and does not benefit much from global interactions. Another reason for selecting L_base close to L_train is that there are no ERF issues when processing sequences of length L_train, as the global layers (starting from the first decimation layer) are trained on sequences of the same length.

Number of Layers to Decimate and Decay Rate Parameters:
The main goal of these parameters is to improve efficiency, not performance. In the paper, we introduce the option to compress the sequence further by applying additional DeciMamba layers (by setting number_of_decimation_layers > 1 and decay_factor < 1). We emphasize that this has negligible effects on performance, as shown in the requested ablation below, which sweeps the two parameters for the passkey retrieval task. In addition, this is supported by other results, such as Multi-Document Retrieval and LongBench, which were run with this option disabled. We further emphasized this in the revised manuscript.

Decay Factor (Beta) Sweep:

BL is the performance of the baseline Mamba model.
The setting is the same as the one in the paper. We use L_base=2000 and num_deci_layers = 9.

Number of Decimation Layers Sweep:

BL is the performance of the baseline Mamba model.
The setting is the same as the one in the paper. We use L_base=2000 and beta=0.5.

.

Concern 3:

---

Writing should be improved: I found the decimation strategy in Decimamba (Section 4) pretty hard to follow. I think that this should be better explained.

---

We thank the reviewer for their feedback. We will improve the mentioned subsection and clarify the explanation. Specifically, we will clarify at the beginning that our goal is to address ERF issues by generating the most informative hidden state. To achieve this, we focus on prioritizing key content through a data-dependent context compression method and provide a detailed explanation of why Delta_t serves as the most relevant indicator of token importance for future predictions. We incorporated this explanation into the revised manuscript (L300-312).

---