R e M amba: Equip Mamba with Effective Long-Sequence Modeling (2025)

Abstract: While the Mamba architecture demonstrates superior inference efficiency and competitive performance on short-context natural language processing (NLP) tasks, empirical evidence suggests its capacity to comprehend long contexts is limited compared to transformer-based models. In this study, we investigate the long-context efficiency issues of the Mamba models and propose ReMamba, which enhances Mamba’s ability to comprehend long contexts. ReMamba incorporates selective compression and adaptation techniques within a two-stage re-forward process, incurring minimal additional inference costs overhead. Experimental results on the LongBench and L-Eval benchmarks demonstrate ReMamba’s efficacy, improving over the baselines by 3.2 and 1.6 points, respectively, and attaining performance almost on par with same-size transformer models.

## introduction transformers @xcite , which form the backbone of most llms, encounter substantial challenges when dealing with long texts. the quadratic computational demands and the linear memory costs of the attention mechanism become prohibitive as the text length grows. this complexity poses a significant barrier to effectively modeling long texts, which is crucial for the development of llms. to address this, mamba is proposed as a solution @xcite , which utilizes a recurrent inference mode that ensures linear time complexity and compress information into the fixed state size. this results in constant memory demands during inference. furthermore, mamba eliminates the need for positional encoding, theoretically allowing it to handle inputs of any length, while performing competitively against transformers on downstream tasks. shortly after, mamba2 was introduced, simplifying the structured a matrix figure 1 : a comparison of pretrained mamba models and transformers of equivalent size across speed, shortcontext, and long-context performance metrics. speed is measured under conditions of 6k input tokens and 1k output tokens. "short scores" represents the average accuracy across six tasks (hellaswag, piqa, arc-e, arc-c, winogrande, openbookqa) evaluated within the lm evaluation harness @xcite . "long scores" corresponds to the average scores on the longbench-e benchmark @xcite . notably, all longbench evaluations employ a maximum token length of 2k to align with the model's training configuration. despite these advantages, some studies reveal that mamba models do not perform as well as expected when dealing with long texts reaching 2k tokens or more @xcite . as depicted in figure 1 , our experimental findings reveal that the pretrained mamba model surpasses pretrained transformers of comparable size, such as llama-3b @xcite , on short-context tasks. conversely, a substantial performance degradation is observed for mamba on long-context tasks relative to transformers. this performance disparity underscores a significant limitation of mamba models in practical long-context applications. this long-context deficency issue of mamba is usually attributed to its rnn-like nature. this kind of architecture exhibits limitations in preserving crucial information from earlier input sequences as the context length increases due to the fixed-size memory @xcite . hybrid architectures @xcite @xcite have sought to mitigate this issue by integrating attention mechanisms from transformers. however, these approaches often lead to decreased computational efficiency and increased memory consumption. the key challenge in mamba is the excessive degradation of distant information. remamba addresses this by employing an effective compression strategy that condenses the information and reduces the context length. specifically, it selects the top-k hidden states during the first forward pass and integrates them into the state space using mamba's selective mechanism in the second pass. remamba introduces minimal computational overhead (one additional forward pass) and maintains low, constant memory consumption. experimental results demonstrate that our approach significantly improves mamba's long-context performance, bringing it close to the performance of transformers. our remamba model achieves a 3.2 improvement over the baseline on longbench @xcite and 1.6 improvement on l-eval @xcite . furthermore, our methodology exhibits transferability to mamba2, yielding a 1.6 improvement on long-bench. ## preliminary study kv cache compression is widely used in transformers to reduce memory consumption and improve inference speed. however, prompt compression often results in performance degradation compared to the full context lengths generation in transformers. unlike transformers, mamba does not employ a kv cache; instead, it utilizes a fixed-size state space in each layer to preserve context memory. a potential issue with mamba is its tendency to forget distant information. in our preliminary study, we hypothesize that the state space update in mamba is insufficient for effectively compressing context information, and that techniques like prompt compression could help relieve this issue. to explore this, we apply a simple prompt compression method: we replace part of the context tokens with a few randomly selected hidden states from the last layer of mamba, creating a shorter prompt (referred to as random mamba). this approach is similar to that of @xcite , which used soft prompts for information compression. intuitively, this random compression method may lead to significant information loss and degraded performance. however, our results show that the average scores for different context lengths on longbench between normal mamba and random mamba are similar when both are trained on the same long-context dataset. furthermore, random mamba outperforms normal mamba at certain context lengths, as shown in figure 3 of 5.4.1. the random_select (sft) represents the fine-tuned random mamba, while mamba (sft) represents the fine-tuned vanilla mamba. this observation suggests that information loss in mamba when handling long contexts is substantial. to relieve this, we propose selective compression and selective adaptation through leveraging mamba's state space update mechanism. ## methodology remamba consists of two forward stages. in the first stage, three feed-forward networks are employed to help determine the significance of hidden states from mamba's final layer. these hidden states are selected based on their importance scores. the second stage integrates these compression hidden states with the input context, adapting mamba's selective mechanism to incorporate them into the state space. our proposed method draws some spirits from techniques employed in kv cache compression @xcite @xcite @xcite by leveraging the language model itself to aggregate information via hidden states and employing a scoring mechanism to select the most salient representations. nevertheless, different from transformers, remamba's compression strategy focuses on two key objectives: 1) compressing and selectively retaining crucial information to minimize information degradation, and 2) reducing the frequency of state space updates to further alleviate the information loss. the selection method employed in remamba is a simplified key-query-value-based approach, commonly used in retrieval-augmented generation and summarization tasks @xcite . in this context, we select the most important hidden states from the last layer of mamba to mitigate information loss during the compression process. ## conclusions this study investigates the long-context efficiency challenges posed by mamba models, hypothesizing that distant information within these models is subject to substantial degradation. in response, we introduce remamba, a novel approach that compresses and selectively preserves critical information during an initial forward pass. this compressed information is subsequently integrated into the state space during a second forward pass, capitalizing on mamba's inherent selective mechanism. notably, remamba incurs minimal computational overhead while substantially enhancing mamba's long-context performance, thereby offering a promising avenue for advancing the mamba model family.