DeciMamba: Exploring the Length Extrapolation Potential of Mamba
Authors: Assaf Ben-Kish
Itamar Zimerman
Shady Abu-Hussein
Nadav Cohen
Amir Globerson
Lior Wolf
Raja Giryes
Venue: ICLR 2025 Poster (2025)
Keywords: Mamba
Long context
Context extension
Long-range language modeling
Primary area: other topics in machine learning (i.e., none of the above)
TL;DR: We explore Mamba's length extrapolation potential, identifying its limitation due to a restricted effective receptive field, and introduce DeciMamba, the first context-extension method specifically designed for Mamba.

Abstract: Long-range sequence processing poses a significant challenge for Transformers due to their quadratic complexity in input length. A promising alternative is Mamba, which demonstrates high performance and achieves Transformer-level capabilities while requiring substantially fewer computational resources. In this paper we explore the length-generalization capabilities of Mamba, which we find to be relatively limited. Through a series of visualizations and analyses we identify that the limitations arise from a restricted effective receptive field, dictated by the sequence length used during training. To address this constraint, we introduce DeciMamba, a context-extension method specifically designed for Mamba. This mechanism, built on top of a hidden filtering mechanism embedded within the S6 layer, enables the trained model to extrapolate well even without additional training. Empirical experiments over real-world long-range NLP tasks show that DeciMamba can extrapolate to context lengths that are significantly longer than the ones seen during training, while enjoying faster inference. We will release our code and models.

Published as a conference paper at ICLR 2025
DECIMAMBA: EXPLORING THE LENGTH EXTRAPO- LATION POTENTIAL OF MAMBA
Assaf Ben-Kish', Itamar Zimerman', Shady Abu-Hussein', Nadav Cohen', Amir Globerson Lior Wolf', Raja Giryes' Tel Aviv University, Gooole Research
ABSTRACT
Long-range sequence processing poses a significant challenge for Transformers due to their quadratic complexity input length. A promising alternative Mamba, which demonstrates high performance and achieves Transformer-lveel capabilities while requiring substantially fewer computational resources. In this paper we explore the mingth-generllizatie to be relatively limited. Through a series of visualizations and analyses we iden- tify that the limitations arise from restricted effective receptive field, dictated by the sequence length used during training. To address this constraint, we intro- duce DeciMambu, context-extension method specificallyy designed for Mamba. This mechanism, built on top of a hidden filtering mechanism embedded within the S6 layer, enables the trained model extrapolate well even without additional training. Empirical experiments over real-world long-range NLP tasks show that DeciMamba can extrapolate to context lengths that are significantly longer than the ones seen during training, while enjoying faster inference.
capabilities of Mamba, which we find
INTRODUCTION
Lengthy sequences, which can span up to millions of tokens, are common in real-world applications including long books, high-resolution video and audio signals, and genomic data. Consequently. developing Deep Learning (DL.) sequence models capable of effectively managing long contexts is critical objective. Transformers (Vaswani et al. (017)), despite their current dominance in general DL tasks, still face challenges in processing long sequences. Specifically, their quadratic complexity in sequence length makes them computationally demanding, restricting the ability to train them over long sequences and very large datasets.
In recent years, substantial efforts have been made in order to tackle this challenge. The most significant advancements include efficient implementations that increase the model's context length during training (Dao al., 2022 Liu et al.), and context-extension methods (Chen et al., 2023b Peng et al. (203b) designed to effectively expand the context after training. However, recent studies suggest that long-range processing still an unresolved problem (Li al. 2024a Liu et al. 2024a).
One promising approach in this domain is the development of attention-frre networks with sub- quadratic complexity, which can be trained more efficiently over long sequence data. In recent line of works (Gu et al. 2021 Gu et al. Gupta et al., 2022), the family of state-space layers has been introduced, These layers can seen as theoretically grounded linear RNNs that can be efficiently computed in parallel via convolutions, thanks to closed-form formulation of their linear recurrent rule. recent advancement by Gu Dao (2024) presented Mamba, which builds on top of an expressive variant of SSMs called Selective State-Space Layers (S6). These layers match exceed the performance of Transformers in several domains, such as NLP (Lieber et al. 2024 Zuo et al,
2024 Pióro et al. 2024 Wang al. 2024a). image classification (Zhu et al.) Liu et al. 2024b).
audio processing Shamss et al. 2024) genomic data (Schiff' et al.) 2024). and more.
In this paper, first explore the length- generalization abilities of Mamba and identify that they relatively limited. Although Mamba layers are theoretically capable of capturing global interactions
the layer level, we show, through series of visualizations, analyses, and empirical measures, that

---

Published as a conference paper at ICLR 2025
LongBench
0-4K
Mamba
Passkey PasskeyDeepp Depth [%)
20-
20
15-
15-
0.5
0.75
1K
8K
16K 32K 128K
MPQA GovR
MN TOA LOC Rep8i Avg
MPQA GowR
MN TrOA
LOC.
Avg
DeciMamba
4-8K
8k+
30-
30
Passkey Depth [%]
DeciMamba-Inst-2.88
25-
25
Mamba-Inst-2.88
0.25
20
20-
Length
0.5
0.75
10-
10-
5-
5-
1K
8K 62K
MN
RepB Avg
MFOA TIOA
RapB Avg
Context Length [Tokens]
Figure 1: Improving Mamba Extrapolation with DeciMamba. We present a novel decimation mechanism tailored for Mamba. With method Mamba is able to process sequences that much longer than the ones trained on while enjoying reduced inference costs. (Left) Zero-shot evaluation of an instruction-tuned 2.8B Mamba model on subset of LongBench tasks. We show both LongBench and LongBench,e (three length groups: 0-4K, 4-8K, 8k+). MFQA, GovR, MN, TrQA. LCC, RepB, and Avg stand for MultiFieldQA. GovReport, MultiNews, TriviaQA, LCC, RepoBench-p, and Average; (Right-Top) Passkey Retrieval for Mamba-1300 (Right-Botton) Same for DeciMamba- 130m. All models (130m, 2.8b) were trained on lengths of 2K tokens.
the main barrier is Mamba's implicit bias towards sequence lengths that were seen during training, phenomenon that we call 'limited effective receptive field' (ERF). Next, based on the assumption that long-context data usually sparse, we present DeciMamba, the first context-extension method specifically designed for S6. Our method relies a dynamic data-dependent pooling method that utilizes a hidden filtering mechanism intrinsic to the Mamba layer. We leverage this mechanism to introduce global compression operator, which expands Mamba's 's ERF by discarding unimportant tokens before the S6 layer. In particular, we interpret the norms of the per-token selective recurrent gate (A,) indicators of each token's importance. This metrie allows most impactful tokens, enabling direct application of the SSM to these tokens only. The proposed method (Figure 2) significantly increases the effective context length of Mamba by several orders of magnitude while requiring a smaller computational budget.
identify the top-k
Our main contributions encompass the following three aspects: (i) identify that Mamba has limited nongh--xxtrapolation capabilities via thorough experiments in controlled environments. (ii) Through a series of visualizations, analyses, and empirical measures, recognize that although Mamba can theoretically capture global interactions via the recurrent state, its limited ERF prevents significant length-ex eiathheettppolatinn, (iii) building on this insight, introduce DeciMamba, the first context-extension technique specifically designed for Mamba models. This approach leverages existing filtering mechanism embedded within the S6 layer. As illustrated Fig. D our method effectively enhances Mamba's 's tength-extrapplattion abilities, and is applicable to real-world long-context NLP tasks.
PRELIMINARIES
Long-range models evolve in two main directions: (i) adapting transformers, the most dominant
chitecture today, to be more suitable for such tasks; (ii) developing architectures with sub-quadratic complexity in sequence length such as Hyena (Poli al. 2023), RWKV (Peng et al.] 2023a),
Hawk (De et al. 2024), XLSTM (Beck et al. 2024), and Mamba, the focus of our paper.
Mamba. Given input sequence U (u1, u2 ... uz) E RLxd of length L such that u, Rd Mamba block with channels is built on top of the S6 layer via the following formula:
= a (Linear(UU),
X = Conv (onnea((Ln))r(())),
Y = S6(X),
O Y G
2

---

Published as a conference paper at ICLR 2025
DeciMamba Block
A Strategy
Decimation: Scope & Ratlo
h = expl - + A,Ba St = Cihb
DeciMamba
Mamba Block
Simplified Recurrent View Mamba
ooo
SSM
Cutput
Mamba Block
C2
DeciMamba Block
A Based Decimation
hr
Recurrent Gate
BA C
AA
DecIMemba Block
Conv1D
Input Gate
Mamba Block
1B
Mamba Block
A. controls the impact on future tokens
Figure 2: DeciMamba (Left) Schematic overview; (Middle) By carefully inspecting the recurrent view of Mamba, we revealed the implicit filtering mechanism embedded in the recurrent gate and controlled by A;; (Right) Visualization of DeciMamba model. The grey lines represent the quence length at the input and output of each layer. Layers with a large ERF are decimated.
where G represents the gate branch, is elementwise multiplication, is the SILU activation, Linear and Conv ID are standard linear projection and I-dimensional convolution layers. The S6 layer is based on variant SSM, which can be elaborated by the following recurrent rule:
h4 A + BR V/ Cgh+
(2)
where X F2,, ...., FL) is the input sequence of representative channel, A, E R RNN,, B. E RNx1, and C, RIXN are the system, input, and output discrete time-variant matrices, respectively. S6 conditions the discrete time-varian. matrices based on the input as follows:
A. ())SAX)). B. SBX,, (ScX,)" A, = exp(A).)),
B, B B,
such that A, is the discretization step, Sft represents the softplus function, and Sa, SB, Sc linear projection layers. Ali et al. (2024) demonstrated that S6 layers, similar to attention models, can be interpreted data-controlled linear operators. Specifically, the S6 layer computation can be represented using the following linear operator a, controlled by the input (via Eq. 3):
==
a. = C. [TT
(Ti B,
(4)
V1
B
0 CaBr
0
#1
y/2
CrA2BI
Ø2
(5)
Lr
C CLITEE
O A
C₁
BL
In this formulation, each output is weighted sum of all inputs, where the 'attention weights' of data-driven. We utilize this perspective to further investigate
all inputs zy. i.e.. the {ag,) the effective receptive field of Mamba layers.
Context Extension & Length Extrapolation. Several methods were proposed enhance the ef- fective context length of transformers and improve their extrapolation over longer sequences. Press al.] demonstrated that models built on top of original sinusoidal, rotary (Su et al. 2024). and T5 bias (Raffel et al. 2020) positional encoding have poor length generalization. It proposed to miti- gate this issue by incorporating distance-basee linear biases into the attention matrix for promoting locality KKzzmnçjad et al.) (2024) showed that transformers without positional encoding (NoPE) exhibit better length extrapolation capabilities in downstream tasks. Recently, CoPE (Golovneval al. 2024) utilized context-aware positional encoding and Peng et al. (2023b); Chen al. (2023a) suggested post-training positional interpolation.
A recent direction involves architectural modifications to pre-trained models followed by short fine- tuning. includes LongLora (Chen et al. 2023b), which proposes shifted sparse attention, and Landmark Attention (Mohtashami & Jaggi 2023), which applies attention in chunks and inserts
3

