MambaExtend: A Training-Free Approach to Improve Long Context Extension of Mamba
Authors: Seyedarmin Azizi
Souvik Kundu
Mohammad Erfan Sadeghi
Massoud Pedram
Venue: ICLR 2025 Poster (2025)
Keywords: Mamba
Long Context Generalization
Discretization Step
SSM
Primary area: foundation or frontier models, including LLMs

Abstract: The inherent quadratic complexity of the attention mechanism in transformer models has driven the research community to explore alternative architectures with sub-quadratic complexity, such as state-space models. Mamba has established itself as a leading model within this emerging paradigm, achieving state-of-the-art results in various language modeling benchmarks. However, despite its impressive performance, Mamba's effectiveness is limited by its pre-training context length, resulting in a pronounced degradation when the model is tasked with handling longer contexts. Our investigation reveals that Mamba's inability to generalize effectively to long contexts is primarily due to the out-of-distribution (OOD) discretization steps. To address this critical limitation, we introduce _**MambaExtend**_, a novel framework designed to significantly enhance the context extension capabilities of Mamba. Specifically, MambaExtend leverages a _**training-free**_ approach to calibrate _only_ the scaling factors of discretization modules for different layers. We demonstrate both gradient-based and gradient-free zeroth-order optimization to learn the optimal scaling factors for each Mamba layer, requiring orders of magnitude fewer updates as opposed to the parameter fine-tuning-based alternatives. 
Using this approach, we achieve a training-free context extension of up to 32x, expanding the context from 2k to 64k tokens with minimal increases in perplexity. In contrast to existing fine-tuning methods, MambaExtend selectively calibrates the scaling factors, requiring up to $\mathbf{5.42 * 10^6} \times$ fewer parameter updates and incurring up to $\mathbf{3.87} \times$ lower peak memory usage, while delivering comparable or superior long-context performance across multiple tasks. Codes and checkpoints are available here$^1$.

Published as a conference paper at ICLR 2025
MAMBAEXTEND: A TRAINING-FREE APPROACH TO IMPROVE LONG-CONTEXT EXTENSION OF MAMBA
Seyedarmin Azizi"t, Souvik Kundu'i, Mohammad Erfan Sadeghi", and Massoud Pedram" "University of Southern California, Los Angeles, USA
Intel Labs, USA
i Equal contribution authors
[seyedarm, sadeghim, pedram @usc. edu, souvikk. kundu@innent com
ABSTRACT
The inherent quadratic complexity of the attention mechanism in transformer models has driven the research community explore alternative architectures with sub-quadratic complexity, such as state-space models. Mamba has estab- lished itself leading model within this emerging paradigm, achieving state- of-the-art results in various language modeling benchmarks. However, despite impressive performance, Mamba's effectiveness is limited by its pre-training context length, resulting in a pronounced degradation when the model is tasked with handling longer contexts. Our investigation reveals that Mamba's 's inability to generalize effectively to long contexts is primarily due to the out-of-distributiin (OOD) discretization steps. To address this critical limitation, we introduce Mam- ba baExtend, a novel framework designed to significantly enhance the context exten- sion capabilities of Mamba. Specifically, MambaExtend leverages a training-free approach calibrate only the scaling factors of discretization modules for dif- ferent layers. We demonstrate both gradient- based and gradient-free zeroth- optimization learn the optimal scaling factors for each Mamba layer, requiring orders of magnitude fewer updates as opposed to the parameter fine-tuning based alternatives. Using this approach, we achieve a training-fiee context extension of up to 32x, expanding the context from 2k to 64k tokens with minimal in- creases perplexity. In contrast to existing fine-tuning methods, MambaExtend selectively calibrates the scaling factors, requiring up to 5- 42.2 10 x fewer pa- rameter updates and incurring up to 3.87x lower peak memory usage, while de- livering comparable superior long-context performance across multiple tasks. Codes and checkpoints are available here
INTRODUCTION
Despite the widespread applications of trans- former (Vaswani 2017) based large language models (LLMs) Touvron et al.) 2023). their quadratic compute and memory demand with se
Mamba-130M
MambaExtend-1300 [ours]
Mamba2-780M
MaamaaZExtedd-70 7-0M [ours)
sequence length has enforced research for
Perplexity
107
emerging alternative architectures. For exam- ple, works including Linformer (Wang al. 2020) and Longformer (Beltagy et al. (2020) presented different approaches approximate attention to reduce the quadratic memory cost, Other works (Kitaev al.) 2020) leveraged locality-based hashing to avoid attention com- putation. Recently, state-space models (SSMs) (Gu et al.) 2022 2020) have emerged an al- ternative to attention-based models, offering a
10a
64k
Context length
Figure Long-context understanding on Pile. Compared to the pre-trained alternatives, Mam- baExtend provides up to ~8145 improvement in perplexity score, via training free calibration.

---

