LongMamba: Enhancing Mamba's Long-Context Capabilities via Training-Free Receptive Field Enlargement

Authors: Zhifan Ye, Kejing Xia, Yonggan Fu, Xin Dong, Jihoon Hong, Xiangchi Yuan, Shizhe Diao, Jan Kautz, Pavlo Molchanov, Yingyan Celine Lin
Primary area: foundation or frontier models, including LLMs
Keywords: Large Language Models, State Space Models, Long Context Understanding

Abstract:
State space models (SSMs) have emerged as an efficient alternative to Transformer models for language modeling, offering linear computational complexity and constant memory usage as context length increases. However, despite their efficiency in handling long contexts, recent studies have shown that SSMs, such as Mamba models, generally underperform compared to Transformers in long-context understanding tasks. To address this significant shortfall and achieve both efficient and accurate long-context understanding, we propose LongMamba, a training-free technique that significantly enhances the long-context capabilities of Mamba models. LongMamba builds on our discovery that the hidden channels in Mamba can be categorized into local and global channels based on their receptive field lengths, with global channels primarily responsible for long-context capability. These global channels can become the key bottleneck as the input context lengthens. Specifically, when input lengths largely exceed the training sequence length, global channels exhibit limitations in adaptively extend their receptive fields, leading to Mamba’s poor long-context performance. The key idea of LongMamba is to mitigate the hidden state memory decay in these global channels by preventing the accumulation of unimportant tokens in their memory. This is achieved by first identifying critical tokens in the global channels and then applying token filtering to accumulate only those critical tokens. Through extensive benchmarking across synthetic and real-world long-context scenarios, LongMamba sets a new standard for Mamba’s long-context performance, significantly extending its operational range without requiring additional training. Our code is available at https://github.com/GATECH-EIC/LongMamba.

Decision: accept (poster)
Reviewer scores: 8, 5, 8, 6

Reviews and discussion:
<thread>
creation date:
1730538345616

thread_type:
official_review

<thread_object>
creator: reviewer_1ahb

<content>
summary:
This paper introduces LongMamba, which is a training-free method to enhance Mamba long context performance. It identified that the channels in Mamba have distinct receptive field lengths: local and global. And it hypothesize that this failure of the global channels
to capture global information under extended sequence lengths is the root cause of Mamba’s poor performance in long-context understanding. LongMamba expands the global channels' receptive fields by removing less important tokens.

strengths:
1. The discovery and explanation of  Mamba models’ limited performance in long contexts are quite interesting.

2. The paper conducts extensive experiments, and LongMamba show its effectiveness

weaknesses:
1. While I really appreciate the discovery and root analysis of Mamba's limitations regarding long contexts, I find that LongMamba's approach—pruning less important tokens to increase the global channel's receptive field—seems somewhat incremental compared to DeciMamba.

2. LongMamba shows some improvements for long contexts, but it still seems to be somewhat behind transformer-based LLMs designed for long contexts. I understand that the authors are currently using a training-free approach. If training were allowed, would the performance be further enhanced?

3. It would be more convincing if the authors can provide an evaluation on the Ruler benchmark[1].

[1] https://github.com/NVIDIA/RULER

questions:
see the weakness section
</content>
</thread_object>

</thread>
<thread>
creation date:
1730625775515

thread_type:
official_review

<thread_object>
creator: reviewer_p417

<content>
summary:
This paper introduces LongMamba, a training-free techniques to enhance Mamba model's ability to handle long sequences beyond the training context lenght.  

The authors discovered that Mamba's hidden state channels can be categorized into two types: local channels that focus on nearby context, and global channels that attend to the entire input sequence. 
The key limitation they identified is that global channels fails to be functional when processing sequences longer than the training length and they hypothesize the exponential decay is the root cause.   

To this end, they propsoe LongMamba that addresses this by enlarging receptive fields of global channels by filtering out the less important tokens while maintaining similar decay distributions as during training.
Through extensive experiments on tasks like Passkey Retrieval, Document Retrieval, Language Modeling, and LongBench, LongMamba shows promising result without additional training and outperform vanilla mamba and previous methods.

strengths:
- The finding about the categroies of channels of Mamba models(local and global) is interesting and meaningful, which could bring insight to the community for better understanding the Mamba model's behavior.

- The experiment results are strong as it outperforms the vanilla Mamba model by large margin.