---

Published as a conference paper at ICLR 2025
global unique tokens into the input sequences between those chunks. Our work focuses taking such an approach for Mamba models, rather than transformers. More related work is in Appendix D Furthermore, refer to Appendix D.4, where motivate the study of Long Context LLMs in parallel to Retrieval Augmented Generation (RAG). Lastly, we note that in order to assess the long range performance of the model, in our experiments measure perplexity only on the farthest
labels of the context window, causing the perplexity values to be larger than their typical values (usually, all tokens in the context window are aggregated).
3 EXTRAPOLATION LIMITATIONS OF MAMBA
In this section we explore the mongth-generalizatinn capabilities of Mamba models. Specifically, we assess the limited effective receptive field (ERF) of S6 layers, a phenomenon that leads to poor information propagation when processing sequences that are longer than the ones trained on. For a detailed description of the ERE, we refer the reader to Appendix D.3] Consequently, we develop methods to measure the portion the context utilized by the model, and show that is dictated by the training sequence lengths.
We start by visualizing Mamba's 's hidden attention (Eq. (5 during extrapolation. In Fig. 3 (left,
center) we display the attention matrices of layer 17 of Mamba- 130M while evaluating sequences of length 2K and 16K (the model trained on the Passkey Retrieval task with sequence lengths of 2K). Notice that while the attention is very dense for the 2K sequence (no extrapolation), the attention for the 16K sequence (extrapolation x8) is much more sparse, and its elements vanish as we move towards the lower rows. This exhibits a limited ERF, as information from the first tokens in the sequence does not propagate to the final tokens in the output sequence.
2K Tokens
16K Tokens
101
Sum Value [Log Scale]
10
26
16K
10?
101
2K
109
2K
4k
BK
16k
16K
Context Length
Figure 3: Detecting and Quantifying Limited ERFs (Center, Left) Recordings of Mamba At- tention Matrices with and without extrapolation (Mamba-130m. layer 17, trained on seq. lengths of 2k). Mamba unintentionally learns a limited ERF during training (highlighted by the dashed rect- angle) which disrupts its extrapolation abilities. (Right) Quantifying Mamba's Information Loss by Measuring Eh AІ Divergence. To show Mamba's 's sensitivity to increasing context lengths we measure the first occasion of information loss, as described in Sec. 3 We observe that in the most semantic layers (16 and 17, see Passkey Retrieval in Sec. 5) Σ k=2 Dg diverges exponentially fast,
causing a fast decay in the attention values, leading to limited ERFs like in the center image.
Measuring ERFs via Mamba Mean Distance. To quantify how well Mamba utilizes the context during inference, we introduce a quantitative measure called 'Mamba Mean Distance". This measure is analogous to the receptive field in CNNs and the attention mean distance in transformers described by Dosovitskiy et al. (2020). For a causal transformer model, the attention mean distance for the i-th output token is computed by:
E d(i.))
Σ Ang (1 35E
(6)
where the normalized attention matrix that defines a probability distribution over the distances for various tokens. We tailored this measure to S6 by leveraging the implicit attention representation,
4

---

Published as a conference paper at ICLR 2025
which we normalize using the following function:
N(x1. 1........ ;LL) === [x
As we usually interested in the last tokens, and as we will see shortly, they are the first to suffer from this phenomenon, we find that it is sufficient to compute our metric for the last token only:
E d(L.j) =
[N()) j).
SSL.
JEL
Fig. visualizes this measure by presenting the 'Mamba Mean Distance' for different hidden at- tention matrices, depicted by the horizontal distance between the red line and the main diagonal.
Figure 4: Mamba Mean Distance. Each panel contains an attention matrix along with its cor- responding 'Mamba Mean Distance', depicted by the horizontal distance between the red diagonal line and the main diagonal. The matrices are extracted from pre-trained model of size 2.8B, trained on the Pile dataset (Gao et al. 2020).
Normalized Mambb Mean Distance
Lowe = 1024
70
Lywe 512
Mamba Mean Distance
0.25
Limo
60
128
64
0.20
40-
Irow 1024
0.15
Liao 512
30
Liww 256
0.10
Lrass 128
Irew
200
400
600
800
1000
200
400
600
800
1000
Leval
Leval
Figure 5: Mamba Mean Distance Quantifies Effective Context Utilization. (Lefft Mamba Mean Distance as a function of the context length during inference, for various training context lengths. (Right) Same, but normalized by the training context length.
By introducing an empirical measure for Mamba's ERF, we can explore it more robustly. Fig. portrays the relationship between the input context length and the Mamba Mean Distance. Here, Mamba-based language models with 80M parameters are trained across various context lengths for next-token prediction on the WikiText-103 benchmark (Merity al. al.] 2016). We average the Mamba Mean Distance over all channels and layers using 100 test examples. More details on Mamba Mean Distance are in Appendix D.S Surprisingly, in the left panel of Fig. 5 we see that the Mamba Mean Distance increases with context length. Although it does grow, it is hard to tell whether the growth is proportional the growth context length. Therfore, in the right panel, we normalize the Mamba Mean Distance by the sequence length during inference, revealing the exact proportion of context utilized in practice. is now evident that the context utilization drops dramatically as the evaluated context length increases.
To identify the internal mechanism that leads to this behavior, we return to the analytic expression of Mamba's hidden attention defined in Eq. 5 We observe that the product of transition matrices
IT Âg includes more elements the sequence length I. increases, and be written for the
j' j'th element in the last
II k=j+1
A
II k=j+1
exp(AAx) exp(A Σ A+) k=j+1
(9)
5

---

Published as a conference paper at ICLR 2025
We note that this product always converges since A, 2 0 and A < by design. Yet, if it converges too fast, the attention element will collapse undesirably. To measure the attention ele- ments' collapse we focus attention element æ due to the following reasons: (i)
usually interested in the information at the end of the sequence; (ii) this element is the first in the row to collapse, capturing the first occasion of information loss because Ek=2 As 2 Ekej+1 A for all L-1 j 2 1, and A is always negative, so the power of its exponent will always be the smallest. As can be seen in Eq. 9 the only input-length dependent value is Σ Ag. Therefore, in Fig. (right), we compute it for varying input lengths L during inference (per layer, averaged over the channels dimension). As can be seen in the figure, the values of Σk= 4 start small and increase exponentially fast in layers 16 and 17, which are the most global layers (Fig. [2]. This aligns with the observations in Fig. 3 (left, center) since these layers require a large ERF, the fast collapse of the attention elements restricts global information propagation, leading distorted performance. In the other layers, the values of Σk=2 As start large and grow linearly. Since these layers perform local processing (diagonal attention matrices), the fast decay is desired.
To conclude, the limited length- -eneralleaiiioonnnnnn abilities arise due to decay rates and As sums (Eq. ( that are sufficient for the training data of length Lerain. but not flexible enough to provide effective length-extrapollatoon for lengths Leval > Ltrain. As the transition matrix product con- verges too fast, observe collapse of the attention matrix, specifically in regions that convey information from the beginning of the sequence to its end.
4 METHOD: DeciMamba
In Sec. 3 we identified that for a pre-trained model, the ERF of Mamba is dictated by the context size used during training Livain. This creates blind spot for sequences that exceed Ltrains as dependencies originating in the far input segments are not captured by the model, resulting in poor length generalization. To solve this problem we propose embedding filtering mechanism within the pre-trained Mamba layers, with no need to re-train the model. The mechanism' 's core task is to reduce the amount of tokens that the S6 layer processes, and it does so by discarding tokens of fewer importance. The method is described in Fig. 2 and Alg. [ It encapsulates three aspects: (i) Decimation Strategy, (ii) Decimation Rate, and (ini) Decimation Scope. We refer the reader to Appendix C where we further discuss each hyperparameter' 's role and selection strategy.
Decimation Strategy (The Role of A.) To pool the most relevant tokens, we must assign each token a relative importance score. We select the A parameter as proxy for this score and motivate our choice by explaining its operation. First, we apply the SSM parametrization to Eq.
hr Adhe-1 + Br,, = Ado + A,
(10)
Notice that when A, + 0 the layer discards the input token and preserves the previous hidden state. When A, > 0 the hidden state 'attends' the input token (the attendance is proportional to A,) and adds it to an attenuated version of the previous hidden state (this attenuation is also proportional to At). We note that this is always the case because A, 2 and A i] 0 by design. Hence, A can be interpreted as the controller of the recurrent gate, determining which tokens should impact future tokens. The final importance score for each token is the mean value over all channels of its respective A parameter. Note that this selection induces small interference in Mamba's 's operation, as Mamba already attempts to ignore input tokens with small Á,
Decimation Ratio For each decimating layer we propose to keep the Top-P, tokens with the largest importance scores, where P, decreases gradually as we go deeper into the network according to the following formula:
Ps = Linse 8*, B E (0, 1), Liase E N,
(11)
where B, Loose are saperpparameters representing a decay factor that controls the decimation rate (as we progress in depth) and the maximal length of the sequence after the first decimating layer. Sec. E.2 further discusses the selected pooling strategy.
Decimation Scope (Layer Selection) We turn to describe the decimation scope, which defines which layers should use the embedded decimation mechanism. We begin by explaining the guiding principles behind our method, followed by the criteria that reflects those principles. Our goal is to

---

Published as a conference paper at ICLR 2025
expand the ERF of pre-trained Mamba model. Traditionally, DL models focus various scales of token interactions at different layers. Therefore, a natural design choice is to decimate layers that already focus long-range dependencies. This approach can potentially increase their ability to learn global dependencies without negatively impacting the layers that are associated with short-term features. As criterion for measuring the scale embedded within each layer, we measure the ERF of Mamba based on the Mamba Mean Distance defined in Eq. 8 We thus select the layers with the highest Mamba Mean Distance, with the number of selected layers being hyper-parameter. Note that during inference we perform decimation only in the pre-fill fill phase. Section E.4 demonstrates that its use in the decoding phase leads negligible improvement and thus not required.
Inference Efficiency Mamba can be computed in parallel with a prefix-saan or step-by-step like RNNs. For efficient inference, the context is processed in parallel mode (pre-fill), and then the response is generated via auto-regressive decoding in recurrent mode. DeciMamba performs pooling only during the pre-fill phase, leveraging the parallel mode to compute the global pooling operation efficiently. During decoding, Mamba and DeciMamba have identical forms (no pooling).
Complexity In the parallel view of Mamba, the time complexity for processing a sequence of length I., is (L log L). In DeciMamba, we first compute the delta-based importance score for each token, identify the top-P elements, and then apply S6 to those elements. This results an overall complexity of OLL + P Plog P). We ignore the complexity of selecting the top-P elements (which is lo P)) since this computation does not depend on the state size or the number of channels H. making it negligible compared to the other operations in both Mamba and DeciMamba.
Algorithm Decimated SSM
Require: x: (B Batch size, Tokens, D Channels)
I: A: (D, N) & Parameter
2: B: (B, L, N) + Sp(x)
3: C: (B, L, N) <- Sc(x)
4: A: LB,L, D) (araarmmeer Sa(x))
Ap < Top--((A)).indiis
6: A, B,B.C,& < A[Ar], 4[Ar] B[Ap), ([Ar.aAA))
7: A, (B, L, D, N) + AiscrrtizeeA,A.
8: return y: (B, L, D) + SS, B, Cx))
5 EXPERIMENTS
We evaluate DeciMamba as context-extension approach across multiple tasks. First, we demon- strate the aangrange understanding and retrieval capabilities of our method in the NLP domain. Next, we assess the retrieval capabilities of our method using the Passkey -etrival task and exam- ine its language modeling capabilities over PG-19. Finally, we analyze and measure DeciMambu's inference efficiency and provide ablation studies to support our claims.
LongBench This benchmark contains variety of zero-shot long context tasks, such as: Sin- gle/Multi Document QA, Summarization, Coding, etc. We evaluate an instruction-tuneet Mamba- 2.8b (model: "xiuyu/maamba-2.58 (...aamaa-..bbzepphe.e 2.eb-ry from huggingface) with and without DeciMamba (Zero- Shot) and report the results in Table [ 0-4k, 4-8k, 8k+ are the scores for each context length group LongBench-FF, LB is the LongBench score, and 'N/A' means that the task does not ex- ist in LongBench-E. In almost all tasks, DeciMamba improves the performance of the model for all context lengths. Most noticeable are the improvements at 4-8k, where Mamba starts suffering from ERFs, while DeciMamba extends its performance significantly. E.g. in TriviaQA we improve the LB result from 3.93 to 12.61, performance gain of 220%, and improve each context length group by 24% (0-4k), 150% (4-8k), and 152% (8k+) respectively. Despite the improvement, when DeciMamba is applied to zero-shot scenarios it does not completely prevent degradation when the context length increases, motivating future research of this behavior. See Sec. A.3 for more details.
Document Retrieval In this task the model receives a query and Ndoce randomly assorted doc- uments, and its objective is return the id of the matching document. Our data is sampled from SQuAD v2 (Rajpurkar et al., 2018). We train both Mamba- 130m and DeciMamba- 130m to retrieve out of Ndocs 11 documents (about 2k tokens), and evaluate their performance with
7

---

Published as a conference paper at ICLR 2025
Table 1: LongBench. The results are for an instruction-tunee Mamba-2.8b model with and without DeciMamba (Zero-Shot). Avg Len = Average Length in words.
Avg
Mamba
DeciMamba
Type (Metric)
Benchmark
Len
0-4k
4-8k
8k++
LB
0-4k
4-8k
8k+
LB
MultiDoc-QA (FI)
2wikimqa
4887
7.04
1.58
0.6
3.92
11.54
6.2
3.08
9.06
MultiDoe-QA (F1)
Hoipotqa
9151
4.7
1.18
0.18
1.45
8.29
3.75
3.19
4.46
MultiDoc-QA (F1)
Musique
11214
N/A
N/A
N/A
0.85
N/A
N/A
N/A
1.73
SingleDoc-QR (F1)
Narrative QA
18409
N/A
N/A
N/A
0.87
N/A
N/A
N/A
1.74
SingleDoc-O5 QA (F1)
Qasper
3619
6.91
4.4
1.36
5.97
8.75
9.2
2.41
8.91
SingleDoc-QA (FI)
Multifield QA
4559
18.26
5.71
2,45
11.16
26.05
13.67
4.63
18.58
Summarizing (Rouge-L)
GovReport
8734
24.76
10,9
5.2
9.84
27.4
19.38
7.81
14.86
Summarizing (Rouge-I.)
QMSum
10614
N/A
N/A
N/A
8.18
N/A
N/A
N/A
7.08
Summarizing (Rouge-L)
MultiNews
2113
24.58
9.79
5.34
23.15
25.45
17.2
4.43
24.58
Few-Shot (F1)
TriviaQA
8209
10.38
5.59
4.11
3.93
12.9
14.02
10.36
12.61
Few-Shot (Rouge-L)
SAMSum
6258
9.58
7.07
7.76
8.56
9.17
6.76
6.88
7.34
Few-Shot (Accuracy)
TREC
5177
0.0
0.0
0.0
0.5
1.0
0.0
0.0
0.5
Code (Edit Sim)
LCC
1235
8.12
5.61
4.17
8.13
9.4
14.25
7.63
8.67
Code (Edit Sim)
RepoBench-p
4206
7.52
5.74
4.63
7.15
10.08
10.49
6.86
10.96
Synthetic (Accuracy)
Passage Count
11141
2.0
0.0
0.0
3.0
0.0
0.5
Synthetic (Accuracy)
Passage Ret. en 9289
0.0
0.0
0.0
0.0
9.0
1.0
0.0
1.5
Ndocs E [11, 160) (2k to 32k tokens). The results are presented in Table (Retrieval). While the performance Mamba and DeciMamba similar when evaluated on sequence lengths close to those used during training, DeciMamba gains clear advantage as the number of documents exceeds this length, achieving improved length generalization. See Appendix A.2 for more details.
Multi- Multi-Documentt Document QA. Next, we stay in setting similar to Sec. (Document Retrieval), yet increase the level of difficulty by asking the model to answer the query in free text (instead of retrieving the id of the most relevant document). We train Mamba- 1330m and DeciMamba- 130m the same dataset and present the Fl score between the generated response and the ground truth answers in Tab. 2 As in the previous task, when the amount of documents is close the amount trained on, the models perform quite similarly. When the document amount increases we can see that DeciManba develops clear advantage. Note that the performance of both Mamba and DeciMamba is relatively modest compared to prior work, due to the use of smaller models.
Table 2: Multi-Document Retrieval and QA. The table shows the score of Mamba and DeciMamba (+Deci) 130m models as we increase the amount of documents during evaluation. For retrieval we show the accuracy and for free-form QA we show the F1 score between the ground truth annotations and the predictions. Both models were trained using 11 documents.
# Docs:
10
20
40
60
80
100
120
140
160
Mamba
68.3
64.7
45.3
24.7
5.7
1.0
0.3
Retrieval (Acc)
+Deci
77.3 68.7
65.3
49.7
37.0
26.3
16.7
5.3
Mamba
25.3 25.3 22.8 23.6 25.8
19.3
7.5
4.1
0.6
QA (F1)
+Deci
26.2 24.4 25.3 22.74
19.1
18.1
12.9
10.0 5.0
Passkey Retrieval. We fine-tune Mamba-130M and DeciMamba- 330M to retrieve a random 5-digit code hidden in a random location within contiguous 2K token sample from Wiki-Text (Merity) et al., 2016). To test the models' extrapolation abilities, during inference we increase the sequence lengths exponentially from 1K to 128K and record their performance for a variety of passkey loca- tions within the context. Full implementation details are in Appendix A. As can be seen in Fig. [] (right), DeciMamba- 130m significantly increases the extrapolation abilities of Mamba-130m from 16K to 128K, when trained on sequence lengths of 2k tokens only. Interestingly, we identify that the A, values in Layer 16 of Mamba- 130m capture the exact location of the passkey, as can be seen in Fig. I0 Using this finding, we display the distortion caused the embedded sequence due the ERF in Fig. II Tokens at locations < 10k have meaningful A, values, yet tokens at locations > 10k have noisy A, 's, possibly due the poor information flow caused by the limited ERF.
8

---

Published as a conference paper at ICLR 2025
Language Modeling Following Chen et al. (2023b); Mehta al. (2022) evaluate method on long-range language modeling using the PG- dataset in both zero-shot and fine-tuned regimes.
Zero-Shot. We test method the test set of PG-19 using the larger Mamba models (1.4b, 2.8b) in Fig. (right). We observe that the selectivity scores learned during pre- are-training are quite effective method can
and extend the context naturally without any training for both models. Furthermore, maintain the low perplexity achieved in the short-context region. In the ablations section we further discuss our proposed decimation mechanism and show its benefit w.r. other alternatives.
Fine-Tuning. We display DeciMamba's s performance on PG-19 perplexity in Fig. 6 (left). We train both Mamba-130% and DeciMamba- 130M with a sequence length of 2K and test their extrapolation abilities. The full training details can be found in Sec. A.4 in the appendix. While Mamba can ex- trapolate to context lengths that are at most x5 times longer than the training sequences, DeciManba can extrapolate to sequences that are about x20 times longer, without any additional computational resources. Furthermore, we plot a lower bound by calculating the perplexity of a Mamba- 130M model that was trained on the same context length it evaluates a different model). We observe that DeciMamba is close to the lower bound and diverges from quite slowly, while utilizing significantly less computational resources. Nevertheless, for train quences longer than 30K the 'lower bound' models reach an Out Of Memory (OOM) error (Nvidia RTX-A6000, 48GB of RAM), which is caused by Mamba's O(L) training complexity. This further emphasizes the importance good extrapolation abilities when scaling to longer sequences.
(each point on the green curve is
Perplexity Over PG-19 Test Set
60
Mamba-2.88
80
DeciMaaba-2..8
Mamba-1.4B
40
Perplexity
60
DeciMamba-1
Perplexity
30
40
00M
20
Mamba
10
DeciMamba
Length Training
10K 20K 30K 40K 50K 60K 70K 80K Context Length [Tokens]
10K
20K
30K
40K
50K
60K
70K
BOK
Context Length [Tokens]
Figure 6: Perplexity Over PG-19 Test Set. (Left) The dashed purple line shows the train sequence Each point on the green curve shows a different
length for Mamba- 100m and DeciMamba-1300 Mamba- 1 130m model trained on the respective context length (Ltrain Levat); for Ltroin > 30K an Out Of Memory (OOM) error occurs. (Right) Zero-Shot perplexity comparison.
Comparison With Transformers. We evaluate equivalent Transformer models from the Pythia suite (Biderman et al. 2023) the experiments described above. We find that vanilla Transformers have inferior length generalization abilities compared to Mamba. Full results are in Appendix B.
Inference Efficiency. We benchmark both DeciMamba and Mamba with a Nvidia RTX A6000 GPU and compare their inference speeds for the Passkey Retrieval task. From Table that DeciMamba is twice faster than the baseline Mamba model, which is expected. This is because the decimation is performed at layer 12 (out of 24 layers in total), hence only half of the blocks process the whole context length The other half processes the decimated sequence (of length P. << L), resulting in negligible computation time. Table D (Appendix) shows that per block, DeciMamba does not induce an additional computational overhead when compared to Mamba (first that the computation time becomes negligible for blocks that
3 it is evident
two rows). In addition, placed after the first DeciMamba block (rows three and four), no matter the input context length.
Ablation study The ablation studies can be found in Appendix E Briefly, these studies examine three key aspects: layer selection, pooling strategy, and decimation mechanism. For layer selection, it seems that the choice of the first decimating layer significantly affects performance, with layer 12

---

Published as a conference paper at ICLR 2025
Passkey Retrieval Success Rate
1.0
DeciMamba
Mamba-1.4B
Mamba
+ Max Norm Pool
0.8
+ Random Pool
60
+ Ae Pool (Ours)
Perplexte perplxxity
0.6
40
0.4
0.2
0.0
0 10K 20K 30K 40K 50K 60K Context Length [Tokens]
6
10 12 14 16 18
20
Min Decimating Layer
Figure 7: Ablation Studies. (Left) Model sensitivity to decimation layer selection. We show the score achieved by DeciMamba- 110 130M in the Passkey Retrieval task when trained on sequences of length 2K. Each point on the graph is model's score, while decimating from layer 'Min Decimat- ing Layer' to layer 20. (Right) A Zero-Shot perplexity measurement comparing between different pooling mechanisms: ours (based on A,), random and max-norm pooling.
Table 3: Inference Efficieny. Inference time (seconds) of the entire model, benchmarked on Nvidia RTX A6000 GPU. DeciMamba is about twice as fast when compared to Mamba.
Context Length
Model
8K
16K
32K
64K
128K
256K
512K
Mamba-1300 DeciMamba- 130m
0.17
0.3
0.69
1.27
2.12
4.03
8.17
0.14
0.19
0.31
0.59
1.08
2.01
4.22
yielding the best results in Passkey Retrieval task. The pooling strategy study compares Top-Ps approach with a Top-K% approach, demonstrating that the former is more effective in extending the model's extrapolation abilities. The decimation mechanism study compares various pooling methods, where bb-aee pooling performs best, though max-norm pooling also shows promise. In Appendix [E.4 we justify our design choice of pooling only during the pre-fill phase of inference. We demonstrate this by ablating an additional pooling mechanism during the decoding phase.
CONCLUSIONS
In this paper, we explore the the characterization of the ERF of S6. This characterization reveals that the ERF of Mamba is significantly constrained by the context length during training, which is counterintuitive given that Mamba theoretically has unbounded receptive field due to its recurrent selective memory. Based on these insights, we develop DeciMamba, a unique data-dependent compression operator that is built on two key insights: (i) There exists hidden filtering mechanism within the Mamba layer, manifested by the selective time-steps A,, which can be interpreted as the controller of the recurrent gate. (ii) Long-range and short-range interactions between tokens are captured by different layers in the model, which can be identified using our Mamba Mean Distance metric.
abilitics of S6 layers. Our first contribution is
Looking ahead, we plan to explore different transformer context-extension methods, including enagth-extrapllatiin PE (Press et al.) Golovneva et al., 2024), hierarchical models, and architec- tural improvements (Sun et al. 2022). Furthermore, while the techniques described in this paper are designed for Mamba models, detailed discussion extending our proposed methods to other models such as Griffin can be found in Appendix [F] Finally, we will further use our Mamba Mean Distance Metric to investigate the mechanism which leads to limited ERFs in Mamba.
7 LIMITATIONS
Similar to other context-extension methods, our model modifies pretrained Mamba model but does not propose an improved Mamba architecture to address the underlying issue. Despite empirical evidence demonstrating the effectiveness of our decimation-based method in capturing long-range interactions and its efficiency, the approach can be suboptimal. For example, pooling and com- pression methods may miss critical information in challenging scenarios. Therefore, designing an improved Mamba variant with enhanced length-geus seathh-ennraalzation abilities that can effectively capture global interactions within single layer (without pooling the sequence) is important next step.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
This work was supported by grant from the Tel Aviv University Center for Al and Data Science (TAD), and was partially supported by the KLA foundation. We would like to thank Yael Vinker and Maor Ivgi for fruitful discussions and valuable input which helped improve this work.
8 KEPRODUCEABILITY STATEMENT
First, we provide the source code used for the key experiments. Second, in appendix A we provide full configurations
experiments including instructions on how to train and evaluate the models. We also explicitly state all used datasets, models and hardware, and describe the data processing pipeline and exact metrics used in each experiment.
ETHICS STATEMENT
This work analyzes and improves long-context understanding, which is crucial when deploying LLMs real-world systems. This improvement is anticipated to have a positive impact on the use of LLMs in society. However, acknowledge that LLMs can propagate biases. We emphasize the necessity of further research into these biases before our work can be applied reliably beyond the research environment.
REFERENCES
Ameen Ali, Itamar Zimerman, and Lior Wolf. The hidden attention of mamba models, 2024.
Maximilian Beck, Korbinian Pöppel, Markus Spanring, Andreas Auer, Oleksandra Prudnikova, Michael Kopp, Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. xlstm: Extended long short-term memory. arXiv preprint 2rxiv:2405.04517. 2024.
Stella Biderman, Hailey Schoelkopf, Quentin Anthony, Herbie Bradley, Kyle O'Brien, Eric Hal- lahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, Aviya Skowron, Lintang Sutawika, and Oskar van der Wal. Pythia: A suite for analyzing large language models across training and scaling, 2023. URL http:: www.v..xivvv..///ii//aa. org/ abs 2304 0330..03373 01373
Shouyuan Chen, Sherman Wong, Liangjian Chen, and Yuandong Tian. Extending context window of large language models via positional interpolation. arXiv preprint arXiv:2306. 15595, 2023a.
Yukang Chen, Shengju Qian, Haotian Tang, Xin Lai, Zhijian Liu, Song Han, and JiayaJia Jia. Longlora: Efficient fine-tuning of long-context large language models. arXiv preprint arXiv:2309 12307, 2023b.
Alexis Chevalier, Alexander Wettig, Anirudh Ajith, and Dangi Chen. Adapting language models to compress contexts, 2023. URL https:/:: arxiv. org/ ab////////// 2305 14788
Ta-Chung Chi, Ting-Han Fan, Alexander Rudnicky, and Peter Ramadge. Dissecting transformer length extrapolation via the lens of receptive field analysis. arXiv preprint arXiv:2212. 10356, 2022.
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. arXiv preprint arXiv:2009. 14794, 2020.
Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. arXiv preprint arXivv2307.0866/1 2023.
Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory- efficient exact attention with io-awareness. Advances Neural Information Processing Systems, 15:4444-116359 2022.
Soham De. Samuel L Smith, Anushan Fernando, Aleksandar Botev. George Cristian-Murard Al- bert Gu, Ruba Haroun, Leonard Berrada, Yutian Chen, Srivatsan Srinivasan, al. Griffin: Mix- ing gated linear recurrences with local attention for efficient language models. arXiv preprint arXiv:2402. 19427, 2024.
11

---

Published as a conference paper at ICLR 2025
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010. /929, 2020.
Weizhi Fei, Xueyan Niu, Pingyi Zhou, Lu Hou, Bo Bai, Lei Deng, and Wei Han. Extending context window of large language models via semantic compression, 2023. URL https://arxxi
org/abs/ 2312. 09571
Quentin Fournier, Gaétan Marceau Caron, and Daniel Aloise. practical survey on faster and lighter transformers. ACM Computing Surveys, 55(14s):1-40) 2023.
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang. Horace He, Anish Thite, Noa Nabeshima, al. The pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2101.00027. 2020.
Tao Ge, Jing Hu, Lei Wang, Xun Wang, Si-Qing Chen, and Furu Wei. In-context autoencoder for context compression in a large language model, 2024. URL https::/axxiv..or///oooo
2307.0694
Olga Golovneva, Tianlu Wang, Jason Weston, and Sainbayar Sukhbaatar. Contextual position en- coding: Learning to count what's important. arXiv preprint arXiv:2405.18719 2024.
Albert Gu and Tri Dao. Mamba: Linear-ting sequence modeling with selective state spaces. In First
Conference on Language Modeling. 2024. URL https:://operrevee... forttt/fommmeeem
LEYskwlVY2
Albert Gu, Karan Goel, and Christopher Re. Efficiently modeling long sequences with structured state spaces, In International Conference on Learning Representations.
Albert Gu, Isys Johnson, Karan Goel, Khaled Saab, Tri Dao, Atri Rudra, and Christopher Ré. Com- bining recurrent, convolutional, and continuous-time models with linear state space layers. Ad- vances neural information processing systems, 34:572-585 2021.
Ankit Gupta, Albert Gu, and Jonathan Berant. Diagonal state spaces are effective as structured state spaces. Advances in Neural Information Processing Systems, 35:22982 25222992-22994, 2022.
Farnoush Rezaei Jafari, Grégoire Montavon, Klaus-Robert Müller, and Oliver Eberle. Mambalrp: Explaining selective state space sequence models. arXiv preprint arXiv:2406.0755 07592, 2024.
Huiqiang Jiang, Qianhui Wu, Chin-Yew Lin, Yuqing Yang, and Lili Qiu. Limlingua: Compressing prompts for accelerated inference of large language models, 2023. URL https://arxiv.
org/aba/ 23100 05736
Huiqiang Jiang, Yucheng Li, Chengruidong Zhang, Qianhui Wu, Xufang Luo, Surin Ahn, Zhenhua Han, Amir H. Abdi, Dongsheng Li, Chin-Yew Lin, Yuqing Yang, and Lili Qiu. Minference 1.0: Accelerating pre-filling for long-context Ilms via dynamic sparse attention, 2024a. URL
https::/aaxxt org/ abs / 2407. 02490
Huiqiang Jiang, Qianhui Wu, Xufang Luo, Dongsheng Li, Chin-Yew Lin, Yuqing Yang, and Lili Qiu. Longllmlingua: Accelerating and enhancing 1lms in long context scenarios via prompt com- pression, 2024b. URL (tthttp::/// arxiv. org/ abs abs/2310. 06839
Amirhossein Kazemnejad, Inkit Padhi, Karthikeyan Natesan Ramamurthy, Payel Das, and Siva Reddy. The impact of positional encoding on length generalization in transformers. Advances in Neural Information Processing Systems, 36, 2024.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization, 2017.
Yuri Kuratov, Aydar Bulatov, Petr Anokhin, Ivan Rodkin, Dmitry Sorokin, Artyom Sorokin, and Mikhail Burtsev. Babilong: Testing the limits of lms with long context reasoning-in aaassonnn--in--hanee -n--ssysaca-a 2024.
12

---

Published as a conference paper at ICLR 2025
Yann LeCun, Léon Bottou, Yoshua Bengio. and Patrick Haffner. Gradient- learning applied to document recognition. Proceedings of the IEEE, 86(11):2278- 2324, 1998.
Tianle Li, Ge Zhang, Quy Duc Do, Xiang Yue, and Wenhu Chen. Long-context Ilms struggle with long in-context learning. arXiv preprint arXiv:2404. 2024a.
Yucheng Li, Bo Dong. Chenghua Lin, and Frank Guerin. Compressing context to enhance inference
efficiency of large language models, 2023. URL http::tp///// http//rxivvv.o./ www.xxv..og/a///222/0 /ltg///bb//3111 2ttp//331..0001111..
Zhuowan Li, Cheng Li, Mingyang Zhang, Qiaozhu Mei, and Michael Bendersky. Retrieval aug- mented generation long-contess llms? comprehensive study and hybrid approach, 2024b.
URL https::/aaxii..ott org/ /etps//407..8888
Opher Lieber, Barak Lenz, Hofit Bata, Gal Cohen, Jhonathan Osin, Itay Dalmedigos, Erez Safahi, Shaked Meirom, Yonatan Belinkov, Shai Shalev-Shwartz, al. Jamba: A hybrid transformer- mamba language model. arXiv preprint arXiv:2403.19887, 2024.
Hao Liu, Matei Zaharia, and Pieter Abbeel. Ringattention with blockwise transformers for near- infinite context. In The Twelfth International Conference Learning Representations.
Nelson F Liu, Kevin Lin, John Hewitt, Ashwin Paranjape, Michele Bevilacqua, Fabio Petroni, and Percy Liang. Lost in the middle: How language models use long contexts. Transactions of the Association for Computational Linguistics, 12:157-173, 2024a.
Yue Liu, Yunjie Tian, Yuzhong Zhao, Hongtian Yu, Lingxi Xie, Yaowei Wang, Qixiang Ye, and Yunfan Liu. Vmamba: Visual state space model. arXiv preprint arXiv:2401. 2024b.
Yuhan Liu, Hanchen Li, Yihua Cheng, Siddhant Ray, Yuyang Huang, Qizheng Zhang, Kuntai Du, Jiayi Yao, Shan Lu, Ganesh Ananthanarayanan, Michael Maire, Henry Hoffmann, Ari Holtzman, and Junchen Jiang. Cachegen: Kv cache compression and streaming for fast large language model serving, 2024c. htttpss:/aaxiii..itttrtt/
2310.1000 07240
Harsh Mehta, Ankit Gupta, Ashok Cutkosky, and Behnam Neyshabur. Long range language model- ing via gated state spaces. arXiv preprint arXiv:2206. 13947, 2022.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models, 2016.
Amirkeivan Mohtashami and Martin Jaggi. Landmark attention: Random- access infinite context length for transformers. arXiv preprint arXiv:2305. 2023.
Jesse Mu, Xiang Lisa Li, and Noah Goodman. Learning to compress prompts with gist tokens, 2024.
URL
org/ abs/ 204.. 68467
Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Stella Biderman, Huanqi Cao, Xin Cheng, Michael Chung, Mattco Grella, et al. Rwkv: Reinventing mns for the transformer era. arXiv preprint arXiv:2305. 13048, 2023a.
Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole. Yam:: Efficient context window extension of large language models. arXiv preprint arXiv:2309.0000 2023b.
Maciej Pióro, Kamil Ciebiera, Krystian Król, Jan Ludziejewski, and Sebastian Jaszczur. Moe- mamba: Efficient selective state space models with mixture of experts. arXiv preprint arXiv:2401.04081 2024.
Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel Fu, Tri Dao, Stephen Baccus, Yoshua Bengio, Stefano Ermon, and Christopher Ré. Hyena hierarchy: Towards larger convolutional language models. In International Conference on Machine Learning, pp. 28043-28077 PMLR, 2023.
Ofir Press, Noah Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. In International Conference on Learning Representations.
Zhen Qin, Songlin Yang, and Yiran Zhong. Hierarchically gated recurrent neural network for se- quence modeling. Advances in Neural Information Processing Systems, 36, 2024.
13

---

Published as a conference paper at ICLR 2025
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of machine learning research, 21(140):1-67, 2020.
Maithra Raghu, Thomas Unterthiner, Simon Kornblith, Chiyuan Zhang, and Alexey Dosovitskiy. Do vision transformers see like convolutional neural networks? Advances in neural information processing systems, 34:12116-122 12128, 2021.
Pranav Rajpurkar, Robin Jia, and Percy Liang. Know what you don't 't know: Unanswerable questions for squad, 2018.
Yair Schiff, Chia- Kao, Aaron Gokaslan, Tri Dao, Albert Gu, and Volodymyr Kuleshov. Caduceus: Bi-directional equivariant long-range dna sequence modeling. arXiv preprint arXiv:2403.03234, 2024.
Siavash Shams, Sukru Samet Dindar, Xilin Jiang, and Nima Mesgarani. Ssamba: Self-supervised audio representation learning with mamha state space model. arXiv preprint arXiv:2405.1 1831,, 2024.
Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: En- hanced transformer with rotary position embedding. Neurocomputing. 568: 127063, 2024.
Yutao Sun, Li Dong, Barun Patra, Shuming Ma. Shaohan Huang, Alon Benhaim, Vishrav Chaudhary, Xia Song, and Furu Wei. A length-extrapplltable transformer. arXiv preprint arXiv:2212.10554 2022.
Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey. ACM Computing Surveys, 55(6))1-28, 2022.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural informa- tion processing systems, 30, 2017.
Jesse Vig and Yonatan Belinkov. Analyzing the structure of attention in transformer language model. arXiv preprint urXiv:1906.000 04284, 2019.
Junxiong Wang, Tushaar Gangavarapu, Jing Nathan Yan, and Alexander M Rush. Mambabyte: Token-free selective state space model. In First Conference Language Modeling, 2024a. URL net. / forum? 1d=x1 xNsukssb
htttss://cpenreviit
Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint 04760 2020.
Xiangfeng Wang. Zaiyi Chen, Zheyong Xie, Tong Xu, Yongyi He, and Enhong Chen. In-context former: Lightning- fast compressing context for large language model, 2024b. URL https:
//arxiv.org/abs/22440 13618
David Wingate, Mohammad Shoeybi, and Taylor Sorensen. Prompt compression and contrastive
conditioning for controllability and toxicity reduction in language models, 2022. URL https:
//arxiv. wttxxvv.org//aa//220..0. 03162
Peng Xu, Wei Ping, Xianchao Wu, Lawrence McAfee, Chen Zhu, Zihan Liu, Sandeep Subramanian, Evelina Bakhturina, Mohammad Shoeybi, and Bryan Catanzaro. Retrieval meets long context large language models, 2024. URL
03025
Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim. Gated linear attention transformers with hardware-efficiet training. arXiv preprint 2xXiv:23/2.06633: 2023.
Lianghui 7hu, Bencheng Liao, Qian Zhang, Xinlong Wang, Wenyu Liu, and Xinggang Wang. Vision mamba: Efficient visual representation lcarning with bidirectional state space model. In Forty-first International Conference on Machine Learning.
Itamar Zimerman, Ameen Ali, and Lior Wolf. A unified implicit attention formulation for gated- linear recurrent sequence models. arXiv preprint arXiv:2405. 16504, 2024.
14