Published as a conference paper at ICLR 2025
different approach to handling long sequences sub-quadratic complexity. Unlike transformers, SSMs are grounded in continuous-time time dynamics and offer the potential to handle much longer se- quences without blowing out the memory and compute demand. Mamba (Gu & Dao 2023 Dao & Gu 2024), a popular SSM variant built leveraging the selective state-space layers (S6), has shown impressive performance on various NL.P, image, and medical genomies benchmarks (Schiff al.,
2024). The key advantage of Mamba stems from the sub-quadratic compute complexity of theoreti-
cally grounded linear RNN layers.
LLMs for long-contens understanding have recently found many useful applications, including sum- marizing long documents and answering long questions (Chen et al., 2023c). However, transformer- based LLMs that are pre-trained on fixed-length contexts yield lower generative performance when used on longer sequences during inference time (Chen et al.) 2024 2023b). This shortcoming of the transformers is tied to the inability of the positional embedding to generalize well on longer sequences, causing such sequences to appear as out-of-distributiin (OOD) sequences (Chen et al.) 2023c Jin et al. 2024). Interestingly, Mamba models, despite their theoretical ability capture global interactions, also fail to generalize to long sequence or context lengths (Ben-Kish et al. 2024 This phenomenon has been tied the Mamba model's implicit bias to limited effective receptive field (ERF) governed by the training data sequence length (Ben- et al.) 2024).
For transformer-based LLMs, the OOD sequence length generalization has been explored exten- sively, including fine-tuning to longer sequences (Chen et al.] 2023c) and allowing sophisticated modification the transformer's positional embedding (Jin et al. [202 Ding el. 2024 Golovneval
iet al. 2024). Unfortunately, such solutions are not directly applicable to Mamba models. This is primarily due to the absence of an explicit positional embedding for Mamba models to generalize. Moreover, unlike transformers, the potential root cause of Mamba's performance deterioration for long sequence processing is yet to be discovered.
A contemporary work, namely DeciMamba (Ben-Kine Kish et al. 2024), has presented a selective token decimation strategy to reduce the number of tokens be processed per layer. This approach poten- tially increases the model's 's ERF, enabling better long-context information flow and understanding. However, DeciMamba requires memory- and compute-inten intensive fine-tume fine-tuning of the model, result- ing in significant time and effort to perform parameter updates of the pre-trained model. Thus, such an approach does not scale to larger models, especially for limited memory or computing resources.
Our Contributions. To mitigate the aforesaid issues, we first investigate the impact of OOD long- context extension on the discretization steps of Mamba (A, values) We have empirically observed two key findings: (1) minimizing the sum of A 's across all tokens can enhance generalization when increasing context length during inference, and (2) applying a fixed down-scaling to A,' 's gener- ally does yield optimal generalization. Building on this insight, we introduce MambaExtend, a framework designed to extend Mamba's context length without any re-training of the model weights. MambalExtend utilizes a calibration function (CF) optimize sizes of the discretization steps (At) across various Mamba layers by incorporating a learnable scaling factor for each layer's s A. This enables the proposed scaling parameters for A, to be learned while keeping the model weights fixed at their pre-trained values. significantly reducing both memory requirements and the number of updatable parameters. Moreover, we introduce a zeroth-order (7.O) optimization- CF to calibrate using only forward passes, potentially yielding further savings in memory and com- putation. Specifically, we used ZO methods based on the simultaneous perturbation stochastic ap- proximation (SPSA) (Spall 1992) to update the scaling factors. As shown in Fig. PPL by up to ~8145,, as evaluated on context lengths of up to 64k.
To show the efficacy of MambaExtend, we performed extensive experiments on perplexity eval- uation, LongBench, and long-contex: retrieval tasks with both Mamba and Mamba2 variants. For example, PG19, only via Z0-based scaling factor update, MambaExtend can improve the context length extension ability of a pre-trained model from 2k to 64k, while not incurring any significant perplexity (PPL. increase. Compared DeciMamba, yield up to 40.6% reduced PPL while requiring up to ~5.42 108x x fewer parameter update with up to 3.87x x lower peak-memory de- mand.
' represents the step size that transforms continuous-time parameters into the corresponding discrete state-space variables.
2

---

Published as a conference paper at ICLR 2025
2 PRELIMINARIES
2.1 THE S6 LAYER AND MAMBA
At its core, each Mamba block utilizes the selective SSM (S6) layer (Gu Dao 2023). which is specifically designed to handle sequential data by preserving structured state dynamics across the input sequence.
The s6 layer: Using linear recurrent system with the hidden state h4, input 2,, and output o at discrete time instant I, the s6 layer's sequence generation and state update can be simplified as:
h Ah Bzg Oz Ch+
The P-length sequence of representative channel is given == (21, #2, ;P)). À RNN,, B E TRXX1, and C' E RRxN are discrete time-variant system, input, and output matrices, respec- tively, governing the discrete state transitions and output sequence generation. The layer pro- duces the 'per-time" (t) discrete time-variant matrices from input and "continuous parameters" as:
A, exp(A,A), B, A, where A, SFFT (FT(Atrr (z))), B, = Wn(#.), C. (Wo(z))
(2) Here, 24 D with D being channel dimension and A, be the discretization step used time t. A Aproy WB. and We are linear projection layers. SFT and exp represent the softplus and point- wise exponential operation, respectively. After the discretization step, the $6 layer's input-output behavior via time-unrolling can be described as:
= aZ with 0i,j C;
II A=j+1
A
B₁
(3)
Thus, for a context length of P, the entire output 0 = f01, 02,
opp is computed as follows:
B
21
02
CrA2BI
C₂B₂
22
(4)
:
OP
A CPI A.
CpBp)
Zp
This matrix formulation shows that each output o, a weighted sum of the inputs 21, 22, ZP, with the weights determined by the state-space matrices A, B, and C. The model can thus integrate information across different time steps while maintaining computational efficiency. This matrix resembles the attention score map in transformer-basee models (Ali et al. 2024). In other words, S6 layers may be interpreted as data-controlled linear operators.
Notably, as these matrices are dynamically adjusted based the input sequence, they enable the model to efficiently capture temporal dependencies across various time steps. This approach allows Mamba to maintain computational complexity that scales linearly with the context length.
Mamba block. One the critical aspects of Mamba's 's architecture is how Mamba block relates its input sequence X (2, ... .PP) to its output sequence Y (y1, V2,..,,,2 ,rr) with P corresponding to the sequence or context length. The relationship between the input and output of the Mamba block is expressed through vire-narying SSM described below:
G o[ proj X), Z ConviD(W (n.))
(5)
0 -= 0=s)( Y 0 G
(6)
Here, G is gating function derived from linear transformation of the input sequence X followed by SILU function, σ. The element wise multiplication O between and allows the model to selectively emphasize or attenuate parts of the input to focus on relevant input information. The input to the S6 a linearly transformed version of the original input X followed by a convolution.
As demonstrated in these equations, the relationship between the last token Op and the first token is governed by the term a Cr II B₁ Cpexp(A DP A1) This means that the exponent of summed A, determines the impact of the first token in the generation of the pth token.
3

---

Published as a conference paper at ICLR 2025
21
212
211
21
23
46
28
8k
27
(A.)))
24
23
21
23
22
21
10
20
50
Layer index
Figure 2: Layer-wise behavior of (A,) for different context length during test-time. We used the Pile dataset on Mamba 1,4B for the evaluation.
3 MOTIVATIONAL CASE STUDIES
The behavioral change of D. We first investigate the behavior of the accumulated discretization matrix A, in the pre-trained Mamba- 1 1.4B model when exposed to inputs of different context lengths.
Using 100 samples from Pile, for each Mamba layer, we compute the II Ae)ll2 for different evaluation context lengths P', where |1-112 represents the 12-morm of a tensor. We plot this analysis in
Fig.
2 which reveals how the accumulation of A, scales with increasing context length. Specifically, Fig. discloses that for each layer of the model, the magnitude of Σ A, increases with the increase
in context lengths P'. According to Equations 3 and 4 and given that all entries of are always of A, appears the exponent in the exp function. Consequently, the term exp(- A, effectively governs the decay influence from any previous token. A larger value of A, results greater forgetfulness, decreasing the model's reliance on earlier tokens. In contrast, smaller A, values enable the model to retain information from more distant tokens. Therefore, exp(- -E= Ar) can be interpreted as parameter that potentially regulates the retention level for the n-th input to compute the token at p'
negative (Gu & Dao 2023), observe that the negative
Inffuence of scaled A. For transformer-based LLMs. a popular method for addressing the out-of- distribution (OOD) context length 1r >> P. (where P represents the training context length) posi- tional interpolation (PI) (Chen et al., (023a). The PI method accomplishes this by multiplying the to- ken index value in RoPE by pr. This rescaling en- sures that the positional indices remain within a valid range, effectively mitigating the OOD problem asso- ciated with longer contexts without retraining.
Baseline PPL (no scaling)
250
Perplexity (PPL)
200
PPL reduces with Increased scaling
PPL increases with Increased scaling factor
150
100
(uniform scaling)
Inspired by this, propose straightforward ap- proach for Mamba to address the accumulated out- of-distribution (OOD) discretization steps by scaling the discretization matrix A, by a fixed scalar value < across all model layers. This method aims to mitigate the OOD effects associated with longer context sizes. We utilized pre-trained Mamba 1.4B validate this approach, conducting grid search over various values of s. We then evaluated the model's performance on the test set of the Pile dataset (Gao al. 2020) for an evaluation context length of 32k tokens, reporting the average perplexity. The results, presented in Fig. 3] demonstrate that scaling A, can significantly reduce the model's PPL from approximately 268 to around 23.5. However, the findings also indicate that the relationship between the choice of scaling value and performance improvement is not straight- forward. As shown in Fig. B while increased scaling helps reduce the perplexity at lower values of (s), the PPL rises after reaching certain threshold. This complex interplay encourages to investigate the model's 's capacity to learn the optimal scaling. Additionally. this uniform scaling factor cannot restore the model's performance for longer contexts to the level observed at its
0.0
0.2
0.4
Scaling factor of De
Figure Impact of different values of uni- form Ar scaling on the perplexity (PPL.) evaluation metric.
4

---

Published as a conference paper at ICLR 2025
pre-trained context length. For instance, the model achieves a PPL of 3.7 for 2k context length, which remains significantly lower than the best PPL obtainable through uniform scaling.
Variable impact of A on different layers. Another important observation in Fig. 2 is that for given test-time context length P', different layers of the model produce significantly different Σ A values (even when viewed on a logarithmic scale). This underscores the point that each layer should not employ the same scaling factor to reduce the impact of A, This observation motivates us to implement heterogeneous (layer-specific) scaling mechanism across the various layers of the model to effectively address the OOD Σ A..
MAMBAEXTEND METHODOLOGY
Motivated by the need to mitigate the OOD effects, we introduce MambaExtend, a training-free method for scaling the discretization steps of each layer. For an L-layer Mamba model, our primary objective is to determine the optimal scaling factors for each layer, denoted as 51, 82,...,, ,, which will be used adjust the discretization matrix A. Note that for layer 8 E R", where indicates that 8; is a scalar, and m > indicates that s; is a vector. Without loss of generality, for m the discretization adjustment inference. The goal in way that is both memory- and compute-efficient and does not involve any additional training or fine-tuning of the model parameters. These constraints will enable such calibration be feasible on resource-limited edge devices.
be expressed as A calibrate the newly introduced learnable parameters s; for all i 1,...,
s; with D applied during
Algorithm MambaEXtend Algorithm
1: Input: An L-layer Mamba model parameterized by M, set of calibration samples C, calibration function CF
2: Output: Scaling factors S = [$1, S2, sz]. where 8; E IR"
3: for < L do
4: 8; < initt((()
5: end for
6: freeze(M)
7: S + CF(
8: return S
Algorithm [ outlines the MambaFxtend framework, which takes pre-trained Mamba model as input, along with a small set of calibration samples from the target task and a specialized function known as the calibration function tors. Importantly, unlike DeciMamba, which allows fine-tuning of the weights, MambaExtend keeps the model weights fixed to their pre- traitrd values (as indicated in Line of Algorithm (1] through- out the calibration process. This approach makes MambaExtend significantly more compute- and memory-efficient compared DeciMamba.
fumnitooionfuuction (CF). ). As its name implies, CF calibrates the learnable scaling fac-
Calibration via propapraon (CF Gradient-based backpropagation a widely used op- timization method for updating the free (unfrozen) parameters on calibration set. However, to minimize computational and memory overhead, we ensure parameter efficiency by restricting up- dates the scaling factors S only. Algorithm 2 summarizes the BP algorithm for finding the optimal scaling factors. We utilize Adam as the optimizer for backpropagation (as noted in Line of Algorithm [ The Evaluate() function in Line computes the loss of the model, which parameterized by frozen weights and the learnable scaling factors S.
Calibration via zeroth-orda optimization (CF zo). Zeroth-order optimization (Spall 1992 Mal-
ladi et al.) 2023b) offers an efficient yet noisier method for calibration, as it relies solely on forward passes to approximate gradients. Algorithm Joutlines the process for optimizing the scaling factors S in CF Specifically, this is multi-iteration process in which, at each iteration, the scaling factors are randomly perturbed using a random variable sampled from a Rademacher distribution.
The magnitude of the perturbation and the learning rate for the updates parameters e and 17. respectively. We employ the two-sided variant of the simultaneous perturbation stochastic approximation method (SPSA) (Spall 1992), which obtains gradient approximations by
controlled by the hyper-
5

---

Published as a conference paper at ICLR 2025
Algorithm 2 CF Algorithm
1: Input: An L-layer Mamba model parameterized by frozen weights M, set of calibration sam-
ples C, the initialized scaling factors
2: Input:
Learning rate 7/ number of iterations K [$1, $2. $L], where s; E RT
3: Output: Learned Scaling factors
4: optimizer= = Adam(S,1)
5: for k S K' do
6: L= EvvalatteMMa.ss.
7:
C.backward
8:
optimizer.stepi
9:
S /S.clamp (min = 0.001) # make sure scaling factors remain positive
end for
11: return S
Algorithm 3 CF Algorithm
1: Input: An L-layer Mamba model parameterized by M, set of calibration samples C, the initial-
ized scaling factors
2: Output: Learned scaling factors S = [s1, 82, $2, .... sr.), where E R
3: Specify learning rate 7/, perturbation magnitude c, number of iterations K
4: for k K do
5: e Rl ~ Rademacher() ()
6:
S S 6, S = -
7:
L = Eval Evaluate(Ma,ss-.0
= E-valuateMMaxxss.
8:
Vs =L L - M(2c))
9:
S +- --1 v
10: S S.clamp (min = 0.001) # make sure scaling factors remain positive
11: end for
12: return
applying both positive and negative perturbations the parameters simultaneously. The two-sided SPSA approach yields gradient estimates with lower variance than the one-sided version, thus en- hancing accuracy, especially in noisy environments (Spall 2005).
The convergence of the zeroth-order calibration method, CFzo. is affected by the number of pa- rameters being optimized, specifically the size of S. Classical lower bounds indicate that conver- gence slows linearly as the number of parameters increases (Nemirovskij & Yudin 1983 Duchi ct al. 2015). Consequently, a natural strategy in our context to employ the backpropagation-banl method, CF BP,, when optimizing larger set of parameters in (S), while reserving CF zo for smaller parameter sets.
Our experiments show that long-cor congecontett evaluation tasks, based the perplexity measure, and the LongBench tasks require relatively fewer scaling factors. Specifically, for each layer 8; E IR+ setting of 1 is sufficient to improve PPL long-context inputs. Here, TR+ represents the set of positive real numbers, as scaling factors cannot take negative values in our case. Any that updates negative value is clamped to a very small positive number ensure this condition in our algorithm. We set m = D for the passkey retrieval task, thereby increasing the number of parameters
calibrated or updated. We empirically find that for the long-context tasks, CF zo performs nearly as well as BF. However, for the passkey retrieval task, we prefer CF due to faster convergence trend compared to the zeroth-order method. future work, we plan to address the tuning of the zeroth-order approach to achieve better convergence rate for relatively high parameter counts.
EXPERIMENTS
This section evaluates the performance and efficiency of our proposed MambaExtend. Specifically, first describe the models and datasets used for experiments. We then present extensive

---

Published as a conference paper at ICLR 2025
empirical results outline our findings regarding the long-context performance of the Mamba model variants. We finally discuss the compute, time, and memory requirements for MambaExtend.
5.1 EXPERIMENTAL SETUP
Models and datasets. To evaluate the performance of MambaExtend, we use both long-context understanding and long-context retrieval ability tasks. For long-context understanding,
the Pile (Gao et al. 2020) and PG-19 (Rae et al. (2199) datasets and assess the performance of the Mam- baExtend in terms of perplexity scores various context lengths. We use Mamba-130M, Mamba- 1.4B (Gu & Dao 2023), and Mamba2-780M (Dao & Gu. 2024) for these evaluations. Additionally, the LongBench benchmark (Bai et al. (023) to evaluate the performance accuracy of the Mamba-1.4D and Mamba2-780M models. In specific, we use seven tasks, namely Qasper (single- document QA), HotpotQA, 2WikiMulthhoQAA (multi-document QA), TREC, TriviaQA (few-shot learning), LCC, and RepoBench-P (code completion). For the passkey retrieval task, we follow the setup described in (Ben- Kish al.] (2024) and evaluate the performance of the Mamba- 130M and Mamba-1 1 .4B models in retrieving 5-digit code embedded at a random sequence depth within sam- ples from the WikiText- 103 dataset (Merity et al. 2016). In our retrieval setup, the input sequence
lengths range from KK to 64K tokens.
Baseline and SoTA comparison. We use the pre-trained Mamba (Gu & Dao 2023) and Mamba2 (Dao & Gu 2024) models to evaluate the baseline performance as we increase the evaluation context length P.. We use DeciMamba (Ben-Kish et al.] 2024), a contemporary work that uses memory- intensive fine-tuning to update all the parameters while improving the effective receptive field.
5.2 EXPERIMENTAL RESULTS
Perplexity evaluations on PG-19 and Pile. To evaluate perplexity (PPL) on the Pile and PG-19, we
use twenty calibration samples from the corresponding training set for a given context length. We
use these samples to learn the scaling factors in Mam- baExtend, then evaluate perplexity on the test set for given context length. As stated earlier for the per- plexity evaluation, for each layer i, we use a single scaling factor s; It, per laye which scales the
100
200 Po0kcom contere
pPL: 2896
D9L: 905328
@970k contont
Mamba-130M o
DeciMambe-1 30M X
s0
Mambaxitend-1330MP (ours)
Perplexity (PPL)
Mambo-14BB V
DediMamba 1.4 X
A, tensor uniformly for that layer. Therefore, in an L-layer Mamba model, we optimize L scaling fac- tors for these datasets. Given the small number of pa- rameters optimize, we use CF 20 as the calibration function.
MambaExtand 1.48 (ours)
40
20
9.75
Fig. depicts the performance of MambaExtend compared to the pre-trained Mamba variants and DeciMamba. Specifically, at 70k context length, MambaExtend-1300 yields a PPL of 30.62, a ~32506 x improvement over the baseline counter- part that fails to provide very high PPL of 995328. Compared to the DeciMamba, it shows consis- tent improvement with reduced PPL of up to 40.66%
20
30
40
50
60
70
Context Length (k)
Figure 4: Perplexity comparison on PG-19. The / and X identify the fine-tuning quirements to be false and true, respectively.
Table reports the PPL values of MambaExtend models and compares them to those of the pre- trained models on Pile. As shown in the table, MambaExtend, through only minimal calibration, allows the models to maintain their performance even with increasing context lengths. Specifically, MambaExtend can improve the PPL by up to ~8145x showing higher improvement trends at longer contexts.
LongBench. LongBench Bai et al. 2023) is a benchmark for bilingual, multitask, and compre- hensive assessment of long-context understanding. For ManbaExtend, we use seven popular tasks from LongBench. Due to the lack of training data, we used 10 samples from the 4K-8K split of each dataset as calibration data and the remaining samples from the same split to evaluate. We apply the
'This may be attributed to the relatively more straightforward nature of long- congecontext understanding op- posed to long-context retrieval, since for the latter we need more fine-grain scaling, increasing the number of calibration parameters.
7

---

Published as a conference paper at ICLR 2025
Table 1: Perplexity for Mamba models over different evaluation context lengths on Pile dataset.
Mamban- 10M
Mamnaa11
1.BB
Mambaa-780I
Conteat Length
24
4k
88
166
32k
64k
2K
46
8K
16k
32K
64%
2k
4k
8k
16k
32k
64k
Pre ree-traied Model
700
6.18
6.22
46592
4.34
3.78
4.19
14.4
260
0.004
4.78
4.02
22.
79
185
378
MamnhFateed
7.06
5.16
8.7
4.31
3.78
3.48
3,62
6.93
Table 2: Mamba vs MambaExtend performance on representative LongBench tasks.
Model
Qasper
HotpotQA
2WikiMutthoOOA
TREC
TriviaQA
LCC
RepoBench-P
Average
Mamba- AB
7.0
11.00
9.75
29.00
L.7
20.12
11.67
12,88
MambaEExend
14.29
13.82
35.0
7.57
18.84
18.91
Mamha2 780M
7.50
6.06
9.48
17.0
0.1
22.1
14.01
10.89
MambaExtend227 7-7MMM
7.96
10.95
18.33
28.00
6.83
28.27
17.71
Manio-1100 FTll Score 71.4%
DesiManbb-1300 130 Ratioval Seare - 65.7%
Mamballstud--
(005
Prodaey Doptn (9)
100%
Partbay Deyn PM]
100%
75%
75%
50%-
50%
5055
28%
0%
0%
0%5
IK 2X
1X 1RK
648
1K
2K 4K 4K
16K 32K 64K
3K 4K 3X
1 33K 64K
Lengls
Context Length
Contexi Langh
Mambu-1440 Fall FT. & Scre R2.9%
DeciMamba-1 4lt, Rarizaa Samm 74.2%
MambaFsstmmd 148 Ratioval
Pookay Doyth [N]
100%
Pankky Dapth (%2]
100%
75%
75%
75%
50%
25%
25%
25%
2K
15K 64K
1K
4K
32K
IK
4K
5K 16K
64K
Contect Laugh
Context Leggh
Context Length
Figure 5: Passkey retrieval performance after fine-tuning (FT) (for Mamba and DeciMamba) or calibrating (for MambaExtend) on samples of 4k context length.
CF zo calibration function to learn the scaling factors. Similar to the calibration setup for perplexity evaluation, we calibrate one scaling factor per layer shared over the whole A, tensor for that layer. As demonstrated in the Table 2 MambaExtend can improve the average LongBench accuracy by up to 6.03%.
Passkey Retrieval. Previous works have demonstrated that tasks requiring exact retrieval are more challenging than achieving low perplexity in longer context (Liu et al. 2024), so we use more fine- grained sharing of scaling factors to optimize. For A, tensor of a layer i. per channel yielding total D scaling factors per layer (s; E R.). Unless otherwise stated, we use CF BP for one epoch to calibrate standard fine-tuning with the same context length for one epoch as we get significant failure in the retrieval. For DeciMamba to have fair comparison, we fine-tune for the same epochs our
one scaling factor
dataset with 4k context length. For the haseline, performed
The evaluation is conducted across context lengths of 1K, 2K, 4K, 8K, 16K, 32K, and 64K, with the target digit hidden depths of 0%, 25%, 50%, 75%, and 100% of each of these sequence. Assuming that each correct retrieval receives a score of and each incorrect retrieval receives a score Total coreect retrievals
of O. we compute the retrieval score percentage (%) as Total (correct + incanecti retrievals 100, across all the depths overall context lengths. The result is demonstrated Fig. 5 Although MamnaExtend
calibrates approximately 3500x and 7100x fewer parameters for Mamba- 130M and Mamba- 1.4B, receptively, performs better or very similarly to the other two alternatives.
5.3 COMPUTE, TIME, AND MEMORY COST ANALYSIS
Fig. 6 demonstrates a comparison of full finetuning of baseline Mamba, DeciMamba, and calibra- tion tuning with MambaExtend for the passkey retrieval task. Note here that to have a fair compar- ison and to demonstrate efficacy at extreme lost cost tuning. we set the epoch to one for all. For
IIn the original paper (Ben- Kish et al. 2024) the model was tin-ttuned for longer duration, however we focus on limited resource calibration and thus keep experiments limited to fine-tuning for one epoch. Please Appendix for fine-tuning results with longer epochs.
8

---

Published as a conference paper at ICLR 2025
DeciMamba-1300 30M--KK
Memba-130M-4k
Hambalatend-130M--
MambaExtend 130M
2.00
2 1.25
Models
Figure 6: Comparison normalized {peak memory, calibration time, and number of parameter updates) between Mamba, DeciMamba, and MambaExtend for passkey retrieval task. We use Mamba-1300 130M model and for each method, we train for one epoch either with 4k or with 8k context length. For these three measurement types, we normalize each value by the corresponding value of MambaFxtend- 300bbbb--xedd1-00-44.. 130M-4k.
MambaExtend, we show results for only perform experiments with tuning with contexts. Notably, MambaExtend requires up to 2.12x fewer memory for tuning with similar context; in other words, it can support calibra- tion with higher context of up 2x. Regarding per epoch calibration time, MambaExtend can be faster by up 1.69x while requiring up 3532.6x fewer parameters to update. To measure the retrieval success, we compute the Interestingly, despite having significant calibration efficiency, 4k tuned MambaExtend provides up to 20% improved accuracy. We yield even better efficiency for CF based calibration. In specific, compared to DeciMamba, MambaExtend requires up to ~5.42 106x fewer parameter updates and costs up to 3.87x lower peak-memory (details pro- vided in Appendix A.3
with both 4k and 8k contexts, while for others, we
5.4 DISCUSSION AND ABLATION STUDY
Understanding the impact of learned scaling on A, To understand the benefits of the learned scaling the A, discretization tensor, we compute the normalized sum of A,
Mamba
2000
IK 12 Here, n refers to the to- ken index whose impact we want to study on the output context length P p' is set to 32k for this analysis. The Fig. demon-
Tok.m Index
5000
4000
8000
strates the heatmap of the 11 A,)
3000
MamnbaExtend
for different token index (n) at different layers of the model. Notably, as discussed earlier, high I| IE At)ll value may be associated with a stronger decaying effect on the output token p'. As we can original Mamba, particularly for later layers, induces significant decaying ef- fect for the earlier tokens (see the value for token index 2000 for layer index 40). This finding aligns with that of Kish the con-
2000
2,000
Token Index
5000
1000
10
30
Index
Figure 7: Impact of the calibrated scaling factors A, (Top) layer-wise normalized sum of A, for pre- trained Mamba. (Bottom) layer-wise Normalized sum of A, layer-wise for a MambaExtend calibrated model. We used Mamba-1.44 on Pile with 32K context.
al. (2024). MambaExtend, trary, reduces this effect significantly, both overall and for later layers.
Ablation on the granularity of scaling factor sharing. Table presents the re- sults with various levels of sharing of the scaling factor layer's A.- Specifically, we allow per-channel, per-token, and per- tensor sharing where scaling factor shared over channel, token, and the whole tensor for a layer's Dc, respec- tively. We calibrate for one epoch for three scenarios. As we can see from the table, per-channel sharing improve the retrieval score significantly. While per-tensor sharing quires considerably fewer calibration parameters, it fails to yield good score, making per-channel sharing optimal choice.
Table 3: Passkey retrieval on Mamba-130M, with a dif- ferent granularity of the scaling factor sharing.
Sharing granularity Params. Retrieval Score (%)
Per-channel
36.8K
91.4
Per-token
98.3K
62.8
Per-tensor
24
22.8

---

Published as a conference paper at ICLR 2025
Ablation BF ZO For simpler long-context understanding tasks we demonstrated 2o to yield signif- icantly improved PPL. In Table 4 we now demonstrate direct comparison of the two calibration functions, namely, CF zo and CF BP for Pile dataset with Mamba-130M. As we can see, the perplexities for the three evaluation context lengths are similar for both of these methods. This experiment demonstrates the efficacy of CF zo along with its efficiency over BP-based alternative.
Table 4: Perplexity result with CFzo z0 v.s. CF Pile dataset.
CF CContext Length
4K
8K
16K
BP
6.10
5.11
4.79
zo
6.18
5.03
4.84
6 RELATED WORK
Long-context understanding for LLMs. (Chen al.,
2023a)
introduced positional interpolation
to mitigate the issue of OOD positions for contexts exceeding the pre-training length in RoPE-based transformers. In parallel, works such as (Han et al.
2024 Jin al., 2024) proposed zero-shot tech- niques that constrain positional indices to discrete integer values when handling extended contexts in transformers. Additionally, (Chen al.
2024 employs evolutionary search to design non- uniform position interpolation and initialization strategy for fine-tuning on longer contexts. The YaRN method (Peng et al.) 2024) further advances this line of work by combining positional in- terpolation with dynamic NTK-aware scaling, which dynamically adjusts the scaling of high- and low-frequency components of positional embeddings based on sequence length. Despite signifi- cant progress in transformer-basee LLMs, long- long-context understanding for SSMs it yet to be fully unveiled. Only recently. inspired by the success of LongLoRA [Chnn et al. (023c)). DeciMamba BBen-Kish et al. 2024) has proposed fine-tuning- -ased context-extension for pre-trained models.
Zeroth-order optimization. Zeroth-order (ZO) optimization refers to a class of optimization al- gorithms that do not have backpropagation-bain based gradient computation. Instead, the ZO methods estimate gradients indirectly by querying function values through only forward passes. Over the past years, several techniques have been developed for ZO gradient estimation. Randomized Gradient Estimation (RGE) (Nesterov & Spokoiny) 2017) approximates gradient by randomly perturbing the input in multiple directions and examining the function value change. The perturbation is typically drawn from random distribution, such as Gaussian or Rademacher. It potentially requires fewer function evaluations than alternative methods like finite differences (FD) (Shi et al.. 2021). Simulta- neous perturbation stochastic approximation (SPSA) (Spall 1992) is highly efficient ZO method for minimizing multivariate loss functions. Unlike the RGE and FD method, which requires multi- ple evaluations per iteration, SPSA perturbs all input dimensions simultaneously, requiring only two function evaluations per iteration, regardless of the problem's dimensionality. This makes SPSA es- pecially attractive for large-scale optimization tasks. Recently, various algorithms, including MeZO (Malladi et al.) 2023a). further improved the memory efficiency of SPSA. In the MambaExtend, we gain efficiency benefits by optimizing small number of parameters.
CONCLUSIONS
We addressed the limitations of Mamba in handling long-context tasks by introducing MambaEx- tend, framework to extend the context length of Mamba models without model training. Through non-uniform calibration the discretization matrix (A,) scaling factors across different model lay- ers, we enabled context extension by up to 32x while maintaining similar perplexity levels. We believe this work opens up new possibilities for efficient, training- free adaptation of state-space models to long context applications, potentially allowing the true potential of sub-quadratic models to be unveiled. We hope our findings will inspire the community to delve further into the global- local ERF (Xiao et al. (0044 aware tuning, as well as the theoretical underpinning of the relation between discretization steps and OOD generalization.
8 ACKNOWLEDGMENTS
We would like to acknowledge the constructive discussions and feedback from the anonymous viewers of the International Conference of Learning Representation 2025. This research was sup- ported in part by grant from the National Science Foundation.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Ameen Ali, Itamar Zimerman, and Lior Wolf. The hidden attention of mamba models. arXiv preprint 0tt.:22433.01990, 2024.
LongMamba Authors LongMamba, Longmamba: Enhancing mamba's long-context capabilities via training-free receptive field enlargement. Openreview ICLR 2025 submission, 2024.
Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, et al. Longbench: A bilingual, multitask benchmark for long context understanding. arXiv preprint arXiv:2308. 14508, 2023.
Iz Beltagy, Matthew E. Peters, and Arman Cohan. Longformer: The long-document transformer, 2020. URL https://aaaaxttt orgg abs 2004 05150
Assaf Ben- Kish, Itamar Zimerman, Shady Abu-Hussein, Nadav Cohen, Amir Globerson, Lior Wolf, and Raja Giryes. Decimamba: Exploring the length extrapolation potential of mamba. CoRR, abs/2406. 14528, 2024. doi: 10.48550/ARXIV.24000 14528. URL https:://d/o /ttp::///oo..r//...... 48550/ 2...0/aXXi..o 2406. 14528
Assaf Ben-Kish, Itamar Zimerman, Shady Abu-Hussein, Nadav Cohen, Amir Globerson, Lior Wolf, and Raja Giryes. Decimamba: Exploring the length extrapolation potential of mamba. arXiv preprint 2rXiv:2406.145 14528, 2024.
Shouyuan Chen, Sherman Wong, Liangjian Chen, and Yuandong Tian. Extending context window of large language models via positional interpolation. CoRR, abs/2306. 15595, 2023a. doi: 10. 48550/ARXIV.2306.1 15595. htthttps: //doi. org/10 48550 www.a. 2306. 15595
Shouyuan Chen, Sherman Wong, Liangjian Chen, and Yuandong Tian. Extending context window of large language models via positional interpolation. arXiv preprint arXiv:2306. /5595, 2023b.
Yukang Chen, Shengju Qian, Haotian Tang, Xin Lai, Zhijian Liu, Song Han, and Jiaya Jia. Longlora: Efficient fine-tuning of long-context large language models. arXiv preprint arXiv:2309. /2307, 2023c.
Yukang Chen, Shengju Qian, Haotian Tang, Xin Lai, Zhijian Liu, Song Han, and Jiaya Jia. Lon- glora: Elficient fine-tuning of long-context large language models. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenRe- view.net, 2024. URL htLttttts:///pnneeeee
net forummpd--6nJJfffa.
Tri Dao and Albert Gu. Transformers ssms: Generalized models and efficient algorithms through structured state space duality. International Conference on Machine Learning, 2024.
Yiran Ding, Li Lyna Zhang, Chengruidong Zhang, Yuanyuan Xu, Ning Shang, Jiahang Xu, Fan Yang, and Mao Yang. Longrope: Extending Ilm context window beyond 2 million tokens. arXiv preprint arXiv:2402. 13753, 2024.
John C Duchi, Michael Jordan, Martin J Wainwright, and Andre Wibisono. Optimal rates for zero-order convex optimization: The power of two function evaluations. IEEE Transactions on Information Theory, 6185)2778822866, 2015.
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, et al. The pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2101.00m0 00027, 2020.
Olga Golovneva, Tianlu Wang, Jason Weston, and Sainbayar Sukhbaatar. Contextual position en- coding: Learning to count what's important. arXiv preprint 11X11:2405.18719, 2024.
Albert Gu and Tri Dao. Mamba: Linear-ting time sequence modeling with selective state spaces. CoRR, abs/2312.00753 2023. doi: 10-4550/ARXI.2312.00555 URL https: /oi. orq/10.
48550 10.50/arXivv 2211.22311220 00752

---

Published as a conference paper at ICLR 2025
Albert Gu, Tri Dao, Stefano Ermon, Atri Rudra, and Christopher Ré. Hippo: Recurrent memory with optimal polynomial projections. In Hugo Larochelle, Mare' Aurelio Ranzato, Raia Had- sell, Maria-Florina Balcan, and Hsuan-Tien Lin (eds.), Advances in Neural Information Process- ing Systems 33: Anmual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, December 6- 12, 2020, virtual, 2020. URL httpt:pp:://poceeedddeee Atstaaaaaeee.bebbbataatat....... Abstract.
nwwrips....... cc/
Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25 29, 2022. OpenReview.net, 2022. httphttpps//ppeeeeeeeee Eorum? Berumi1d=uYLFo2Iv/
Chi Han, Qifan Wang, Hao Peng, Wenhan Xiong, Yu Chen, Heng Ji, and Sinong Wang. Lm-infinite: Zero-shot extreme length generalization for large language models. In Kevin Duh, Helena Gómez- Adorno, and Steven Bethard (eds.), Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguisties: Human Language Technologies (Vol- ume 1: Long Pupers), NAACI. 2024, Mexico City, Mexico, June 16-21, 2024, pp. 3991-4008. Association for Computational Linguisties, 2024. doi: 10. 111166//1122224AAAC--L11000 naa. 10..111..1666.. wtw.://0...aa1111araaa..
URL
Hongye Jin, Xiaotian Han, Jingfeng Yang, Zhimeng Jiang, Zirui Liu, Chia-Yual Chang, Huiyuan Chen, and Xia Hu. Lim maybe longlm: Self-extend Ilm context window without tuning. arXiv preprint arXiv:2401.01322 2024.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020. URL Vk
net. net/forumiid
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, 2024.
Sadhika Malladi, Tianyu Gao, Eshaan Nichani, Alex Damian, Jason D. Lee, Danqi Chen, and Sanjeev Arora. Fine-tuning language models with just forward passes. In Alice Oh, Tris- tan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine (eds.), Ad- vances in Neural Information Processing Systems 36: Annual Conference on Neural Infor- mation Processing Systems 2023, NeurIPS 2023, New Orleans, LA, USA, December 10 16, 2023, 2023a. URL 4627810151be4al1f900 wtt.aeeeaaa331111188868181200010 Absttsrrtt---tstttt------a-
nips wwtp:esnnpss..c.ppppeeeeeeepee fiwes/paper/e/ filess/paper2222//ees
Abstract Conference. html]
Sadhika Malladi, Tianyu Gao, Eshaan Nichani, Alex Damian, Jason Lee, Danqi Chen, and Sanjeev Arora. Fine-tuning language models with just forward passes. Advances in Neural Information Processing Systems, 36:53038 -30038,5377,,11 2023b.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. arXiv preprint arXiv:1609.000 07744, 2016.
Arkadij Semenovič Nemirovskij and David Borisovich Yudin. Problem complexity and method efficiency in optimization. 1983.
Yurii Nesterov and Vladimir Spokoiny. Random gradient- free minimization of convex functions. Foundations of Computational Mathematics, 17(2):527-566, 2017.
Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole. Yarn: Efficient context win- dow extension of large language models. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna. Austria, May 7-11, 2024. OpenReview.net, 2024. URL https::/
net/ iddrummmidwwiBBfhhu..
Jack W Rac, Anna Potapenko, Siddhant M Jayakumar, and Timothy P Lillicrap. Compressive transformers for long-range sequence modelling. arXiv preprint arXiv: /9/ 0.5507, 2019.
12

---

Published as a conference paper at ICLR 2025
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.
Hao-Jun Michael Shi, Melody Qiming Xuan, Figen Oztoprak, and Jorge Nocedal. On the numerical performance of derivative-free optimization methods based on miite-difference approximations. arXiv preprint arXiv:2102.09762, 2021.
James Spall. Multivariate stochastic approximation using a simultaneous perturbation gradient approximation. IEEE transactions automatic control, 37(3):332-344 33,-341, 1992.
James C Spall. Introduction to stochastic search and optimization: estimation, simulation, and control. John Wiley & Sons, 2005.
Hugo Touvron, Louis Martin, Kevin Stone. Peter Albert, Amjad Almahairi, Yasmine Babaci, Niko- lay Bashlykow, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open founda- tion and fine-tuned chat models. arXiv preprint arXiv:2307.09288 2023.
A Vaswani. Attention is all you need. Advances in Neural Information Processing Systems, 2017.
Shida Wang. Longssm: On the length extension of state-space models in language modelling. arXiv preprint arXiv:2406.000080 2024.
Sinong Wang, Belinda Z. Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity, 2020. URL
Guangxuan Xiao, Jiaming Tang. Jingwei Zuo, Junxian Guo, Shang Yang. Haotian Tang, Yao Fu, and Song Han. Duoattention: Efficient long-context Ilm inference with retrieval and streaming heads. arXiv preprint 2401022410.10819, 2024.
A APPENDIX
A.1 DETAILED HYPERPARAMETERS
C0 hyperparameters. For Pile, PG,19, and LongBench dataset calibration, we set the ZO opti- mization hyperparameters to 17 = 0.001, c 0.1, and K 50.
CF hyperparameters. For the passkey retrieval task, we train the models for one epoch using Adam optimizer with a learning rate of 0.1 for MambaExtend. For DeciMamba, and full fine-tuning, we use the learning rate to be le 4, suggested by the authors Kish et al. (2024). For all three cases, we use batch size of 32, gradient clipping of 1.0. weight decay of 0.1, and train on sequence length of 6144.
A.2 PRE-TRAINED MODEL CHECKPOINTS USED
The pre-tr model checkpoints of Mamba are taken from the Hugging Face model Hut
state-ssaccs//mambee 3m0
state- sateee-pppese/amaaaaaaa 1. 4b
state-spaces/mambbr 8m
A.3 MORE RESULTS
Fig. 8 shows the performance comparison of DeciMamba and MambaExtend in terms of compute, memory, and time. For DeciMamba, we use the total training time of 5 epochs, to evaluate the normalized FT time. For MambaFxtend, we use 70 for the calibration, we report the time as- sociated with the 50 iterations of calibrations. Notably, for MambaExtend, we calibrate separately for each evaluation context length, while DeciMamba does one fine-tuning for 5 epochs with 2k
13

---

Published as a conference paper at ICLR 2025
Homaltedd Memory
3.5
Mambaastten 130M
DeciMaba 130M
Hormaloed Call Puram (loe)
-2p
8
20.0
Dec Mamie 130M
107
2.0
2.5
10P
MantaBieerd
Dedikambz
Normalzzd
100
32.5
124
3.5
10)
3:0
5:0
05
ugo
2,5
00
24
46 Length
Longettiinghh
Contest Length
Figure 8: Comparison of normalized (peak memory, number of parameter updates, and
codibrattioffime-tnniit (FT) time (total)) between DeciMamba, and MambaExtend for PG-19. We
Mamba- 130M model for this evaluation.
Table 5: PPL comparison with cransorrmr-basee LLM for long-context understanding on Pile.
Model
ZK
4K
8K
16K
32K
64K
Tinyl LaMA1. 1B (2K)
4.6
62.6
426.6
1243.7
2684.6
3372.04
LLMAL IB-PI
4.6
9.56
50.34
116.47
168.84
229.46
MumbaExtend-I
7.06
6.18
500
484
5.16
5.72
context length. This causes the peak memory and fine-tuning time to increase for MambaFxtend while keeping them constant for DeciMamba. For each evaluation metric, we normalize the cor- responding value for MambaExtend at the context length under consideration. As Fig. g shows, MambaExtend requires ~5.42 * 106 x fewer parameter updates and costs up to 3.87 x lower peak memory. Additionally, MambaExtend provides up to 20.9x faster calibration opposed to the fine-tuning duration of DeciMamba.
A.4 COMPARISON WITH TRANSFORMER- BASED LLMS
Supporting longer contexts during inference is critical challenge for both transformer-basse and Mamba-band LLMs. In this section, we compare performance with transformer- based models to provide a broader perspective on the long-context extension results achieved with Mamba models. Specifically, we use the TinyLLaMA-I.I model, which is trained on 2K context length, as our baseline transformer model. We also consider a positionally interpolated version of the same model. It is important to note that positional interpolation (PI) is widely used training-free method for extending the con- text transformer models. As shown in Ta the ManbaExtend model, despite being smaller, consistently outperforms TinyLLaMA-1.11 at longer context lengths, both with and without PI. Furthermore, we compare the performance of TinyLLaMA-1111 (with and without PI) and Mam- baExtend on PG19, another popular benchmark for perplexity (PPL) evaluation long contexts. The results, presented in Table [6] clearly demonstrate the significant performance advantages of MambaExtend compared to transformer-base alternatives. It is noteworthy that MambaExtend sig- nificantly surpasses the TinALLaMA model variants, whether they are smaller or of similar size, highlighting its efficacy.
Table 6: PPL comparison with transformer-based LLM for long-context understanding on PG19.
Model
16K
32K
64K
TinyLLaMAI, I.B (2K)
2236.98
4205.64
8664. I
TinyLLaMAI. IB-PI
226.69
300.46
375.49
MambaExtend-
19.25
20.3
25
MambaExtend--LAkk
14
14.34
16.12
A.5 MORE COMPARISON WITH DECIMAMBA
While in the main manuscript, demonstrate the benefits of MambaExtend over the baseline Mamba on Pile dataset, we now show a comparison with DeciMamba (Ben- Kish et al. 20244 on the same. In specific, Table 7 demonstratss the efficacy of MambaExtend in maintaining the PPL better than DeciMamba, particularly at longer contexts with context length 2 SK. Additionally, we show
Table 7: PPI. comparison between DeciMamba and MambaFxxtend on Pile.
Model
2K
4K
8K
16K
32K
64K
DeciMamba- 130M
5.36
5.21
6.99
8.19
10.62
MambaExtend- 130M
7.06
6.18
5.03
5.16
5.72
results LongBench to compare with that generated by DeciMamba in zero-shot fashion. In
14

---

Published as a conference paper at ICLR 2025
Table Fl scores on HotpotQA and Qasper from LongBench on DeciMamba and MambaExtend, respectively. Italicized numbers identify the results taken from (Authors LongM
2024 paper.
Model
HotpotQA
Qasper
DeciMamha- 1.48
13.88
14.24
MambaExtend-Lan
14.29
16.67
Marbo
2000
5000
4000
6000
6000
4000
5000
3000
2000
2000
2000
1000
0000
1000
1000
5000
8000
20 Layer
30
Layer
16)
Figure 9: Impact of the calibrated scaling factors on A, (a) Mamba vs. MambaExtend, and (b) Mamba vs. DeciMamba evaluated Pile 32K context length. (Top) of both (a) and (b) shows layer-wise Normalized sum A for pre-trained Mamba-1 1.4bb (Bottom) layer-wise Normalized sum of A, for (a) MambaEXtend--1. 1.44 calibrated model, and (b) DeciMamba- 1.4B fine-tuned model. To fine-tune DeciMamba 1.4B model we adhered to the setup described in (Ben-Kish et al. 2024).
specific, Table shows that MannaExtend can yield reasonably improved performance as evaluated on HotpotQA and Qasper, respectively.
Comparing the impact learned scaling and full fine-tuning on A., MambaExtend employs learned scaling policy to adjust the discretization steps, A,. In contrast, DeciMamba (Ben-Kish al., 2024) performs full model fine-tuning enhance performance on longer contexts. We now visualize the impact of these two approaches on the normalized sum of A per layer. Specifically, in
Fig. 2 we present direct comparison of MambaExtend a)) and DeciMamba (2))). Interestingly. both approaches exhibit similar influence the normalized sum of A,'s, ss significantly reducing
its values the later layers. This experiment demonstrates that both methods effectively recalibrate the A, ss However, approach achieves similar benefits with greater computational and memory efficiency and reduced latency.
55
Lase
3.5
Scaling Facker of
Figure 10: Impact of different values of uniform A, scaling on the loss landscape of the model.
A.6
THE Loss LANDSCAPE FOR GRID-SEARCHED SCALING FACTORS
Fig. in the main manuscript demonstrates the impact of uniform A, scaling per layer terms of PPL value. We now plot the loss landscape of the model with uniform scaling factor values in the same range as that of Fig. specific, [00 shows the loss landscape to have a convex nature as we sweep over the scale factors (s) in
15

---

Published as a conference paper at ICLR 2025
Table PPL comparison with TBTT Wang 2024 fine-tuned model on Pile.
Model
TBTT
line-tuned
4K
8K
16k
32K
64K
Mambaz--SSh (baseline)
No
4.62
22.
19
183
378
Mambu2 780M
Yes
4.62
4.34
3,89
4.92
5.16
Mamiba2Extend-8900
No
3,95
3,89
4,25
5,56
5
Table 10: Comparison with TBTT (Wagg 2024
fine-tuned model on Passkey retrieval task.
Model
TBTT line-tuned
Arg. Accuracy
Mambaz- 780M (basclinc)
No
Mamhu2-7800
Mamba Excend-SM
91.34
Table 11: Comparison between fine-tuning and calibration for longer epochs on Passkey retrieval.
Model
Passkey retrieval acc. (%)
DeciMamba- 130M
93.1
MambalExtend 1300
94.3
A.7 COMPARISON WITH MODELS THROUGH TIME
VIA TRUNCATED BACKPROPAGATION
Contemporary research on Mamba2 models trained using truncated backpropagation through time (TBTT) has demonstrated promising generalization capabilities on longer contexts (Wang 2024 To compare MambaEXtend with the TBTT fine-tuned model, we apply fine-tuning for three epochs based on the TBTT approach on a pre-trained Mamba2-780M using 0.8B tokens from the PG19 train split. We then evaluate performance on the Pile and Passkey retrieval tasks and present comparisons with MambaExtend in Tables 9 and 10 respectively. Interestingly, we observe a no- table performance improvement for the Pile dataset for longer contexts, nearing the performance of MambaExtend. However, in the critical benchmark of long-contex: retrieval (Table [0), the TBTT fine-tuned model exhibits negligible retrieval accuracy. In contrast, MambaExtend achieves a sig- nificant accuracy boost solely by calibrating the scaling factors.
Important Notes on TBTT Training. The TBTT-based fine- tne-tuning method shares similarities with the DeciMamba approach (Ben- et al.] 2024), which also proposes full fine-tuning to enhance long-context understanding (albeit without TBTT). However, we emphasize that the key advantage of scaling-based calibration in MambaExtend is that it is an orthogonal method compared to such full fine-tuning approaches. This not only leads improved accuracy but also offers significant computational and memory benefits, potentially enabling calibration in resource-constraine envi- ronments. Furthermore, as illustrated in Figure 7 of the original LongSSM paper (Wang) 2024), training relatively large models, such the 140M S5, with previously initialized state (under the TBTT policy) can lead to severe stability issues. This raises concerns about the scalability of such approaches, identified by the authors.
A.8 FINE-TUNING VS. CALIBRATION FOR LONGER EPOCHS
Table II shows results of fine-tuning with DeciMamba for five epochs on passkey retrieval. For fair comparison, we show the results of MambaExtend with scaling factors calibrated for the same epochs. As we can see, MambaExtend can still retain improved performance over the other. However, please note, in this work we aim to achieve long context generalization with minimal compute and calibration overhead, thus aim to focus fine-tuning for only epoch.
A.9 HARDWARE AND API RESOURCES USED
We used an Nvidia A6000 GPU with 48 GB memory for all the experiments. To perform calibration and fine-tuning finettunig, we used Pytorch API to write the corresponding code.
16


Reviews:
[{"review_id": "g4D0ARAxx5", "reviewer": "Reviewer_jd9z", "summary": "This paper introduces MambaExtend, a novel framework designed to enhance the long-context capabilities of Mamba models, a leading type of state-space model (SSM) known for achieving impressive speed in language mode traininig/inference. this paper identifies the limitations of Mamba in generalizing to longer contexts due to out-of-distribution discretization steps and proposes a training-free approach to address this challenge. The authors have shown that their training-free method successfully increases the context length of Mamba models while preserving computing efficiency. This paper's primary finding indicates significant enhancements in perplexity scores on the Pile and PG-19 datasets, achieving increases of up to 8145x and 32506x, respectively.", "questions": "see weaknesses above", "limitations": "", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "* The authors empirically found that scaling down the discretization step size (Δt) can improve generalization on increased context lengths at inference time.\n\n* This paper enhances Mamba's context length without retraining the model weights by leveraging a calibration function (CF) to optimize the discretization step sizes (Δt) across various Mamba layers, which significantly reduces the number of parameter updates and peak memory demand compared to traditional fine-tuning methods", "weaknesses": "* Limited Baseline Comparisons The study lacks comparisons with non-Mamba architectures like transformers on Pile and LongBench datasets; Including established techniques like Positional Interpolation would provide valuable context about relative performance improvements; A broader comparison framework would better highlight MambaExtend's contributions within the larger landscape of long-context models\n\n* Insufficient Coverage of Related Work: Lack of discussion of non-mamba techniques that addressing the long context LLM techniques, such as  NTK-Aware scaling, LongLora. A discussion of how these existing approaches compare or could complement MambaExtend would strengthen the paper's positioning\n\n* Deeper Explanation of Discretization and Its Role in OOD: While the paper states that OOD discretization steps limit Mamba's long-context performance, it could benefit from a more thorough explanation of what discretization is, how it specifically relates to Mamba's S6 layers, and why it becomes problematic for longer sequences.  the paper mentions that scaling down Δt improves long-context generalization but doesn't fully explain the underlying rationale. While the paper demonstrates that reducing Δt improves performance, the theoretical underpinning and connection to effective receptive field (ERF) warrant deeper exploration.\n\n* How to choice of CF is non clear: The rationale for choosing between CFZO and CFBP calibration functions lacks clarity, The paper mentions using CFZO for simpler tasks and CFBP for more complex ones due to the number of scaling factors involved. This could be confused, would you elaborating more detail, guideline of criteria of selecting the calibration function?", "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": "e4WTaOSLpw", "reviewer": "Reviewer_SxC7", "summary": "This paper introduces MambaExtend, an approach to improve the long-context capabilities of Mamba models.\nKey contributions are:\n\nThe authors leverage the out-of-distribution discretization steps to improve Mamba's poor generalization to longer contexts. They show that scaled-down discretization steps can improve performance on longer sequences.\nThey propose MambaExtend, which introduces learnable scaling factors for the discretization steps (Δt) in each layer while keeping the pre-trained model weights frozen. The framework includes both gradient-based and gradient-free (zeroth-order) optimization approaches to learn these scaling factors.\nWith their method, they enable context extension up to 32× (from 2k to 64k tokens) without significant perplexity degradation\nand achieve up to 40.6% reduced perplexity compared to alternative approaches.", "questions": "Figure 3 indicites, as its also quite intutitive, that line searching over the different scaling factors would be sufficient - althought potentially costly. Generally, it would be interesting to analyse the loss landscape wrt to the discretization steps in general and confirm its convex nature. \n\nCan you highlight more the differences to DeciMamba, and work on showing qualitative / quantiative results. Afaiu DeciMamba tunes all weights, and therefore induces more memoery and compute costs? Does GD on all weights nevertheless tunes (primarily) the discretization steps, nevertheless? This could be analysed quite easily imo and would highlight your contribution, showing that GD on all parameters anyway tunes what you propose to tune.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The paper is quite well writtin and easy to follow. The authors propose a simple approach to scale the out-of-context generalization of Mamba by tuning the discretization steps of the architecture. The method is intutitve and quite well analysed.", "weaknesses": "I find the term \"training-free\" quite misleading. Afaiu the method does train the discretization steps, with backprop or approximation. Please clarify what you mean with this. \n\nI find it also hard to judge what is novel in this work compared to the (recent) DeciMamba, in terms of insights, as they already, to the best of my knowledge, identify the discretization steps to be crucial for solving longer context generalization. Generally, I do not find these insights surprising.", "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": "5Zm9ckGNWr", "reviewer": "Reviewer_rYkc", "summary": "This paper aims to address the length extrapolation problem in Mamba. The authors are inspired by the $\\Delta_t$ discussion in DeciMamba and further investigate the role of $\\Delta_t$, finding that the longer sequence will have stronger forgetting for the previous sequence. Thus, they propose an approach to adjust the $\\Delta_t$ for length extrapolation. The proposed method is smooth to follow and the 2 different optimization methods provide more computationally efficient solutions.", "questions": "See Weakness", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Originality: this paper further investigates the role of $\\Delta_t$ in Mamba length extrapolation, and proposes an easy but novel algorithm to learn the scaling of $\\Delta_t$ using both first and second-order methods. Provided a computationally efficient and general method for the targeted problem.\n2. Quality and Clarity: writing is easy to follow, and experiments are comprehensive and clear.\n3. Significance: this work is like PI in a normal transformer and should be a significant contribution.", "weaknesses": "1. Experiments are not fully compared with existing work (e.g. DeciMamba). In Tables 1 and 2, the authors only compare with the vanilla Mamba.\n2. Figure 5: the passkey experiment, is different from the DeciMamba[1] Figure 8 and LongMamba[2] Figure 3, where DeciMamba performs better compared to Fig 5's result. Could authors explain the possible reason for the difference? \n\n[1] DeciMamba: Exploring the Length Extrapolation Potential of Mamba\n\n[2] LONGMAMBA: ENHANCING MAMBA’S LONGCONTEXT CAPABILITIES VIA TRAINING-FREE RECEPTIVE FIELD ENLARGEMENT", "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": "0psmGzQogN", "reviewer": "Reviewer_UZhN", "summary": "The paper proposes MambaExtend a computationally efficient way to improve Mamba's ability handle longer contexts. \nThe method is based on the observation that when presented with longer sequences than during training, the time discretization parameter of Mamba assumes values which the model cannot handle. \nAs a result, the authors propose to use a similar approach to position interpolation in transformers and scale the discretization parameter per layer independently to a range which Mamba has seen during training. \nThe authors propose two methods to estimate the scaling factor for the discretization parameter of each layer a gradient-based and a zero-order optimization method using finite-differences to estimate the scaling factor and ensure that the scaling factor can only be positive.\nThe authors demonstrate improved performance on pass-key retrieval tasks and show significantly lower perplexity at high context lengths.", "questions": "- Please see the main weakness. In a sense, I think it would be important to understand whether the performance degradation you observe for Mamba is an artifact of the training using a transformer pipeline where the data is shuffled into non-consecutive chunks compared to an RNN training pipeline where the data isn't chunked and the last hidden state of the previous batch is the first hidden state of the current batch.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The proposed method allows Mamba to be used with much higher context lengths than seen during training. Given that long-context handling is the computational advantage of Mamba models, it's great that the proposed method allows to alleviate performance degradation.\n- The results on pass-key retrieval are convincing. \n- The method is well-motivated and easy to understand.\n- The authors propose two methods: a more accurate but also more memory heavy gradient based method and a less accurate but less memory intensive zero-order method which makes sense to me as they are applicable depending on the model size and hardware constraints.", "weaknesses": "- While the proposed method is effective at making an existing checkpoint of Mamba more capable of handling long contexts, I feel like it treats more a symptom rather than a cause. As shown in Wang et al. (2024) (Figure 4) Mamba can handle long context effectively if trained using truncated backpropagation through time on consecutive chunks of the data (rather than shuffling). Crucially, the initial state of Mamba is changed from zero to the last hidden state in the sequence. As a result Mamba doesn't need to start from scratch at token position i every time. \nThis is challenging to implement in Mamba 1 as the initial state can't be set but is an exposed option for Mamba 2. I think it would be interesting to understand whether the observed differences in the discretization parameter for longer context still occur in models trained in this way or whether it makes the proposed method less effective. This could be checked experimentally using an existing small checkpoint (340m) and then fine-tuning on a small dataset of around 1B token using TBTT.\n- The code for reproducibility was not provided.\n\n\n\n\nReferences:\n- Wang, S. (2024). LongSSM: On the Length Extension of State-space Models in Language Modelling. arXiv preprint arXiv:2406.02080.", "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)