- This paper is well-written and easy to follow, for example, they reveal the "categroies of channels" phenomena with sufficient analysis and visualization and figure 2 (the method overview) is tidy and clear

weaknesses:
The paper is generally good but I believe the experiments section can be improved by adding direct comparsion with other long-context models beyond Mamba variants.
I find this is necessary because the core motivation of this work is to address the poor long-context ability of Mambas compared to the transformer-based model, as said in the introduction section. Therefore, it is useful to learn how much LongMamba helps bridge the performance gap. More specifical, the author may follow the [Mamba](https://arxiv.org/abs/2312.00752)'s main experiment setting and compare the LongMamba to varients of transformers (like [Pythia](https://arxiv.org/abs/2304.01373) ) or hypird model

questions:
(1) How does the LongMamba perform on the code categroy of LongBench?

(2) Does the token filtering process bring extra compute or memory cost?

(3) How does the "local and global channels" behave on larger Mamba model?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730690257875

thread_type:
official_review

<thread_object>
creator: reviewer_xmy1

<content>
summary:
The paper introduces LongMamba, a method to extend the context length of Mamba models without additional training, focusing on long-context understanding for state-space models. Mamba models, though efficient in handling long contexts due to their linear complexity, struggle with context lengths far beyond training sequences due to exponential decay in the attention value, especially toward early tokens. LongMamba addresses this by categorizing channels into local and global based on receptive field lengths, then enlarging global channels' receptive fields via adaptive decay adjustments through token filtering. This modification allows LongMamba to maintain performance over long contexts without the computational cost of fine-tuning, making it a practical improvement over previous methods like DeciMamba.

strengths:
1 - The paper addresses a critical challenge in long-context modeling, particularly for SSMs like Mamba, which suffer from limited long-sequence capabilities despite their efficient state-space mechanics. Enhancing this capability without retraining aligns well with real-world demands for scalable, efficient language models.


2 - Devising a “training-free” context extension technique is an important aspect of the paper, as directly training the SSMs on longer context requires significant compute/resource and expansive datasets. 


3 - Extensive benchmarks show that LongMamba consistently outperforms vanilla Mamba on long-context tasks, including language modeling on the PG-19 dataset and passkey/document retrieval. LongMamba’s results in real-world tasks, such as LongBench, also showcase its adaptability

weaknesses:
1- A substantial limitation is that LongMamba relies heavily on the DeciMamba[1] framework, raising questions about its independent novelty. Although LongMamba innovatively extends global channels, the reliance on DeciMamba's token handling might limit its perceived originality. In particular, the idea of token filtering based on the averaged delta_t value was proposed by DeciMamba. LongMamba merely improves upon this idea by limiting the scope of this averaging to some specific channels (global channels).

2 -The paper does not systematically address threshold selection for distinguishing global versus local channels. This selection is left to empirical values, which may lack generalizability across different datasets or models. A more formal approach to tuning this threshold would strengthen LongMamba's applicability. This is of importance, as global and local channels are integral parts of the LongMamba framework. 


3 - The method for handling global channels’ decay distribution adjustment is not fully elaborated. It remains ambiguous whether averaging across global channels fully mitigates decay issues over extended contexts, potentially affecting LongMamba’s consistency in varied contexts (Figure 2c, Decay Distribution Adjustment).

questions:
1. Could the authors please address the points mentioned in the weaknesses section?

2. Is a similar exponential decay observed in the Mamba2[2] model? If so, does LongMamba improve the long-context understanding in this architecture as well?

3. Would the proposed approach generalize to other SSMs beyond Mamba, including hybrid Transformer/SSM architectures, or is it specifically tailored to Mamba’s architectural nuances?

4. LongBench is a comprehensive benchmark. How does LongMamba perform on other LongBench tasks, namely 2wikimqa, TREC, TriviaQA, LCC, and RepoBench-P?

5. Have the authors considered alternative strategies for decay adjustment besides token filtering? If so, why was filtering chosen over other methods? 

References:

[1]Ben-Kish, Assaf, et al. "DeciMamba: Exploring the Length Extrapolation Potential of Mamba." arXiv preprint arXiv:2406.14528 (2024).

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

</thread>
<thread>
creation date:
1730820692609

thread_type:
official_review

<thread_object>
creator: reviewer_gwxu

<content>
summary:
The paper proposes a simple modification to Mamba architecture where they prevent channels from exponential decaying by filtering out tokens from the training sequence if the update is small than threshold , for smaller updates and keepin them as is for larger updates. This allows to prevent catastrophic forgetting, and surprisingly allows to reuse existing model on a significantly longer context without finetuning. 
The modification seem pretty straightforward where they essentially identify "Global channels" that seem to be particularly prone to exponential decay and change the SSM update procedure to ignore some of the smaller update to reduce the amount of exponential decay for longer contexts.

strengths:
The main result is very interesting: they show that existing SSM architectures acan be expanded to much longer contexts without having to fine tune it. Overall  a simple mechanism that seemingly enables Mamba to generalize to much longer contexts. Great introduction to Mamba models.

weaknesses:
See questions. I think the empirical analysis section would benefit from some more detail study  of distribution of local vs global channels receptive field (see Q5 below). Sec 5 would benefit from expanding how A' are computed.

questions:
1. End of Section 5: line 333, i think you meant Eq 16 not E 18?
2. It is not clear what A'_t are if updates is greater than threshold.  (e.g. dual of Eq 17)
3. How is g(S) computed? Is it computed to satisfy equation 16?
4. Local/Global channel: is the number 0-300-600-900-1200 are just percentiles of a fixed model? It would be better to specifically say that.
5. Can you include a graph showing the cumulative graph of receptive fields for each each channel? Especially this would be interesting for training seq.len and a much longer seq len?
</content>
</thread_object>

</thread>
<thread>
creation date:
1733097935476

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
Dear ACs and Reviewers,

We sincerely appreciate the time and effort of all reviewers in providing thoughtful and constructive feedback, especially considering the high volume of submissions at a top-tier conference like ICLR. We also hope that you and your families have had a joyful Thanksgiving holiday.

We are encouraged by the reviewers’ recognition of LongMamba for its key insight, extensive experiments, and clear writing. In particular, we are pleased that our insight—highlighting the exponential hidden state decay of global channels as the primary factor constraining Mamba’s receptive field and performance on long-context tasks—has been well received.

In addition to the commendations, we also received insightful questions from reviewers and have addressed all of them in our response to each reviewer. We welcome and look forward to further discussions with all reviewers, especially with tomorrow being the rebuttal deadline.

In the sections below, we provide a brief summary of our responses, particularly highlighting two common insightful questions from reviewers: the benchmark of LongMamba-enhanced Mamba with Transformers and the novelty of LongMamba.

---

## **Benchmark LongMamba-enhanced Mamba with Transformer Models:**

Following the suggestion of Reviewer P417, we benchmarked a 7B Mamba model against Transformer models of similar scale during the rebuttal period. **While the vanilla Mamba model showed lower performance compared to Transformer models, LongMamba demonstrated a significant improvement, narrowing the performance gap** and even surpassing certain Transformer models. Detailed experimental results are included in our response to Reviewer P417’s W1.

---

## **Novelty of LongMamba:**

**The core contribution of LongMamba lies in its insight** that the exponential hidden state decay of global channels is the primary cause of Mamba’s constrained receptive field and limited performance on long-context tasks. Building on this insight, LongMamba specially handles the identified global channels with selective token filtering, which has demonstrated superior performance compared to the prior work DeciMamba [a].

Moreover, **we have validated the generality and applicability** of our insight through the following experiments conducted during the rebuttal period:

1. We have implemented an alternative method to instantiate our insight (i.e., “decay normalization”, as discussed in our response to Reviewer xMy1’s W1 & Q5), which also enhances Mamba’s long-context understanding capabilities.
2. Applying LongMamba to the Mamba-2 [b] model and the hybrid Transformer/SSM model Hymba [c] also improves their long-context understanding capabilities. Detailed experimental results are provided in our responses to Reviewer xMy1’s Q2 and Q3, respectively.
</content>
</thread_object>

</thread>


Metareview:
metareview:
All reviewers found the work interesting, with strong results. The proposed technique for extending Mamba’s receptive field is both novel and clearly explained, supported by thorough experiments on multiple tasks. In particular, mitigating exponential decay by filtering out smaller updates demonstrates a straightforward yet impactful solution to catastrophic forgetting, without finetuning. Their experiments highlight the effectiveness of this approach in real-world settings. The paper is well-written, and its insights on local and global channels are valuable. Overall, a clear accept.