---

Published as a conference paper at ICLR 2025
Jingwei Zuo, Maksim Velikanov, Dhia Eddine Rhaiem, Ilyas Chahed, Younes Belkada, Guillaume Kunsch, and Hakim Hacid. Falcon mamba: The first competitive attention-fret 7b language model. arXiv preprint arXhv:2410.05355, 2024.
15

---

Published as a conference paper at ICLR 2025
A
EXPERIMENTAL DETAILS
All model checkpoints are taken from the Hugging Face Model Hut
state spaces/namba 130m
state spaces/namba 37
state-ssacces//ambbee
state-ssaaces/nambbee 4b
state-spaces/nambes 8b
iuyul/mamba-2.1 2.---2...----p- zephyr
Our code is based on the official Mamba implementation!
2
A.1 PASSKEY RETRIEVAL
Each model is trained for epochs with learning of le-4, gradient clipping of 1, batch size of 32 (used batch accumulation) and A optimizer (Kingma & Ba 2017) with weight decay of
0.1. In each epoch the models train over 6144 sequences of length 2K. For DeciMamba- 130M we use L base 2K, 3 = 0.5, decimating layers = [13.. 21, min seq len = 20. Our code is built over an existing version of BABILong (Kuratov al. 2024).
A.2 MULTI-DOCUMENT RETRIEVAL. AND QA
We train each model with data from SQuAD v2 (Rajpurkar et al. 2018), which provides amples in the form of (Query, Document, Answer). Our training samples have the following form: Ndvcs x Document>; <Answer>, where <Document> can be either the golden doc- ument (which holds the answer to the query) or one of Ndocs - randomly sampled documents. <Answer> holds the id of the golden document. In our setting Ndocs 11, the order of the doc- uments is random, and the query and respective document id are appended to the beginning of each document. During Evaluation we use the same setting but vary the value of Ndocs. between 11 and 160. We note that the length of an average document in SQuAD is bit smaller than 200 tokens, so our average training sample has about 2,200 tokens, and the evaluation samples vary between 2,200 tokens to 32,000 tokens. We train for two epochs (1500) steps in each), use a learning rate of le-4, gradient clipping of I, batch size 64 (used batch accumulation), and Adam W optimizer with weight decay of 0.1. We found that the optimal decimation parameters are decimation layer 12, Lbase=2000 during training and =000 during evaluation. We intentionally decreased Lbase during training so the model could experience decimation during the training period (Ltrain bit higher than Linse,, because otherwise the training of DeciMamba and Mamba would have been identical. The displayed results the average performance 3 different training seeds.
A.3 LONGBENCH
We use the 'iuyul/mamba: 2. 1---8--zze "eb-zephyr' instruction-tuned model from huggingface. deci-
mation layer= 28, Linse varies between different benchmarks, between 2000 2800.
A.4 PG-19 PERPLEXITY
We train each model total of 100M tokens with learning rate of le-4, gradient clipping of 1, batch size of 250 (used batch accumulation) and AdamW optimizer with weight decay of 0.1. During training we sample a single window from each example and train on it (For the extrapolating models the window length is 2K, for the lower bound models the window length is equal to the context length trained on). During evaluation, for each example we evaluate 10 windows with maximal constant stride. We evaluate only the last 100 labels in each window, which represent the extrapolation abilities of the model at sequence lengths the range [ctxlen 100, ctx cen_,
16

---

Published as a conference paper at ICLR 2025
providing approximation to the model's performance at the wanted ctr len. For DeciMamba- 130M we use L_base = 2K, 3 = 0.83, decimating_layer = [12,.... .02, min_seq_een = 20. During evaluation keep the same parameters except setting L base 8K. Additionally, in this specific task DeciMamba was trained with a similar, yet not identical, Language Modeling (LM) loss. We break the labels sequence (length 2K) into chunks. The first 1K lahels trained conditionally on the first IK tokens of the sequence (like vanilla LM). The last IK labels are trained conditionally on the whole sequence (2K), and DeciMamba was configured compress the first 1K input tokens. This way we are able to train DeciMamba to compress context while training on each label in the sequence, making the training much more efficient. We also experimented with chunking the labels into more than two chunks, but only experienced a slowdown in computation while achieving similar performance. For the lower bound models we had to reduce the amount of training steps in order to constrain the training to 100M tokens. Specifically, for each context 100M (ttt///bbtchsizee
length, we followed the following formula: oum-of-stepp 100M\ (250 * ctz. len). For the 1.4b model we used Layer 12 for decimation and Liase 4000. For the 2.8b model we used Layer 22 for decimation and Laase 4000.
ctæJeen) =
B ADDITIONAL EVALUATIONS
B.I COMPARISON WITH TRANSFORMERS
We find that vanilla Transformers of equivalent size (trained the same dataset with similar training recipe), have inferior length generalization abilities compared to Mamba. This is evident in multiple long-context tasks (Tables
Table 4: Comparison With Transformers Passkey Retrieval. The setting is the same in AppendixA.I] All models were trained on sequences of length 2k. For each context length we test performance for different needle locations, and report the success rate (between 0 and 1). v stands for 100% success, X for 0% success.
Context Length Pythia-160M
IK
2K
4K
8K
16K
32K
64K
128K
X
X
X
Mamba- 130M 7
V
0.8
X
X
X
DeciMamba-1300
V
Table 5: Comparison With Transformers Zero-Shot Perplexity. The setting is the same as in Appendix A.4 All models pre-trained on sequences of length 2k (we remind that DeciMamba is applied directly without any tuning). inf replaces any perplexity result larger than 100.
Context Length
IK
2K
4K
8K
10K
20K
30K
40K
50K
60K
70K
80K
Pythia-2.8B
10.24
9.96
inf
inf
inf
inf
inf
inf
inf
inf
inf
inf
Mamba-2.8B
9.39
9.17
11.6
inf
inf
inf
inf
inf
inf
inf
inf
inf
DeciMamba-2.SB
9.39
9.17
11.98
14.58
14.73
17.17 19.83
22.2
24.89
27.57 27.89
27.43
Pythia-1.4B
11.64
11.32
inf
inf
inf
inf
inf
inf
inf
inf
inf
inf
Mamba- 1.4B
10.51
10.31
10.5
14.43
18.4
inf
inf
inf
inf
inf
inf
inf 29.28
DeciMamba-L.4B
10.51
10.31
10.5
14.13
14.5
18.33
23.54
26.82 28.97
30.56 29.97
For LongBench, the equivalent Transformer model (Pythia-2.8B) repeatedly causes an OOM error on GPU (A6000, 48GB of RAM).
C HYPERPARAMETER SELECTION
In practice, Liase is the only parameter we sweep. We typically scan 3 or values that are similar in magnitude to the context length used during training (Ltrain). The reason for selecting Liase close to Lirain originates from assumption 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 at most Tarain important tokens; hence, it makes
17

---

Published as a conference paper at ICLR 2025
sense to pool this number of tokens, regardless of the amount of noise. We find this assumption quite reasonable for many long-content tasks, such as retrieval, multi-document question answering, retrieval/QA, only one document is
and next-token prediction. For example, in multi- 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 Lbase close to Leroën is that there are no ERF issues when processing sequences of length Ltroins the global layers (starting from the first decimation layer) are trained on sequences of the same length.
Regarding the number of layers to decimate and the decay rate (8): The main goal of these parame- ters is to improve efficiency, not performance. Both parameters introduce the option to compress the sequence further by applying additional DeciMamba layers. We emphasize that this has neg negligible effects on performance, as shown in the sweeps below for the passkey retrieval task (Tables 6 In addition, this is supported by other results, such as Multi-Document Retrieval and LongBench, which were run with this option disabled.
Table 6: Decay Rate (3) Sweep. BL is the performance of the baseline Mamba model. The setting is the same as in AppendixA.1 We use Luuse 2000 and 9 decimation layers (13-21).
BL
0.75
0.5
0.25
Success Rate
0.55
1.0
0.975
0.975
0.975
Table 7: Number Decimation Layers Sweep. BL is the performance of the baseline Mamba
model. The setting is the same as in Appendix A.1 We use Lbase 2000 and 3 = 0.5
# of Deci Layers Success Rate
BL
1
2
3
4
5
6
7
8
0.55
1.0
0.975 0.925
0.925
0.975
0.925
0.95
0.95
0.975
D OTHER RELATED WORK
D.1 LONG RANGE TRANSFORMERS.
Transformers have emerged as highly effective models for various tasks, yet their widespread adop- tion has been constrained by their limited long-range modeling capabilities. Thus, applying trans- formers effectively to long-range data remains central challenge in DL, particularly in NLP. A primary factor in this challenge that the effective context of transformers is dominated by the context observed during training, which is limited because training LLMs on datasets with billions of tokens across lengthy sequences is computationally demanding. Hence, three main approaches have been developed to tackle this problem: (i) creating efficient variants of transformers that al- low an increase in the length of sequences during training. (ii) Context extension methods, which enable training on short sequences and evaluation on long sequences, and finally, (iii) hierarchical models that rely on pooling, chunking, and compression. Despite these extensive efforts, several re- cent studies indicate that high-quality handling of long text remains unresolved issue (Liu et al.
2024a Li et al. 2024a).
D.2 EFFICIENT TRANSFORMERS.
Over the years, many approaches have been proposed for making transformers more efficient (Tay et al., 2022, Fournier et al.] 2023). The two most prominent directions are hardware-aware im- plementations such as flash-attention (Dao et al. 2022 Dao 2023) and ring-attention (Liu et al.), which accelerate computations over long sequences by several orders of magnitude. Additionally, developing efficient attention variants with sub-quadratic complexity has become very popular. Two notable examples Linformer ( (Wang et al. 2020),, which utilizes a low-rank attention matrix, and Performer (Choromanski et al., 2020), a variant that approximates the attention operator through a kernel function.
18

---

Published as a conference paper at ICLR 2025
D.3 EFFECTIVE RECEPTIVE FIELD (ERF)
The Receptive Field (RF) (LeCun et al., 1998) of a neuron refers to the region of the input space that influences the neuron's output. traditional CNNs, the RF captures how the area of influence expands as the neuron's 's depth increases. However, in modern architectures like transformers, which process the entire input context at each layer and can theoretically model interactions between any parts of the input, the concept of RF becomes less informative. In these models, the theoretical receptive field is maximal, rendering the term ambiguous.
To address this, the ERF offers a more practical measure by quantifying the empirical sensitivity of model to different input regions. The ERF highlights where the model actually focuses its attention and depends not only on the architecture but also on the data and training process. The seminal work on Vision Transformer (ViT) (Dosovitskiy et al. 2020) and Vig & Belinkov (2019) introduced the
concept of attention mean distance to measure the ERF, as detailed in Eq. Building upon this, we
extend the measurement to Mamba layers in Eq. 8 allowing us to effectively estimate their ERF.
D.4 WHY USE LONG CONTEXT LLMS INSTEAD OF RETRIEVAL AUGMENTED GENERATION?
Performance-wise, recent studies do not show a conclusive advantage of RAG over Long Context LLMs: Li et al. 2024b), Xu et al. (022)). Interestingly, the common conclusion of all these studies (and others, such as Anthropic's 's Contextual Retrieval) is that when combined, RAG and Long- Context LLMs have synergistic effect which improves performance even more. Moreover, each study proposes a different way of combining RAG and Long Context L.LMs (directly, switching be- tween RAG and Long Context, augmenting the retriever chunks, etc.) showing that further research in each area alone, or combined, has potential for an overall improvement in LLM performance.
D.5 JUSTIFICATION FOR MAMBA MEAN DISTANCE
We are aware of two main approaches for measuring the ERF of Mamba and Attention models: (i) via the attention mean distance or our analogy of Mamba Distance (Eq. 10), and (ii) by measuring employed for transformers in Chi et al. (2022)
the gradient norms over the entire model, which and very recently applied to Mamba models (Jafari et al., 2024) (section 6). The main draw back of the second approach is that it measures the ERF in the context of gradients, which are significantly influenced by the type of loss function, task, and label, complicating the analysis. Additionally, single layer or set of layers, the second method be-
while the first method can be applied comes less straightforward when applied to a subset of layers, requiring additional considerations and design choices, such as whether to normalize gradients, the use of gradients instead of relevance scores, and the computation of gradients only for attention blocks that learn token interactions or also consider linear and normalization layers. Due to its flexibility and simplicity, as well as being more popular (several examples are Vig & Belinkov (2019)) Dosovitskiy et al. (2020); Raghu et al. (02T)), D), we chose the first method.
E ABLATIONS
E.1 LAYER SELECTION
We ablate our layer selection methodology, as it has significant affect the model's performance (Figure 7 left). Each point on the curve represents the score achieved by a DeciMamba- 130M model in the Passkey Retrieval task when trained on sequences of length 2K. The only difference between the models is the first decimating layer (x axis). We see that when the decimating layer is too shallow (e.g. layer 8), the model fails completely. As we increase the minimal decimating layer we observe a large increase in performance, until reaching the climax at layer 12. This result aligns with the fact that the A, distribution is farther from O in Mamba's is early layers (all tokens are important). We hypothesize that the tokens still don't have strong representation at this stage, hence it makes less sense to decimate w.r.t their A, value. After layer 12 the performance of DeciMamba starts to drop, and stabilizes at the vicinity of the baseline Mamba model. At this region we start decimating too late, as some long-deu dongdeeend layers suffer from ERFs when processing the longer Passkey Retrieval sequences. The final layers (21-23) also have larger A, values, so they should not be
19

