ReMamba: Equip Mamba with Effective Long-Sequence Modeling

State space sequence models (SSMs) are a class of models that take inspiration from classical continuous dynamical systems. To efficiently compute the state transformations, structured SSMs impose specific constraints on the state transition matrix $\hat{A}$. Mamba models are a prime example, employing a diagonal structure for this matrix. In the one-dimensional situation, discrete structured SSMs transform sequences as follows:

The state space model above chooses the time-invariant $\hat{A}$ (state transition matrix) and $\hat{B}$ (input coefficient matrix) thus lacking expressiveness and flexibility. Mamba (Gu and Dao 2024) proposes to make $\hat{A}$ and $\hat{B}$ dynamically depend on inputs.

Recall that in one Mamba layer $l$ , SSM states $S$ are transformed as follows:

Furthermore, $\Delta_{t-1}^{l}\in\mathbb{R}^{H^{\prime}}$ is the discrete time step corresponding to the selective mechanism in Mamba, where $H^{\prime}$ is the intermediate hidden size. The continuous and discrete state transformation matrices at layer $l$ are given by $A^{l},\hat{A}^{l}\in\mathbb{R}^{H^{\prime}\times N}$, respectively. The continuous and discrete input coefficient matrices are denoted by $B_{t-1}^{l},\hat{B}_{t-1}^{l}\in\mathbb{R}^{N\times 1}$. The state size is represented by $N$. The discretization method for computing $\hat{A}$ and $\hat{B}$ is indicated by “discretize”. The vector $h_{t-1}^{{}^{\prime}l}\in\mathbb{R}^{H^{\prime}\times 1}$ and the SSM state is represented by $S_{t}^{l}\in\mathbb{R}^{H^{\prime}\times N}$. The symbol $\otimes$ denotes element-wise multiplication, and $\hat{B}_{t-1}^{l}\left({h^{\prime l}_{t-1}}\right)^{T}$ represents matrix multiplication. It is important to note that the definitions of $A^{l}$ and $\hat{A}^{l}$ presented here differ from their original definitions due to Mamba’s simplification to diagonal matrices.

Mamba improves sequence modelling ability of the state space model by enabling the model to focus on or filter out inputs into a sequential state.

Dao and Gu (2024) theoretically proves the connections between structured state space models and attention mechanisms. They also simplify structured matrix $\hat{A}$ further into scalar-times-identity structure and thus develop a new state space duality (SSD) framework with multi-head patterns similar to transformers. This modification trades in the expressiveness of the $\hat{A}$ matrix for faster training and enlarged state size.

Selective compression involves selectively compressing the input prompt by leveraging the final layer hidden states of the Mamba model to decrease state updates and consolidate information.

Suppose the sequence length is $L$ and the context token embeddings are $\{t_{i}\}_{i=1}^{L}$. We define the relative range to be compressed as $range:=(s,e)$, where $e=s+p$, with $s$ and $e$ denoting the relative start and end positions, respectively, and $p$ representing the relative length to compress. These values satisfy $0\leq s,p,e\leq 1$. The index set of the context to compress is $\mathcal{R}:=[S,E]$, where $S=L\cdot s+1$ and $E=L\cdot(s+p)$. Consequently, the length of the prompt to compress is $L^{\prime}=E-S+1$. For convenience, we use $\mathcal{R}$ to represent both the set of indices and the set of actual tokens within the context to be compressed. Furthermore, we define the compression ratio $\rho$ and compress the selected context $\mathcal{R}$ into $K:=|\mathcal{R}|\cdot\rho$ hidden representations.

In Figure 2, the compression hyperparameter settings are: $s=0.2$, $p=0.4$, $range=(0.2,0.6)$, $\mathcal{R}=[3,6]$, $\rho=0.5$, $K=2$. In our experiments, we find that $s=0$ yields the best results, which can be attributed to the casual language modeling nature of Mamba (this will be discussed in more details later).

As shown in the Stage 1 of Figure 2, we denote the last layer’s output hidden states as $\{h_{i}\}_{i=1}^{L}$, where each $h_{i}\in\mathbb{R}^{H}$ with $H$ representing the hidden size. We then transform the last hidden state $h_{L}$ into a query hidden state, namely $q$, through a feed-forward layer named $Query$. Additionally, the hidden states to be compressed, denoted as $\{h_{i}\}_{i=S}^{E}$, are transformed into $\{k_{i}\}_{i=S}^{E}$ via a $Key$ layer (this transformation is not shown in Figure 2). Finally, the cosine similarity scores, $Cos=\{cos_{i}\}_{i=S}^{E}$, are computed to serve as importance scores for the hidden states $\{h_{i}\}_{i=S}^{E}$. The calculation of $q$, $k_{i}$, and $cos_{i}$ is formulated as follows:

