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

2025-04-22

AI Summary

Mamba exhibits significant performance degradation on ultra-long-context understanding tasks, primarily due to its limited global channel receptive field, which causes critical information decay in hidden states. Method: This work first identifies that Mamba’s hidden states can be decomposed into local and global channels—and reveals that global channels constitute the bottleneck for long-context modeling. To address this, we propose a training-free, architecture-preserving receptive field expansion method: leveraging channel-wise receptive field analysis, it dynamically identifies and retains key tokens within global channels, enabling selective filtering without any parameter updates. Contribution/Results: The method incurs zero fine-tuning cost and zero training overhead. It substantially outperforms vanilla Mamba on both synthetic and real-world long-text benchmarks, significantly extending effective context length. This establishes a novel paradigm for long-range modeling in state-space models.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Natural Language Processing: (Large) Language Models
• Search and Optimization: Learning to Search

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

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.

Problem

• Enhancing Mamba's long-context understanding without training
• Addressing Mamba's underperformance in long-context tasks
• Improving global channel memory decay in Mamba models

Innovation

• Training-free receptive field enlargement
• Token filtering for critical token accumulation
• Enhancing Mamba's long-context capabilities