---

Published as a conference paper at ICLR 2025
decimated as well. We hypothesize that these layers 'decode' the processed embeddings back into token distributions, so the intermediate representations are yet again not for semantic decimation.
E.2 POOLING STRATEGY
We show the importance of the proposed Top-P, pooling strategy in the following ablation study. Suppose select different pooling strategy, Top-K%, which keeps the top K% tokens with the largest mean A, values (over the channels dimension). We train a DeciMamba-1 130m model for the Passkey Retrieval task sequences of length Itrain 2K, with K 50% and decimation layers [12.... 20). As can be seen in Fig. g (Appendix). Top-X% extends Mamba's 's extrapolation abilities from 8K 32K, yet does not do as well as Top-/', which extends to 128K. While Top- K% decimation decreases the sequence length greatly (the output is x512 times shorter), during extrapolation each SSM layer still faces input sequences that are much longer than the ones it has seen during training, as demonstrated in Fig. 9 We further explain this result; During training Layer 17 saw sequences of length (Ltrain = 2k, followed by a decimation rate of 50% in each layer from layer 12). During inference it sees sequences of lengths 65 and 1002 (where Levat 2K, 32K, left and right images, respectively). Since the sequence length at layer 17 of the latter (1K) is much longer than the one it saw during training (65), the extrapolation is limited by ERF (right image, dashed orange shape). This shows the importance of keeping the lengths of the inputs to each decimating SSM layer similar to the ones seen during training, done by Top-P..
Mamba + Top-K%
Mamba + Top-Ps
Passkey Depth [%]
0-
0.25
0.25
0.5
Train Length
Length
1
1K 2K 4K 8K 16K 32K 64K 128K Context Length [Tokens]
1K
2K 4K 8K 16K 32K 64K 128K Context Length [Tokens]
Figure 8: Pooling Strategy Ablation Results. The figure compares two pooling strategies, Top-P, (ours) and Top- fop-K%. As shown, the Top-K% approach lags behind the Top-P, approach, demon- strating that our strategy allows the model to extrapolate to significantly longer sequences. Results for Mamba-
2K Input Tokens
32K Input Tokens
65
65
1K
Figure 9: Pooling Strategy Ablation. Top-K% pooling leads to limited ERF in layer 17. The size of the attention map is affected by the input sequence length: for Levat = 2K layer will process 65 tokens (left) and for Lenat 32K it will process 1K tokens (right). Since Litrain 2K, layer 17 has only seen training sequences of length 65, therefore suffers from an ERF when Level increases (right, dashed orange shape).
20