We select the top-$K$ hidden states ${h_{j}}$, where $j\in G$, from the hidden states $\{h_{i}\}_{i=S}^{E}$ based on their importance scores, denoted by $Cos$. The index set $G$ is defined as:

Note that the original order of these indices is preserved.

In our model, after selecting the top-$K$ hidden states $h_{j}$, we apply a feed-forward layer, $Value$, to project them into the token embedding hidden space:

Their corresponding cosine similarity scores are $\{cos_{i}^{\prime}\}_{i=1}^{K}$. We then replace the token embeddings $\{t_{i}\}_{i=S}^{E}$ ($\mathcal{R}$) with $\{v_{i}\}_{i=1}^{K}$. Consequently, the new input embeddings for Mamba are replaced by:

One significant challenge in using top-$K$ selection based on importance scores is its non-differentiability, which impedes the ability to train such models effectively. Here we propose a framework that integrates importance scores into the selective mechanisms of the Mamba model.

For hidden states (embeddings) that do not require compression in stage 1, namely $\{t_{i}\}_{i=1}^{S-1}$ and $\{t_{i}\}_{i=E+1}^{L}$, the standard Mamba algorithm is applied during the second forward pass. For embeddings at selected positions, specifically $\{t_{i}^{{}^{\prime}}\}_{i=S}^{S+K-1}$ or equivalently $\{v_{i}\}_{i=1}^{K}$, Equation 2a is reformulated as follows:

Following the forward encoding processes, standard causal language generation is applied using the Mamba architecture. During training, newly introduced parameters within the selective compression mechanism are optimized. These parameters, except for $\Theta$ which is initialized to all zeros, are initialized with a subset of the weights from the first layer’s in_proj matrix. Additionally, for parameters in Mamba, the dt_proj matrix is fully trained, while in_proj, out_proj, embeddings, and lm_head are updated using Low-Rank Adaptation (LoRA) (Hu et al. 2022). In our best implementation, to emphasize the significance of specific information, gradients flowing into the importance scores are scaled proportionally to these scores. This approach intuitively prioritizes the training of more critical representations.

Our model is designed for long-context question-answering tasks, necessitating a substantial corpus of long-context instruction tuning data. To this end, we leverage the OpenOrca dataset (Mukherjee et al. 2023) and LongAlpaca-12k (Chen et al. 2024). The former comprises a rich collection of ChatGPT-augmented FLAN data alignments, while the latter is a long-context alignment dataset. We initially filter long instruction tuning instances from OpenOrca and concatenate them with LongAlpaca. To accommodate device memory constraints, prompts are truncated to a maximum length of 6,000 tokens. This process yields approximately 200,000 long-context training examples. To augment training data diversity, the initial 300,000 standard instances from OpenOrca are incorporated. During training, the hyperparameter $s$ is fixed at 0. The hyperparameter $p$ is randomly sampled from the interval [0.1, 0.3], while $\rho$ is randomly sampled from the interval [0.05, 0.2].

We finetune the baseline Mamaba 2.8b model and our ReMamba model on the same dataset using LoRA with the same hyperparameter setting (r=32, others default). We also finetune a llama2-3b (Geng and Liu 2023) models for reference. Given the 2k maximum positional encoding limit of llama2-3b, we conduct fine-tuning experiments using two configurations: one matching the Mamba data and another with data truncated to a maximum length of 2K. Comparative results reveal superior performance for the truncated model, which is subsequently adopted as the default setting.

We conduct comparative analyses of our model against baseline Mamba2.8b (both of finetuned and pretrained) on the widely adopted LongBench benchmark (Bai et al. 2024) and LEval benchmark (An et al. 2023), which encompass a diverse set of challenging real-world long-context tasks. For consistency, the same prompt templates and greedy decoding configurations are employed across all models. To provide a reference point, the performance of a similarly sized transformer architecture (llama2-3b) is also included. It is important to note that due to llama2-3b’s 2k maximum positional encoding limitation, its performance in the 6k setting is subpar, necessitating the utilization of its optimal 2k configuration for evaluation.