---

Published as a conference paper at ICLR 2025
E.3 DECIMATION MECHANISM.
To emphasize the benefit of pooling w.r.t the values of A. we examine two additional decimation mechanisms: max-norm decimation (tokens with maximal norm are kept) and random decimation (tokens are randomly kept). We equip Mamba- 1.4b with each of the decimation mechanisms (deci- mation layer= 12, Luese 8000), and compare each model's zero-shot perplexity in Fig. (right). Max-norm pooling achieves a similar trend to A pooling (DeciMamba), yet performs somewhat worse. The result suggests that A ken's norm also holds some information about its importance. We hypothesize that method that combines both might achieve better performance than any one alone, but leave this question open for future work. Random pooling induces strong distortion the embedded sequence, yet achieves better perplexity for longer contexts. This surprising result demonstrates how sensitive Mamba is to limited ERFs in longer sequences, it is actually better to randomly decimate the sequence rather than process its full length.
better candidate for pooling, but also shows that the to-
E.4 ADDING ADDITIONAL POOLING DURING DECODING.
By design, DeciMamba performs pooling only during the pre- fill phase of inference. Although the
approach is motivated by many previous works (Jiang et al.] 2202aa Ge et al.) 2024 Jiiagg et al.
2024b [ et al. 2023 Jiang et al. 2023 Liu et al. 2024c| Wingate et al., 2022 Mu et al. 2024 further motivate by ablating an ad-
Chevalier et al. 2023 Fei al. 2023] Wang al. (0024b),
ditional pooling mechanism during the decoding phase. We tested this on the longest generation task in LongBench: GovReport, which involves summarizing long documents with an average length of 10,000 tokens. The additional decimation during decoding was performed by iteratively combining prefill and decoding for each generated chunk, where the chunk size is 50 tokens. Table 8 presents the results and shows that adding decoding pooling has a negligible effect on performance.
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 that cause limited ERFs (it needs to be longer than the context length used to train Mamba). Thus, the above suggests that decimation is not needed during decoding. Note also that while in GovReport the generated responses are relatively long, for most of the remaining 15 long-context tasks in LongBench, such as Multi-Document QA, the model is required to generate only 10 to 30 tokens per answer. Therefore, pooling is not required during their decoding phase as well.
Clearly, if new task emerges where the number of generated tokens is larger than the ERF, it might be useful to add pooling to the decoding phase.
Table 8: GovReport: Long Summary Generation With and Without Decoding Decimation. LB LongBench; 0-4k, 4-8k, 8k+ are the three LongBench-e context length groups; and +DD is DeciMamba with additional Decoding Decimation.
Model
LB
0-4k
4-8k
8k+
Mamba
9.8
24.8
10.9
5.2
DeciMamba
14.9
27.4
19.4
7.8
DeciMamba + DD
14.7
26.6
18.4
8.14
F EXTENSIONS TO OTHER MODELS
While this paper focuses on exploring the cength-generalization abilities and introducing the first context-extension method specifically designed for Mamba models, our proposed techniques are broadly applicable to other architectures. First, the Mamba Mean Distance metric can be generalized to other models that rely on implicit attention. Examples include HGRN (Eq. in Qin et al. (0244). Griffin (Eq. 14 in De et al. (2024)) RWKV (Eq. 18 in Zimerman et al. 20224)) and others. These tools can provide valuable insights into the ERF of these models, offering a deeper understanding of their srngth-generalizatten
technique introduced in Sec. can be adapted for other models. Specifically. models with per-token data-dependent recurrent gates can leverage these gates to compute token importance scores, enabling decimation
behavior. Furthermore, the entire context-
21

---

Published as a conference paper at ICLR 2025
strategy during pre-fill. For instance, in Griffin, the recurrent gate rt (Eq. 1 in De al. (2024) be used to assign token importance scores, facilitating an effective decimation strategy. Similarly, in HGRN and GLA, the recurrent gate A, (Eq. 2 in Qin et al. (2024) and G, (Eq. 3 in Yang et al. (2023) accordingly can be interpreted as token importance measures, allowing our approach be extended to design context-extension techniques for these architectures. By extending these techniques, we provide foundation for exploring wongth-generalization in various sequence models.
Table 9: Inference Time Per Block. In milliseconds, benchmarked on Nvidia RTX A6000 GPU. DeciMamba and Mamba blocks have similar processing time. Notice how the processing time drops for blocks that are placed after a DeciMamba block, no matter the input sequence length.
Context Length
Model
8K
16K
32K
64K
128K
256K
512K
Mamba Block
3.2
6.1
11.6
22.7
52.5
95.9
204
DeciMamba Block
3.3
6.1
11.8
23.1
45.9
93.1
188
Mamba Following DeciMamba
1.6
1.6
1.5
1.9
1.8
1.9
2.4
DeciMamba Following DeciMamba
1.3
1.4
1.5
1.5
1.5
1.5
1.5
Needle Token Start 1270
Needle Token Start 400
Figure 10: Demonstrating the "Importance-Sooringg Abilities of A. We evaluate Mamba- 130M on the Passkey Retrieval task, and record the values of A, for all channels of layer 16. Each pair of images is identical, except that the left one marks the location of the passkey with dashed pink line. The horizontal and vertical axes indicate the token number and the channel respectively. can be seen from the two cases examined above, the A, activation captures the needle location successfully, demonstrating the effectiveness of its 'importance scoring' mechanism.
2
1000
4K
8k
100
16k
1000
32k
t
Figure 11: Measuring the Effeets of Limited ERFs. We show the A, values across the channels
for layer 16 the Mamba-1300 130M model, examined on different context lengths. The horizontal and vertical axes indicate the token number and the channel respectively. As can be observed from the results above. the passkey can be detected clearly until the ERF ends (for t > 10K).
22

---

Published as a conference paper at ICLR 2025
mo
333
100
Figure 12: Normalized Mamba Attention Map. Displayed in log scale for each layer of the Mamba-130M model.
23


Reviews:
[{"review_id": "ME8gzctzal", "reviewer": "Reviewer_PVqe", "summary": "This paper proposes a new method to extend the context length of Mamba. The authors start with an explanation of why Mamba in its original form cannot extend the context length. They propose viewing the S6 block as applying an “attention operation” (Section 2) and based on these attention weights, they compute the Mamba mean distance to measure the effective receptive field of Mamba (Section 3). They basically show that the ERF decreases as the context length at inference time increases and that the main culprit is the fast decrease of the sum of discretization steps. To tackle this issue, they introduce Decimamba, a filtering mechanism that discards tokens of lesser importance (Section 4). They finally show on multiple information retrieval benchmarks that the method better performs on long contexts than Mamba.", "questions": "I asked my questions in the weakness section.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "Overall, I like this paper, I think that Mamba is a very appealing method due to its low inference cost and getting methods that allow extending the context length for Mamba is a very important question. I appreciate the fact that the authors considered diverse benchmarks and the gap between Mamba and Decimamba is pretty consistent in some cases. I also appreciated the scientific approach in the paper that consists in isolating  the problematic component in Mamba and proposing a method to alleviate the issue.", "weaknesses": "I have a few concerns regarding this paper that I list below: \n\n- **Hyperparameter choices**: I agree with the fact that the fast decay of the sum of the discrete time steps may explain the lack of length generalization. However, 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 later and the number of layers to decimate. And it does not seem very clear how to make these choices without a gridsearch?\n- **Ablations**: have you tried to do an ablation with respect to the number of layers to be decimated or with respect to the decay rate? \nWriting should be improved: I found the decimation strategy in Decimamba (Section 4) pretty hard to follow. I think that this should be better explained. \n- **Tasks where Decimamba offers much higher gains?**: we see that Decimamba leads sometimes to big improvements (Squad and Passkey retrieval) but on Multi-Document QA, the decay of Mamba seems to also be important. Do you have an explanation/intuition of the type of tasks where Decimamba leads to substantial improvements?\n\n\n\nMinor points:\n\n**Improving the introduction**: I think the author should maybe better explain the Decimamba method in the introduction (and maybe add Figure 5 to the introduction?). \n- **Lack of length generalization of Mamba**: I was a bit confused about the fact that the authors were surprised by the lack of length generalization of Mamba. I agree with the authors that compared to Transformers, the context cache (the hidden states) do not grow with the context length and thus Mamba can scale the context length to infinity. However, this state is **bounded** and thus the model cannot store all the information it may need . In an information retrieval setting, when the number of documents is large, it is easy to imagine that Mamba has trouble deciding which tokens to store. Anyway, this all to say that the authors should maybe highlight that one of the factors explaining the poor length generalization is the small state space size. I think that this is aligned with the method proposed by the authors in that it looks to decimate some non-important tokens that may be captured in the hidden state.\n- **Influence of hidden state size on length generalization**: one thing that i am curious about is: have you tried to study the length generalization of models of similar scale but with one of the models having a bigger state space?\n- **Comparison with Transformers**: have you tried to run Transformers in the benchmarks of Section 5? Just curious to see the gap of DeciMamba with Transformers.\n- **Out of memory**: Maybe I would add a comment saying that the OOM happens in Figure 6 because the complexity at **training time** is O(L). Most people may have in mind the complexity of Mamba at inference time, that is independent of the context length and may be confused by the OOM you encounter. \n- **Typo**: In equation (8), i think the letters \"j\" and \"L\" were swapped. It should be $d(L,j)$.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "xPKah23ymI", "reviewer": "Reviewer_gREq", "summary": "This paper investigates Mamba's length extrapolation capabilities through a systematic analysis approach. The authors introduce the Mamba Mean Distance metric as a novel way to quantify and analyze Mamba's ability to model long-range dependencies. Their findings reveal inherent limitations in Mamba's capacity to handle longer sequences effectively. To address this challenge, they propose a selective token processing mechanism leveraging the delta t parameter from the Mamba formula. This approach intelligently filters tokens by retaining those with larger delta t values, effectively aligning the input complexity with Mamba's modeling capacity. The method is specifically designed for the prefilling phase and demonstrates promising improvements in Mamba's long-text processing capabilities.", "questions": "- The relationship between distance and performance shown in Fig 1 appears counterintuitive. If Mamba struggles with long-distance relationships, why does the performance degradation manifest primarily in middle lengths rather than showing a clear diagonal boundary from top-left to bottom-right between the red and green regions? Could this suggest a more complex underlying mechanism?\n- Given the increasing importance of long-text generation in applications like complex reasoning and creative writing, have the authors explored possibilities of extending this method beyond prefilling to support long-text generation scenarios? It would be benefitial to discuss any challenges the authors foresee in extending the approach to generation tasks.\n- The experimental setup shows variations in model selection across different figures and analyses (e.g. Fig 3 and Fig 4). Could the authors provide more detailed justification for these choices and discuss how they might impact the generalizability of the findings?\n- 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?", "limitations": "", "rating": 3, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- The investigation of token importance scoring in the context of Mamba represents a valuable contribution to the field, especially given the growing interest in alternatives to attention-based mechanisms\n- The method achieves substantial improvements while maintaining implementation simplicity, making it readily applicable in practical scenarios", "weaknesses": "- 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\n- 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\n- Incomplete solution to fundamental limitations: While the approach provides a practical workaround, it doesn't address the underlying limitations of Mamba in processing long sequences. A more thorough analysis of the Mamba Mean Distance metric could potentially lead to more fundamental solutions", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "SuM4bGRv95", "reviewer": "Reviewer_44XJ", "summary": "This paper presents the limitations of Mamba in terms of length generalization and proposes an algorithm that combines token dropping with the Mamba architecture to effectively enhance its length generalization capabilities.", "questions": "Please refer to the Weaknesses.", "limitations": "", "rating": 3, "confidence": 5, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. The proposed method is straightforward and can be directly combined with Mamba to improve the length generalization abilities.\n2. The authors analyze the shortcomings of Mamba in terms of its length generalization abilities.", "weaknesses": "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.\n2. 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). Additionally, for simpler long-context tasks like Passkey, the authors mention that DeciMamba requires fine-tuning to handle them effectively.\n3. The presentation of the paper could be improved. For instance, proper citation formats (\\citep and \\citet) should be used; table captions should appear above the tables; and in Table 1, the DeciMamba entry is missing the LB score for the LCC dataset.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "BKMb7xtC3G", "reviewer": "Reviewer_txNF", "summary": "This paper studies the question of how to understand and address the limitations of length generalization capabilities of the Mamba architecture. The paper introduces ERF, the effective receptive field, and shows that Mamba is biased towards sequence lengths seen during training. The authors then introduce DeciMamba to extend length generalization capabilities by limiting the number of tokens that the S6 layer processes.", "questions": "- 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?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- The paper introduces the phenomenon of \"limited effective receptive field\" which could be useful for future work studying the Mamba architecture\n- The proposed design is well motivated by the identified ERF phenomenon\n- The proposed architecture shows consistent improvement over Mamba in long context range experiments\n- The paper includes effective visualizations, especially Figure 2", "weaknesses": "- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)