Samba: Simple Hybrid State Space Models for Efficient Unlimited Context Language Modeling
Authors: Liliang Ren
Yang Liu
Yadong Lu
yelong shen
Chen Liang
Weizhu Chen
Venue: ICLR 2025 Poster (2025)
Keywords: Large Language Models;Length Extrapolation;Efficiency;Hybrid State Space Models
Primary area: foundation or frontier models, including LLMs
TL;DR: Simple hybrid state space models outperform SOTA Transformers and SSMs.

Abstract: Efficiently modeling sequences with infinite context length has long been a challenging problem. Previous approaches have either suffered from quadratic computational complexity or limited extrapolation ability in length generalization. In this
work, we present Samba, a simple hybrid architecture that layer-wise combines
Mamba, a selective State Space Model (SSM), with Sliding Window Attention
(SWA). Samba selectively compresses a given sequence into recurrent hidden
states while still maintaining the ability to precisely recall recent memories with the
attention mechanism. We scale Samba up to 3.8B parameters with 3.2T training
tokens and demonstrate that it significantly outperforms state-of-the-art models
across a variety of benchmarks. Pretrained on sequences of 4K length, Samba
shows improved perplexity in context lengths of up to 1M in zero-shot. When
finetuned on 4K-length sequences, Samba efficiently extrapolates to a 256K context length with perfect memory recall on the Passkey Retrieval task, and exhibits
superior retrieval extrapolation on the challenging Phonebook task compared to
full-attention models. As a linear-time sequence model, Samba achieves a 3.73×
higher throughput compared to Transformers with grouped-query attention for user
prompts of 128K length, and a 3.64× speedup when generating 64K tokens with
unlimited streaming.

Published as a conference paper at ICLR 2025
SAMBA: SIMPLE HYBRID STATE SPACE MODELS FOR EFFICIENT UNLIMITED CONTEXT LANGUAGE MODELING
Liliang Ren-.2* Yang Liu
Yadong Lu-1
Yelong Shen' Chen Liang'
Weizhu Chen 2
1 Microsoft University of Illinois at Urbana-Champaign
(liliangren, yaliu10, yadonglu, yelong. .eeloonn.s,n,,n...e whenliange wwwchen)mmicroooft.ccooo
ABSTRACT
Efficiently modeling sequences with infinite context length has long been chal- lenging problem. Previous approaches have either suffered from quadratic compu- tational complexity or limited extrapolation ability in length generalization. In this work, we present SAMBA, a simple hybrid architecture that layer-wise combines Mamba, selective State Space Model (SSM), with Sliding Window Attention (SWA). SAMBA selectively compresses a given sequence into recurrent hidden states while still maintaining the ability precisely recall recent memories with the attention mechanism. We scale SAMBA up to 3.8B parameters with 3.2T training tokens and demonstrate that it significantly outperforms state-of-the-art models across a variety of benchmarks. Pretrained on sequences of 4K length, SAMBA shows improved perplexity in context lengths of up to 1M in zero-shot. When finetuned on -K-length sequences, SAMBA efficiently extrapolates to a 256K con- text length with perfect memory recall on the Passkey Retrieval task, and exhibits superior retrieval extrapolation on the challenging Phonebook task compared to full-attention models. As linear-time sequence model, SAMBA achieves a 3.73x higher throughput compared to Transformers with grouped-query attention for user prompts of 128K length, and a 3.64X speedup when generating 64K tokens with unlimited streaming. Our code for training open source data is publicly available at https: //gi thub. www. wwm.cooooot/aaaaaa......
INTRODUCTION
Attention-basse models (Vaswani et al., 2017: Bahdanau et al., 2014) have dominated the neural architectures of Large Language Models (LLMs) (Radford et al., 2019; Brown et al., 2020; OpenAI, 2023; Bubeck et al., 2023) due to their ability to capture complex long-term dependencies and the efficient parallelization for large-scale training (Dao et al., 2022a). Recently, State Space Models (SSMs) (Gu et al., 2021: Smith et al., 2023; Gu et al., 2022: Gu & Dao, 2023) have emerged as promising alternative, offering linear computation complexity and the potential for better extrapolation to longer sequences than seen during training. Specifically, Mamba (Gu & Dao, 2023), variant of SSMs equipped with selective state spaces, has demonstrated notable promise through strong empirical performance and efficient hardware-aware implementation. Recent work also shows that transformers have poorer modeling capacities than input-dependent SSMs in state tracking problems (Merrill et al., 2024). However, SSMs struggle with memory recall due to their recurrent nature (Arora et al., 2023), and experimental results on information retrieval-related tasks (Fu et al., 2023; Wen et al., 2024; Arora et al., 2024), have further shown that SSMs are not competitive as their attention- based counterparts.
Previous works (Zuo et al., 2022; Fu et al., 2023; Ma et al., 2023; Ren et al., 2023) have explored various approaches to hybridize SSMs with the attention mechanism, but none have demonstrated significantly better language modeling performance compared to state-of-the-art Transformer ar- chitectures. Existing length extrapolation techniques (Han et al., 2023; Xiao et al., 2023; Jin et al.,
Work partially done during internship at Microsoft.
I Equal second-author contribution.

---

Published as a conference paper at ICLR 2025
2024) designed for attention mechanisms constrained by quadratic computational complexity or insufficient context extrapolation performance, particularly when evaluated under perplexity metrics. In this paper, introduce SAMBA, a simple neural architecture that harmonizes the strengths of both the SSM and the attention-based models, while achieving a potentially infinite length extrapolation with linear time complexity. SAMBA combines SSMs with attention through layer-wise interleaving Mamba (Gu & Dao, 2023), SwiGLU (Shazeer, 2020), and Sliding Window Attention (SWA) (Beltagy al., 2020). Mamba layers capture the time-dependent semantics and provide backbone for efficient decoding, while SWA fills in the gap modeling complex, non-recurrent dependencies. A detailed discussion of related work is included in Appendix A.
We scale SAMBA with 421M, 1.3B, 1.7B and up to 3.8B parameters with 3.2T tokens. In particular, the largest 3.8B post-trained model achieves a 71.9 score for MMI.U (Hendrycks et al., 2021), 62.8 for HumanEval (Chen et al., 2021), and 87.6 for GSM8K (Cobbe et al., 2021), substantially outperforming the post-trained Phi-3-mini model under a control of the same training recipes and datasets, as detailed in Table 1. Despite being pre-trained in the 4K sequence length, SAMBA can be extrapolated to IM length in zero shot with improved perplexity on Proof-Pile (Zhangir Azerbayev & Piotrowski, 2022), achieving 256x extrapolation ratio, while still maintaining the linear decoding time complexity with unlimited token streaming, as shown in Figure 2. We show that when instruction tunde in a 4K context length with only 500 steps, SAMBA be extrapolated perfect memory recall in Passkey Retrieval (Mohtashami & Jaggi, 2023). In contrast, the fine-tuned SWA-based model simply cannot recall memories beyond 4K length. We further demonstrate that the instruction-tune tuned SAMBA 3.8B model can achieve significantly better performance than the SWA- based models on downstream long-context summarization tasks, while still keeping its impressive performance on the short-context benchmarks. In a more challenging multiple key- key-value retrieval task, Phonebook (Jelassi al., 2024), we demonstrate that instruction one-tuming enables SAMBA to bridge the retrieval performance gap with full-attention models, while exhibiting significantly better extrapolation ability when retrieving phone numbers beyond the training context length. Finally, we perform extensive amalyzes and ablation studies across model sizes up 1.7B parameters to validate the architectural design of SAMBA. We also offer potential explanations for the effectiveness of our simple hybrid approach through the lens of attention/selecciion entropy. To the best of our knowledge, Samba is the first hybrid model showing that linear complexity models can be substantially better than state-of-ther oftthe-arttttt Transformer models on short-context tasks at large scale, while still being able to extrapolate to extremely long sequences under the perplexity metric.
256K context length with
2 METHODOL.OGY
We explore different hybridization strategies consisting of the layers of Mamba, Sliding Window Attention (SWA), and Multi-Layer Perceptron (Shazeer, 2020; Dauphin et al., 2016). We concep- tualize the functionality of Mamba as the capture of recurrent sequence structures, SWA as the precise retrieval of memory, and MLP as the recall of factual knowledge. We also explore other linear recurrent layers including Multi-Scale Retention (Sun et al., 2023) and GLA (Yang et al., 2023) as potential substitutions for Mamba in Section 3.2. Our goal of hybridization is to harmonize between these distinct functioning blocks and find an efficient architecture for language modeling with unlimited length extrapolation ability.
2.1 ARCHITECTURE
As illustrated in Figure 1, we explore three kinds of layerwise hybridization strategies on the 1.7B scale: Samba, Mamba-SWA-MLLP and Mamba-MLP. We also explore other hybridization approaches with full self-attention on smaller scales in Section 4. The number of layers N is set to 48 for Samba, Mamba-MI.P and Mamba, while 1.7B parameters. We only modify the layer-level arrangement for each of the models and keep every other configuration the same to have apple-to-apple comparisons. More details on the configuration of each layer are explained in the following subsections.
has 54 layers, so each model has approximately
2. MAMBA LAYER
Mamba (Gu& & Dao, 2023) is a recently proposed SSM-based model with selective state spaces. It enables input-dependent gating to both the recurrent states and the input representation for soft
2

---

Published as a conference paper at ICLR 2025
Mamba
Mamba
Mamba
Mamba
MLP
N/3
SWA
MLP
Mamba
N44 x
N44
N/4 x
SWA
MLP
Mamba
Mambu
MLP
MLP
Mamba
Samba
Mamba-SWA-MLP
Mamba-MLP
Mamba
Figure I: From left right: Samba, Mamba-SWi SAAraa-SWA--LLLAA Mamba-MLPP and Mamba. The illustrations
depict the layer-wise integration of Mamba with various configurations of Multi-Layer Perceptrons (MLPs) and Sliding Window Attention (SWA). We assume the total number of intermediate layers to be N, and omit the embedding layers and output projections for simplicity. Pre-Norm (Xiong et al., 2020; Zhang& & Sennrich, 2019) and skip connections (He et al., 2016) are applied for each of the intermediate layers.
selection of the input sequence elements. Given an input sequence representation X R"xd.,, where is the length of the sequence and dm is the hidden size, Mamba first expands the inputs to higher dimension de, i.e.
= XWin in Ruxd,
where Win in E Rd xd. is a learnable projection matrix. Then a Short Convolution (SC) (Poli et al., 2023) operator is applied to smooth the input signal,
U SC(H) = SiLU(pppshiiseConn/H Woom)) R
where W conv E RR kxd, and the kernel size k is set to 4 for hardware-aware efficiency. The Depthwise Convolution (He et al., 2019) is applied over the sequence dimension followed by SiLU (Elfwing et al., 2017) activation function. The selective gate is then calculated through low-rank projection followed by Softplus (Zheng al., 2015),
A = SSotplus UW, + b) E Rnxxd
where W, E Rd. Wg E xd, and d, is the low-rank dimension. E RRd, carefully initialized so that A E [Amins Amax] after the initialization stage. We set [Amin Amax] = |0.001, 0.1]. and find that these values are not sensitive to language modeling performance under the perplexity metric. The input dependence is also introduced for the parameters B and C of SSM.
B = UWb E Rand,
C = UW. E Rndd,
where de is the state dimension. For each time step 1 < < n, the recurrent inference of the Selective SSM (S6) is performed in an expanded state space Z, E Rdd xd. i.e.,
Z, exp( - A O exp ()) Z A O B & U.) ERd.
Y, = -,C + DD E ER
where Zo 0, ① means the point-wise product, @ means the outer product and exp means the point-wise natural exponential function. D E Rd. is learnable vector initialized as D, 1 and A 1d, xd, is learnable matrix initialized A₁ log(s), < j dg. following the S4D- Real (Gu et al., 2022) initialization. In practice, Mamba implements hardware- scan algorithm for efficient parallelizable training. The final output is obtained through gating mechanism similar to Gated Linear Unit (Shazeer, 2020; Dauphin et al., 2016),
parallel
O Y SiLU(XW, W... E nxd.
where W E Xd. and W out E xd are learnable parameters, In this work, we set de 2dm.. d, = d /16, and d, = 16. The Mamba layer in SAMBA is expected to capture the dimedeennndent semantics of the input sequence through its recurrent structure. The input selection mechanism in the Mamba layer enables the model to focus on relevant inputs, thereby allowing the model to memorize important information in the long term.
3

---

Published as a conference paper at ICLR 2025
2.1.2 SLIDING WINDOW ATTENTION (SWA) LAYER
We include Sliding Window Attention (Beltagy et al., 2020) layers to address the limitations of Mamba layers in capturing non-recurrent dependencies in sequences. Our SWA layer operates on window size w 2048 that slides over the input sequence, ensuring that the computational complexity remains linear with respect to the sequence length. RoPE (Su et al., 2021) is applied within the sliding window, with base frequency of 10,000. By directly accessing the contents in the context window through attention, the SWA layer can retrieve high-definition signals from the middle to short-term history that cannot be clearly captured by the recurrent states of Mamba. We use FlashAttention 2 (Dao, 2023) for the efficient implementation of self-attention throughout this work. We also choose the 2048 sliding window size for efficiency consideration; FlashAttention has the same training speed as Mamba's s selective parallel scan the sequence length of 2048 based on the measurements in (Gu& & Dao, 2023).
2.1.3 MULTI-LAYER PERCEPTRON (MLP) LAYER
The MLP layers in S SAMBA serve as the architecture's 's primary mechanism for nonlinear transformation and recall of factual knowledge (Dai et al., 2022). We use SwiGLU (Shazeer, 2020) for all the models trained in this paper and denote its intermediate hidden size as dz- As shown in Figure 1, Samba applies separate MLPs for different types of information captured by Mamba and the SWA layers.
3 EXPERIMENTS AND RESULTS
We pre-train four SAMBA models with different parameter sizes, 421M 1.3B, 1.7B and 3,8B, to investigate its performance across different scales. The details of the hyperparameters for the training and architecture designs are shown in Table 12 of Appendix G. We also train other hybrid architectures as mentioned in Section 2.1, including the baseline Mamba (Gu & Dao, 2023), Llama-3 (MetaAI, 2024; Dubey et al., 2024), and Mistral (Jiang et al., 2023) architecture detailed byperparameters in Table of Appendix G. We do comprehensive downstream evaluations on wide range of benchmarks, focusing four main capabilities of the models: commonsense reasoning (ARC (Clark et al., 2018), PIQA (Bisk et al., 2020), WinoGrande (Sakaguchi et al., 2021), SIQA (Sap et al., 2019)), language understanding (HellaSwag (Zellers al., 2019), BoolQ (Clark et al., 2019), OpenbookQA (Mihaylov et al., 2018), SQuAD (Rajpurkar et al., 2016), MMLU (Hendrycks al., 2021), MMLU-Pro (Wang et al., 2024), GPQA al., 2023)), truthfulness (TruthfulQA (Lin et al., 2022)) and math and coding (GSM8K (Cobbe et al., 2021), MBPP (Austin al., 2021), HumanEval (Chen al., 2021)).
scale of around 1.7B, with
Table 1: Downstream performance comparison between Samba-3.8B-IT and Phi-3-mini-4K on both long-context and short-context tasks. We report 5-shot accuracy (averaged by category) for MMLU, 8-shot CoT (Wei et al., 2022) for GSM8K, O-shot pass@1 I for HumanEval, ROUGE-L for both GovReport and SQUALITY. + Results from the Phi-3 technical report (Abdin et al., 2024).
Model
MMLU
GSM8K
HumanEval
GovReport
SQuALITY
1it-3-mini-4K-instee $
68.8
82.5
58,5
14.4
21.6
Samba-3.8B-I1
71.9
87.6
62.8
18.9
21.2
3.1 LANGUAGE MODELING ON TEXTBOOK QUALITY DATA
We first present results from our largest 3.8B SAMBA model, trained on the same data set used by Phi3 (Abdin et al., 2024) with 3.2T tokens. We follow the same multiphase pretraining strategy as Phi3-mini, and apply both the original Phi-3-mini post-training recipe and the Phi3-mini-June-200 recipe to produce our instruction-tune tuned S SAMBA 3.8B models, i.e., Samba-3.8B-PP and Samba- 3.8B (June) respectively. We report comprehensive benchmark results of the Samba 3.8B base model and Samba-3.87 8B (June) in Appendix B. As shown in Table 1, we evaluate the downstream performance of Samba-3.8B-11 on both long long-context summarization tasks (GovReport (Huang et al., 2021), SQuALITY (Wang al., 2022)) and major short-context benchmarks (MMLU, GSM8K, HumanEval). We can see that Samba has substantially better performance than Phi-3- mini-mk--nstteee on both the short-context (MMLU, GSM8K, HumanEval) and long-context (GovReport) tasks, while
4

---

Published as a conference paper at ICLR 2025
still having the 2048 window size of its SWA layer and maintaining the linear complexity for efficient processing of long documents. Details data statistics and evaluation setup for long context tasks included in Appendix F.
Table 2:
Downstream evaluation of the architectures trained on 230B tokens of the Phi2 dataset. We report the unnormalized accuracy for multiple choice tasks. GSM8K is evaluated with 5-shot examples while other tasks are in zero-shot. Best results are bold, second best underlined.
Llama-3
Mistral
Mamba
Gamba-SWA-MLP
Mamba-MLP
SAMBA
Benchmark
1.6B
1.6B
1.8B
1.6B
1.9B
1.7B
ARC-Easy
76.85
77.02
77.99
76.68
78.91
79.25
ARC-Challenge
43.26
44.20
45.22
46.16
47.35
48.21
PIQA
76.66
75.79
77.31
76.50
78.84
77.10
WinoGrande
70.01
70.72
73.40
73.72
72.38
72.93
SIQA
51.23
52.00
53.12
55.12
54.30
53.68
HellaSwag
46.98
47.19
49.80
49.71
50.14
49.74
BoolQ
68.20
70.70
74.83
74.74
73.70
75.57
OpenbookQA
34.00
32.80
36.60
33.80
35,40
37.20
SQuAD
74.88
72.82
67.66
76.73
63.86
77.64
MMLU
43.84
43.54
45.28
47.39
43.68
48.01
TruthfulQA (MC1)
25.70
25.09
26.81
26.20
26.44
27.78
TruthfulQA (MC2)
40.35
38,80
40.66
40.80
40.04
41.62
GSM8K
32.68
32.45
32.07
44.05
27.52
38.97
MBPP
46.30
47.08
47.86
47.08
47.08
48.25
HumanEval
36.59
36.59
35.98
37.80
31.10
39.02
Average
51.17
51.12
52.31
53.77
51.38
54.33
To examine the different hybridization strategies mentioned in Section 2.1, we train 6 models with around 1.7B parameters on the Phi2 (Li et al., 2023) dataset with 230B tokens and evaluate them in the full suite of 15 downstream benchmarks have holistic assessment of hybrid and purebred architectures. As shown in Table 2, SAMBA demonstrates superior performance on a diverse set of tasks, including commonsense reasoning SQuAD), TruthfulQA and code generation (HumanEval, MBPP). It outperforms both the pure attention-based and SSM-based models in most tasks and achieves the best average performance. By comparing the performance of Mamba-MLPP and Mamba Table 2, we can observe that replacing Mamba blocks with MLPs does not harm common sense reasoning ability, but its performance in language understanding and complex reasoning ability, such as coding and mathematical reasoning, degenerates significantly. We can also see that pure Mamba models fall short on retrieval intensive tasks such SQuAD due to their lack of precise memory retrieval ability. The best results are achieved through the combination of the attention and Mamba modules, as shown with our Samba architecture. We can also notice that GSM8K, potentially resulting from closer collaboration between the Mamba and the SWA layers. The distinct downstream performances of different hy hybridization strategies pose interesting future work for developing task- adaptive dynamic architectures.
language understanding (MMLU,
has significantly better performance on
3.2 EXPL.ORATION ON HYBRIDIZING ATTENTION AND LINEAR RECURRENCE
Since SSMs belong to a broader realm of linear recurrent models (Orvieto et al., 2023; Qin et al., 2023; Yang et al., 2023; Katsch, 2023; Qin et al., 2024; Yang al., 2024), there exist multiple alternatives other than Mamba when combing attention-based layers with recurrent neural networks. We also add architecture ablation studies to justify the design choices of Samba. Specifically, in addition to Llama-2, Mamba, Samba and Hamba-SWA-Met we investigate the comparative analysis of the following architectures:
Alama-2-SWe 2-SWA is a pure attention based architecture that replaces all full attention layers in Llama-2 with sliding window attention.
5

---

Published as a conference paper at ICLR 2025
Table Perplexity on the validation set of SlimPajama for different attention and linear recurrent model architectures trained at 4,096 context length. We use window size 2,048 for Sliding Window Attention (SWA). The perplexity results have fluctuation around +..3%.
Training Speed
Validation Context Length
Architecture
Size
Layers
(x10* tokens/s)
4096
8192
16384
20B training tokens on 8xA GPUs
Llama-2
438M
4.85
11.14
47.23
249.03
Llama-2-SWA
438M
24
4.96
11.12
10.66
10.57
Mamba
432M
60
2,46
10.70
10.30
10.24
Sliding GLA
438M
24
4.94
10.43
10.00
9.92
Sliding RetNet
446M
24
4.32
10.38
9.96
9.87
Mega-S6
422M
3,26
12.63
12.25
12.25
Mamba-SWA-MLP
400M
24
4.21
10.07
9.67
9.59
MLP2-SWA-MLP
417M
24
5.08
10.95
10.50
10.41
SAMBA-NOPE
421M
24
4.48
10.11
28.97
314.78
SAMBA
421M
4.46
10.06
9.65
9.57
100B training tokens on 64 HIOO GPUs
Llama-2
1.3B
40
25.9
7.60
44.32
249.64
Llama-2-SWA
1.3B
7.60
7.37
7.21
Mamba
1.3B
48
17.8
7.47
7.26
7.15
Sliding GLA
1.2B
36
25,9
7.58
7.35
7.19
Sliding RetNet
1.4B
36
23.0
7.56
7.35
7.56
Mega-S6
1.3B
36
17.9
9.01
8.81
8.68
Mamba-SWA SWA-MLR
1.3B
23.5
7.37
7.16
7.00
MLP2-SWA-MLP
1.B
36
26.6
7.81
7.58
7.42
SAMBA-NoPE
1.3
36
25.2
7.33
20.40
326.17
SAMBA
1.3B
36
25.2
7.32
7.11
6.96
Sliding RetNet replaces Mamba layers in the Samba architecture with Multi-Scale Retention (Sun et al., 2023) layers. RetNet is a linear attention model with fixed and input-independent decay applying to the recurrent hidden states.
Sliding GLA replaces Mamba layers in the Samba architecture with Gated Linear Attention (GLA) (Yang et al., 2023). GLA is more expressive variant of linear attention with input-dependent gating.
Mega-S6 replaces replaces all MD-EMA modules in the Mega (Ma et al., 2023) architecture with the Shor combinations from Mamba to adapt Mega to the modern Mamba architecture. Rotary position embedding, RMSNorm and Softmax attention intermediate dimension of the Mega-S6 layer to be dm so that it has a roughly 5df. number of parameters. This represents classical baseline that conducts sequential intra-layer SSM-Attention hybridization.
also adopted. We set the
MLP2-SWA-MLP replaces all Mamba layers in the Samba architecture to SwiGLU layers with 6df, number of parameters.
Samba-NoPE removes the rotary relative position embedding in Samba and does not have any position embedding in the architecture.
We pre-train all models the same SlimPajama (Soboleva al., 2023) dataset under both around 438M and 1.3B settings, and evaluate these models by calculating perplexity on the validation set with context length 4096, 8192, and 16384 tokens to investigate their zero-shot length extrapolation ability. Peak training throughput is also measured as an efficiency metric. The details of the hyperparameter settings are included Appendix G. As shown Table 3, SAMBA consistently outperforms all other models in different context lengths and model sizes. The training speed of SAMBA competitive compared to pure Transformer- based models on the 1.3B scale. Mamba has significantly worse training throughput because Mamba layers have slower training speed than MLP layers, and the purebred Mamba models need to have more layers than other models at the same number of parameters. Comparing Mamba- Seriba---W--M SWA-MLP with Samba, slightly better perplexity scores and higher training throughput. Mamba-SWe SSA-a--SA-MLLL trades off the MI.P layers with more I/O intensive Mamba and Attention layers, leading slower training speed.
that Samba has

---

Published as a conference paper at ICLR 2025
This also indicates that Mamba-SWA-MLL will have slower decoding speed than Samba due to larger total cache size resulting from more SSMs and Attention layers. We can further observe that replacing Mamba with MLP speeds up the training but harms perplexity significantly, indicating the importance of Mamba layers in the Samba architecture. Interestingly, even though we use SWA in Samba architecture, Samba-NoPE still has exploded perplexities beyond its training length without RoPE. We can also find that while RetNet can extrapolate well under the 438M scale, it has an increasing perplexity on 16K length the 1.4B scale, which may indicate that its input-independent decay may need specific tuning at different scales to work well.
Table 4: Downstream evaluation of models pre-trained with 100B tokens from SlimPajama. We measure the Character-normalize accuracy for HellaSwag following Gu & Dao (2023). All tasks are evaluated zero-shot.
ARC-Easy acc f
HellaSwag acc_norm
Wino, acc f
PIQA acc f
LAMBADA
Architecture
Size
Avg.
acc 1
LLaMA-2
1.3B
55.09
$2.32
53.35
71.11
48.52
56.08
LLaMA-2-SWA
1.3B
56.65
52.59
54.93
71.60
47.56
$6.67
Sliding GLA
1.2B
56.94
52.52
56.75
71.38
48.17
57.15
Sliding RetNet
1.4B
57.66
52.64
56.75
71.33
48.34
57.34
Mega-S6
1.3B
50.63
41.91
52.96
68.17
37.88
50.31
Mamba
1.3B
58.08
54.93
53.99
71.98
45.97
56,99
Mamba-SWA-MI SWA-MIP
1.3B
59.64
54.50
55.25
72.42
49.12
58.19
MIPP-SWA-MILP
1.3B
55.18
50.32
52.80
70.67
48.11
55.42
SAMBA-NOPE
1.3B
58.38
54.62
56.51
72.03
51.08
58.52
SAMBA
1.3B
58.21
54.73
55.72
72.36
51.68
58.54
In Table 4, we evaluate all our 1.3B scale models on five typical commonsense reasoning tasks (ARC-Easy, HellaSwag, WinoGrande, PIQA and the OpenAl variant' of LAMBADA (Paperno et al., 2016) ) to understand the effect of architecture designs on downstream performances. We can see that Samba has the best average accuracy, outperforming the LLaMA 2 architectures by large margin. Similar to our perplexity evaluation, Samba and Samba-NoPE have similar average accuracies, whereas Mamba-SWA-MLJ falls slightly behind. We observe that different architectures excel at different tasks. Srabba-SWA-MLL.. performs best on ERC-Easy, while Samba and Samba-NoPE achieve superior results LAMBADA. Hybrid models based Mamba generally outperform hybrid linear attention models and pure softmax-attention models on HellaSwag.
3.3
EFFICIENT LENGTH EXTRAPOLATION
1.38
Lama
400
5.00
Hetrl
3.64*
Portoss
Talenos
2.50
Models
Samtta 1.73
4.00
Mamba 1.00
amss 1.60
SF- Lama-1 1.6
00
Migrat 1.50
256 526 26
45
6ak 1204, 2008 312K 34
2k 46
16
16C
17%
601
128K
Generation Langth
(a) Perplexity on the test set Proof-Pile
(b) Decoding throughput with batch size
Figure SAMBA shows improved prediction up IM tokens in the Proof-Pile test set while achieving a 3.64x faster decoding throughput than the Llama-3 architecture on 64K generation length. We also include an SE-Llama-3 1.6B baseline which applies the SelfExtend (Jin al., 2024) approach for zero-shot length extrapolation. All models are trained with 4K sequence length.
We use the test split of the Proof-Pile (Zhangir Azerbayev & Piotrowski, 2022) dataset to evaluate the length extrapolation ability our models at a scale of around 1.7B parameters. We follow Position

---

Published as a conference paper at ICLR 2025
Interpolation (Chen et al., 2023a) for data pre-processing. The sliding window approach (Press et al., 2021) is used for the perplexity evaluation with a window size of 4096. Besides having the decoding throughput in Figure 2 for the generation efficiency metric, also measure the prompt processing speed in Figure 6 of Appendix B for the models SAMBA 1.7B, Mistral 1.6B, Mamba 1.8B, Llama-3 1.6B and its Self-Extender (Jin al., 2024) version SE-Llama-3 1.6B with the prompt length sweeping from IK to 128K. We set the group size to 4 and the neighborhood window to 1024 for Self-Extension. We fix the total processing tokens per measurement be 128K and varying the batch size accordingly. The throughput is measured on a single A100 GPU with the precision of 66loat16. We repeat the measurements 10 times and report the averaged results. We can see that Samba achieves 3. x higher throughput in prompt processing compared to Llama-3 1.6B at the 128K prompt length, and the processing time remains linear with respect to the sequence length. We can also observe that the existing zero-shot length extrapolation technique introduces significant inference latency overhead on the full-attention counterpart, while it still cannot extrapolate infinitely with perplexity performance comparable to that of Samba. In Figure 2, we can also see that Mamba has a slowly and stably increasing perplexity up to 1M sequence length, which indicates that linear recurrent models can still not extrapolate infinitely if the context length is extremely large.
3.4 LONG-CONTEXT UNDERSTANDING
1.0
0.9
0.8
0.7-
0.6
0.5
280
0.4
800
0.3
35 11
0.2
0.1
Figure 3: Passkey Retrieval performance up to 256K context length for SAMBA 1.7B (Left) vs. Mistral 1.6B (right) instruction tuned on 4K sequence length with 500 steps.
0.0
1000
2000 3000 4000 5000 5000 7000 8000 Sequence Length
Figure Phonebook evaluation accu- racy of different base models.
Beyond its efficiency in processing long context, Samba can also extrapolate its memory recall ability to 256K context length through supervised fine-tuning, and still keeps its linear computation complexity. We fine-tum Samba 1.7B on Passkey Retrieval with a 4K training sequence length for only 500 steps. As presented in Figure 3, SAMBA 1.7B demonstrates a remarkable ability to recall information from significantly longer contexts compared to Mistral 1.6B, a model based solely on Sliding Window Attention (SWA). This capability is particularly evident in the heatmap, where SAMBA maintains the perfect retrieval performance across a wider range of pass-key positions in a long document of up to 256K length. We also draw the training loss curve and the overall passkey retrieval accuracy across the fine-tuming procedure in Figure 7 and Figure 8 of Appendix C. We find that despite the fact that both architectures can reach near-zero training loss in less than 250 steps, Samba can achieve near-pertee pear-perfect retrieval early at 150 training steps, while the Mistral architecture struggles at around 30% accuracy throughout the training process. This shows that Samba can have better long-range retrieval ability than SWA due to the input selection mechanism introduced by the Mamba layers. In Figure 8, we can also notice that the pre-trained base Samba model has retrieval accuracy (at step 0) similar to that of Mistral, highlighting the need for future work to improve Samba's zero-shree shot retrieval capabilities.
The encouraging results on Passkey Retrieval drives us to further explore the limits of our finetuning approach. We perform instruction tuning to the Samba-3.8B base model on Phonebook (Jelassi et al., 2024) with only 100 steps on 4K sequence length and evaluate the resulting Samba-3.8B-F7 model for a sequence length up to 8K. The evaluation setting requires the models to retrieve a random phone number from phone book containing 20 (length 400) to 480 (length 8400) name-number pairs, resulting in a pressure test of memorization to Samba which has a constant memory state size. Surprisingly, as shown in Figure 4, can see that the Samba-3.8B-PP model can close most of its gap with a attenttnoio model (Llama2 7B) that has twice the parameter size within the 4K training length, and achieves much better extrapolation accuracy compared to all other models including
8

---

Published as a conference paper at ICLR 2025
the Phi3 base model which also uses 2K sliding window attention. Since both Passkey Retrieval and Phonebook require models to remember numbers in a long context document, it is interesting one task transfer its ability to the other task in
to investigate if model instruction-tund zero-shot. We directly evaluate the Passkey Retrieval finetuned Samba 1.7B and Mistral 1.6B models (named Samba 1.7B PK-FT and Mistral 1.6B PK-FT respectively) on the Phonebook task. As shown in Figure 4, Samba 1.7B has slightly better retrieval accuracy than Mistral 1.6B, but both models cannot generalize their number recall ability beyond its sliding window size. We leave it for future work to further explore the transferability of long-context capabilities in linear complexity models.
4 ANALYSIS
In this section, we analyze the experimental results of SAMBA by answering the following research questions. The perplexity results on SlimPajama have a fluctuation around +..3% Training speed is measured on 8 x H100 GPUs by default. All the models in this section are trained on SlimPajama with 20B tokens and 4K sequence length, unless otherwise specified. We also have additional analyses on the training of SWA-based models and the effectiveness of short convolution in Appendix D.
Why not hybridize with full attention? Some previous works (Fu et al., 2023; Lieber et al., 2024) suggest a hybrid architecture of Mamba with full attention. However, as shown in Table 5, the extrapolation perplexity is exploding at a context length of 16K even if a single full attention layer is placed at the beginning of the model. Although hybridization with full attention in the second and middle sixth blocks (the fourth row in the table). following Dao et al. (2022b), can bridge the perplexity gap between full-attention hybrids and Samba, they still cannot extrapolate beyond the training sequence lengths. Samba also has much better training throughput compared to Mamba-Mil alternatives because self-attention with the FlashAttention 2 implementation more training efficient than Mamba when the sequence length is 4096.
Table 5: Perplexity on SlimPajama of Mam architectures with full attention layers replacing Mamba layers at different block indices. We define block as two consecutive layers with Mamba/Altention layer followed by an MLP. All the models have 12 blocks in total.
Block Index of Full Attention
Training Speed
Validation Context Length
Architecture
Size
(x10 tokens/s)
4096
8192
16384
449M
11
7.78
10.29
10.53
13,66
449M
7.78
10.10
10.05
12.83
MLP
449M
7.78
10.89
10.55
10.63
443M
1,5
7.93
10.06
10.34
13.57
SAMBA 421M
SWA at odd indices
8.59
10.06
9.65
9.57
How many parameters should be allocated to Attention? Given that Mamba can already capture low-rank information in the sequences through recurrent compression, the attention layers in Samba theoretically will only need to focus on information retrieval where a small number of attention heads should suffice. In Table 6, explore the techniques of query head grouping (Ainslic et al., 2023; Shazeer, 2019), for both the Llama and Samba models. Surprisingly, both the L-aMA-2-SWA architecture and the Samba architecture show improved validation perplexity when there is only one key-value head. We conjecture that this is because small language models can be more easily optimized with fewer KV heads to pay attention the contexts. We can also see that Samba has 2x smaller optimal number of query heads than the SWA model, which confirms our hypothesis that Samba can support smaller number of attention heads.
Potential explanations on why hybrid is better? We examine the entropy of the attention distribu- tions for both the Samba 1.7B and the Mistral 1.6B models. As shown in Figure 5a, the Samba model has larger variance of the attention entropy distributed over the layer indices, with an interesting pattern that the upper and lower layers have entropy higher than the middle layers. This may indicate that the attention layers are more specialized in the Samba architecture, with the middle layers focusing on precise retrieval with low low-entropy attention, and the top and bottom layers focusing on integrating the global information through high-entropy attention. We can also see in Figure 5b that,

---

Published as a conference paper at ICLR 2025
Table Perplexity on SlimPajama of Llama-2-SWA and Samba models the 430M scales trained with different number of Query and Key- Value heads. "KV Size" means the size of Key-Value Value vectors per token and attention layer. Since grouped query attention will reduce the parameters for attention from 4d to roughly 2d,, we increase the intermediate size of MLP from 8/3dm to 3dm. 4608 to have roughly the same number of total parameters as the original models.
Query Head
Key-Value Head
Head Dim.
KV
Model Size
Training Speed
Validation Context Length
Size
(x (x10" tokens/s)
4096
8192
16384
Llama-2-SWA Architecture
128
512
419M
10.01
11.11
10.64
10.56
1
256
512
419M
9.98
11.09
10.62
10.54
12
1
128
256
414M
10.25
10.89
10.44
10.35
12
4
128
1024
428M
9.85
11.11
10.64
10.56
Samba Architecture
12
2
128
512
426M
8.55
10.09
9.68
9.60
1
256
512
426M
8.46
9.99
9.59
9,51
128
256
424M
10.07
9.58
12
4
128
1024
431M
8.57
10.02
9.62
9.55
compared to the Mamba-MLP model, Samba has a higher entropy of input selection probabilities in the middle layers. This indicates that, given the memory recalling ability of the attention layers, the Mamba layers can focus more on modeling the recurrent structure rather than performing retrieval with precise input selections. This kind specialization can be beneficial for the downstream model performance, which may explain the impressive results from the Samba architecture. Details on how entropy is calculated are included in Appendix
Mistral 1.68
Samba 1:70
Awers Entrapy
Average Entropy
6.4
6.2
Mamba-M MLP 1.96
10
15
20
Block Index
(a) Average attention entropy per decoding step
(b) Average S6 selection entropy on full sequences
Figure 5:
The average entropy of the attention mechanism and the Mamba's S6 input selection mechanism at each block of layers on 100 random samples from the GSM8K dataset.
5 CONCLUSION
In this paper, we introduce SAMBA, a simple yet powerful hybrid neural architecture designed for efficient language modeling with unlimited context length. We show that SAMBA substantially outperforms state-of-the-art pure attention-basee and SSM-based models across a wide range of benchmarks including common-sense reasoning, language understanding, mathematics and coding. Furthermore, SAMBA exhibits remarkable efficiency in processing long contexts, achieving substantial speedups in prompt processing and decoding throughput compared the state-of-the-as ofttth-artttt Transformer architecture. The architecture's 's ability to extrapolate memory recall to very long contexts (up to 256K) through minimal fine-tuning underscores its practical applicability for real-world tasks requiring extensive context understanding. This efficient long-term memorization ability is further demonstrated to be useful by our evaluations in downstream long-context summarization tasks. Our analyses also provide insight into the optimal training configurations for hybrid models and underscore the benefits of combining attention mechanisms with SSMs. We find that allocating fewer parameters to the attention mechanism while leveraging Mamba's 's strengths for capturing recurrent structures leads to more efficient and effective language modeling. Our results suggest that SAMBA is a strong neural architecture for language modeling with unlimited context length.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENT
We want to thank Shuohang Wang and Liyun Liu for helping with the training infrastructure, Mojan Javaheripi and the team for the pre-training data, Ziyi Yang, Jianwen Zhang, Junheng Hao and the team for helping with post-training. The first author also wants to thank Songlin Yang for her Triton implementation of Mamba.
REFERENCES
Marah Abdin, Sam Ade Jacobs, Ammar Ahmad Awan, Jyoti Aneja, Ahmed Awadallah, Hany Awadalla, Nguyen Bach, Amit Bahree, Arash Bakhtiari, Harkirat Behl, Alon Benhaim. Misha Bilenko, Johan Bjorck, Sébastien Bubeck, Martin Cai, Caio César Teodoro Mendes, Weizhu Chen, Vishrav Chaudhary, Parul Chopra, Allie Del Giorno, Gustavo de Rosa, Matthew Dixon, Ronen Eldan, Dan Iter, Abhishek Goswami, Suriya Gunasekar, Emman Haider, Junheng Hao, Russell J. Hewett, Jamie Huynh, Mojan Javaheripi, Xin Jin, Piero Kauffmann, Nikos Karampatziakis, Dongwoo Kim, Mahoud Khademi, Lev Kurilenko, James R. Lee, Yin Tat Lee, Yuanzhi Li. Chen Liang. Weishung Liu, Eric Lin, Zeqi Lin, Piyush Madan, Arindam Mitra, Hardik Modi, Anh Nguyen, Brandon Norick, Barun Patra, Daniel Perez-Becker, Thomas Portet, Reid Pryzant, Heyang Qin. Marko Radmilac, Corby Rosset, Sambudha Roy. Olli Saarikivi, Amin Saied, Adil Salim. Michael Santacroce. Shital Shah, Ning Shang. Hiteshi Sharma, Xia Song, Olatunji Ruwase, Xin Wang. Rachel Ward, Guanhua Wang. Philipp Witte, Michacl Wyatt, Can Xu, Jiahang Xu, Sonali Yadav, Fan Yang, Ziyi Yang, Donghan Yu, Chengruidong Zhang, Cyril Zhang. Jianwen Zhang. Li Lyna Zhang. Yi Zhang. Yunan Zhang, and Xiren Zhou, Phi-3 technical report: A highly capable language model locally your phone. arXiv preprint arXiv: 2404. 14219, 2024. URL https: http:/awwwwiagraaaae///eee www.oogg///////eeeee/o. 2224//0/.//2/...../// 14219v1.
J. Ainslie, J. Thorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebr' on, and Sumit Sanghai. Gqa: Training generalized multi-query transformer models from multi-head checkpoints. Conference on Empirical Methods in Natural Language Processing, 2023. doi: 10.48550/arXiw.200555 13245. URL https: //arxis org/abs org/abs/222 /3abs/305........ 0033351132555 13245v3.
Ekin Akyürek, Bailin Wang, Yoon Kim, and Jacob Andreas. In-context language learning: Architectures and algorithms. arXiv preprint arXiv: 2401.12973 2024. https:/tss///rxi org/ abrrg/11///2/...._m/// 2401 12973v2.
Simran Arora, Sabri Eyuboglu, Aman Timalsina, Isys Johnson, Michael Poli, James Zou, Atri Rudra, and Christopher Ré. Zoology: Mcasuring and improving recall in efficient language models. arXiv preprint arXiv: 2312.04927, 2023. URL https:://arti: org/ hts///bs//2/// 2312. 04927v1.
Simran Arora, Sabri Eyuboglu. Michael Zhang. Aman Timalsina. Silas Alberti, Dylan Zinsley, James Zou, Atri Rudra, and Christopher Ré. Simple lincar attention language models balance the recall shal--hhooughput tradcoff. arXiv preprint ar 2iv: 18668, 2024.
Jacob Austin, Augustus Odena, Maxwell Nye, Maarten Bosma, Henryk Michalewski, David Dohan, Ellen Jiang, Carrie Cai, Michael Terry Quoc Le, and Charles Sutton. Program synthesis with large language models. arXiv preprint arXiv: 2703.07732, 2021. URL https://arxiv http://axiivvorrbbb//////// ortg///////b/...... 2ttp://0bb//18.77222.11..
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. International Conference On Learning Representations, 2014. URL https://arxisss org/abs/ 1409 0473v7.
Iz Beltagy. Matthew E. Peters, and Arman Cohan. Longformer: The long- transformer. arXiv preprint arXiv: Arxiv-2004.05550, 2020. URL https::/arxiv..or wwtp:/iwwwor////see..eee abwgg///b//////// http//2004...500..2 05150v2.
Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. PIQA: reasoning about physical commonsense in natural language. In The Thirty-Fourth AAAI Conference on Artificial Intelligence, AAAI 2020, The Thirty-Second Innovative Applications of Artificial Intelligence Conference, TAAI 2020, The Tenth AAAI Symposium on Educational Advances in Artificial Intelligence, EAAI 2020, New York, NY, V34105.6233 URL
USA. February 7-12, 2020. pp. https:///oi..o org/10 1609 v34105. 6239.
AAAI Press, 2020. doi:
Aleksandar Botev, Soham De. Samuel L Smith, Anushan Fernando, George-Cristian Muraru, Ruba Haroun, Leonard Berrada, Razvan Pascanu, Pier Giuseppe Sessa, Robert Dadashi, Léonard Hussenot, Johan Ferret, Sertan Girgin, Olivier Bachem, Alek Andreev, Kathleen Kenealy, Thomas Mesnard, Cassidy Hardin, Surya Bhupatiraju, Shreya Pathak. Laurent Sifre, Morgane Rivière, Mihir Sanjay Kale, Juliette Love, Pouya Tafti, Armand Joulin, Noah Fiedel, Evan Senter, Yutian Chen, Srivatsan Srinivasan. Guillaume Desjardins, David Budden, Arnaud Doucet, Sharad Vikram, Adam Paszke, Trevor Gale, Sebastian Borgeaud, Charlic Chen, Andy Brock, Antonia Paterson, Jenny Brennan, Meg Risdal, Raj Gundluru, Nesh Devanathan, Paul Mooney,
11

---

Published as a conference paper at ICLR 2025
Nilay Chauhan, Phil Culliton. Luiz GUStavo Martins, Elisa Bandy, David Huntsperger, Glenn Cameron, Arthur Zucker, Tris Warkentin, Ludovic Peran, Minh Giang, Zoubin Ghahramani, Clément Farabct, Koray Kavukcuoglu, Demis Hassabis, Raia Hadsell, Yee Whye Tch, and Nando de Frictas. Recurrentgemma: Moving past transformers for efficient open language models. arXiv preprint arXiv: 2404.07839, 2024. URL / httgg///////// 87839v1.
Tom Brown, Benjamin Mann. Nick Ryder. Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot leamner.. Advances neural information processing systems, 33:1877-1901, 2020. URL https://arxiv.. orgg g/abs/ 2005. 14165v4.
Sébastien Bubeck, Varun Chandrasckaran, Ronen Eldan, Johannes Gehrke, Eric Horvitz, Ece Kamar, Peter Lee, Yin Tat Lee, Yuanzhi Li, Scott I Lundherg, Harsha Nori, Hamid Palangi, Marco Tulio Ribeiro, and Yi Zhang. Sparks of artificial general intelligence: Farly experiments with gpt-4. arXiv preprint arXiv: 2303. 12712, 2023. URL https: //tp://aavvvv//e//ee .rgg/ abs/2303. 12212222555
Mark Chen, Jenry Tworek. Heewoo Jun, Qiming Yuan, Henrique Ponde de Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, Alex Ray. Raul Puri, Gretchen Krueger, Michael Petrov, Hcidy Khlaaf, Girish Sastry, Pamela Mishkin, Brooke Chan, Scott Gray, Nick Ryder, Mikhail Pavlov, Alethea Power, Lukasz Kaiser, Mohammad Bavarian, Clemens Winter, Philippe Tillet, Felipe Petroski Such. Dave Cummings, Matthias Plappert, Fotios Chantzis, Elizabeth Barnes, Ariel Herbert- Voss, William Hehgen Guss, Alex Nichol, Alex Paino, Nikolas Tezak, Jie Tang. Igor Babuschkin, Suchir Balaji, Shantanu Jain, William Saunders, Christopher Hesse, Andrew N. Carr, Jan Leike, Josh Achiam, Vedant Misra, Evan Morikawa, Alec Radford, Matthew Knight, Miles Brundage, Mira Murati, Katie Mayer, Peter Welinder, Bob McGrew. Dario Amodei, Sam McCandlish, Ilya Sutskever. and Wojeiech Zaremba. Evaluating large language models trained on code. arXiv preprint arXiv: 2107.03374 2021. URL https: //arxiv.org 220/ 03374v2.
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. URL, https:
///aaxiv.orrg///// ab2//366./ 15595v2.
Yukang Chen, Shengju Qian, Haotian Tang, Xin Lai, Zhijian Liu, Song Han, and Jiaya Jia. Longlora: Efficient line-tuning of long- large language models. International Conference on Learning Representations. 2023b. doi: (0048550/arXi..2300..22 12307. URL htttps://arae g/ abs/ 2309.. 12307v1.
Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. PREPRINT, 2019. URL httpps///arxxt wrw. att////s///// 10509v1.
Christopher Clark, Kenton Lec, Ming-Wei Chang. Tom Kwiatkowski, Michael Collins, and Kristina Toutanova. Boolq: Exploring the surprising difficulty of natural yes/no questions. In Jill Burstein, Christy Doran, and Thamar Solorio (eds.). Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-HLT 2019, Minneupolis, MN, USA, June 2-7, 2019, Volume (Long and Short Papers), pp. 2924 2936. Association for Computational Linguistics. 2019. doi: 10.18653/V 111111/1//199 1300. URL https: //oii org/10 0 (//5//////19-1////
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try are, the ai2 reasoning challenge. arXiv preprint arXiv: 1803.05457, 2018. URL
/10///183... .5455v1..
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Retichiro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems. arXiv preprint arXiv: 2110.14168, 2021. URL 14166 18v2.
Damai Dai, Li Dong, Yaru Hao, Zhifang Sui, Baobao Chang, and Furu Wei. Knowledge neurons in pretrained transformers, ACL, 2022. URL https::t:///et http:///aaaavv.ooaa//////// 2104 08696v2.
Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. arXiv preprint arXiv: 2307,08691, 2023. URL. wtt.ttttss//axi....t org arg/abs//e 230.2/367......... 08691v1.
Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAttention: Fast and memory- efficient exact attention with 10-awareness. In Advances in Neural Information Processing Systems, 2022a.
Tri Dao, Daniel Y. Fu, Khaled Kamal Saab, A. Thomas, A. Rudra, and Christopher Ré. Hungry hungry hippos: Towards language modeling with state space models. International Conference On Leorning Representations, 2022b. doi: 10.48550/arXww.2212 21.2....... URL https: wwtparwwwwwwwr.aa/////aaaae/a//e 14052v3.
12

---

Published as a conference paper at ICLR 2025
Y. Dauphin, Angela Fan, Michael Auli, and David Grangier. Language modeling with gated convolutional networks. International Conference On Machine Learning, 2016. URL https://arxiv. http:///aakivvoo/obb/////s ab611 08883v3.
Soham De. Samuel L. Smith, Anushan Fernando, Aleksandar Botev. George
Albert Gu, Ruba
Haroun, Leonard Berrada, Yutian Chen, Srivatsan Srinivasan, Guillaume Desjardins, Arnaud Doucet, David Budden, Yee Whye Teh, Razvan Pascanu, Nando De Freitas, and Caglar Gulcehre. Griffin: Mixing gated lincar recurrences with local attention for efficient language models. arXiv preprint arXiv: 2402.19427. 2024. URL
19427v1.
Yiran Ding, L. Zhang, Chengruidong Zhang, Yuanyuan Xu, Ning Shang, Jiahang Xu, Fan Yang, and Mao Yang. Longrope: Extending IIm context window beyond 2 million tokens. International Conference on Machine Learning, 2024. doi: 10.485550/rxii.2422...5 137553
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, Anirudh Goyal, Anthony Hartshorn, Aobo Yang. Archi Mitra, Archic Sravankumar, Artem Korenev, Arthur Hinsvark, Arun Rao, Aston Zhang, Aurelien Rodriguez, Austen Gregerson, Ava Spataru, Baptiste Roziere, Bethany Biron, Binh Tang, Bobbic Chem, Charlotte Caucheteux, Chaya Nayak, Chloe Bi, Chris Marra, Chris McConnell, Christian Keller, Christophe Touret, Chunyang Wa, Corinne Wong. Cristian Canton Ferrer, Cyrus Nikolaidis, Damien Allonsius, Daniel Song, Danielle Pintz, Danny Livshits, David Esiobu, Dhruv Choudhary, Dhruv Mahajan, Diego Garcia-Olano, Diego Perino, Dieuwke Hupkes, Egor Lakomkin, Ehab AlBadawy, Elina Lobanova, Emily Dinan, Eric Michael Smith, Filip Radenovie, Frank Zhang, Gabriel Synnaeve, Gabrielle Lee, Georgia Lewis Anderson, Graeme Nail, Gregoire Mialon, Guan Pang. Guillem Cucurell, Hailey Nguyen, Hannah Korevaar, Hu Xu. Hugo Touvron, Iliyan Zarov, Imanol Arrieta Ibarra, Isabel Kloumann, Ishan Misra, Ivan Evtimov, Jade Copet, Jaewon Loe, Jan Geffert, Jana Vranes, Jason Park, Jay Mahadcokar, Jeet Shah, Jelmer van der Linde, Jennifer Billock, Jenny Hong, Jenya Lee, Jeremy Fu, Jianfeng Chi, Jianyu Huang, Jiawen Liu, Jie Wang. Jiecao Yu. Joanna Bitton. Joe Spisak, Jongsoo Park. Joseph Rocca, Joshua Johnstun, Joshua Saxe, Junteng Jia, Kalyan Vasuden Alwala. Kartikeya Upasani. Kate Plawiak, Ke Li, Kenneth Heafield, Kevin Stone, Khalid El-Arini, Krithika Iyer, Kshitiz Malik, Kuenley Chiu, Kunal Bhalla, Lauren Rantala- Yeary. Laurens van der Maaten, Lawrence Chen, Liang Tan, Liz Jenkins, Louis Martin, Lovish Madaan, Lubo Malo, Lukas Blecher, Lukas Landzaat, Luke de Oliveira, Madeline Muzzi, Mahesh Pasupuleti, Mannat Singh, Manohar Paluri, Marcin Kardas. Mathew Oldham, Mathicu Rita, Maya Pavlova, Melanie Kambadur, Mike Lewis, Min Si, Mitesh Kumar Singh, Mona Hassan, Naman Goyal, Narjes Torabi, Nikolay Bashlykov, Nikolay Bogoychev, Niladri Chatterji. Olivier Duchenne, Onur Çelebi, Patrick Alrassy, Pengchuan Zhang, Pengwei Li, Petar Vasic, Peter Weng, Prajjwal Bhargava, Pratik Dubal, Praveen Krishnan, Punit Singh Koura, Puxin Xu, Qing He. Qingxiao Dong. Ragavan Srinivasan. Raj Ganapathy. Ramon Calderer. Ricardo Silveira Cabral, Robert Stojnic, Roberta Raileanu, Rohit Girdhar, Rohit Patel, Romain Sauvestre, Ronnie Polidoro, Roshan Sumbaly, Ross Taylor, Ruan Silva, Rui Hou, Rui Wang. Saghar Hosseini, Sahana Chennabasappa. Sanjay Singh, Sean Bell, Scohyun Sonia Kim, Sergey Edunov, Shaoliang Nie, Sharan Narang, Sharath Raparthy, Sheng Shen, Shengye Wan, Shruti Bhosalc, Shun Zhang, Simon Vandenhende, Soumya Batra, Spencer Whitman, Sten Sootla, Stephane Collot, Suchin Gururangan, Sydney Borodinsky, Tamar Herman, Tara Fowler, Tarck Sheasha, Thomas Georgiou, Thomas Scialom, Tobias Speckbacher. Todor Mihaylov, Tong Xiao, Ujjwal Karn, Vedanuj Goswami. Vibhor Gupta, Vignesh Ramanathan, Viktor Kerkez, Vincent Gonguet, Virginie Do Vish Vogeti, Vladan Petrovic, Weiwei Chu, Wenhan Xiong, Wenyin Fu, Whitney Meers, Xavier Martinet, Xiaodong Wang, Xiaoqing Ellen Tan, Xinfeng Xie, Xuchao Jia. Xuewei Wang. Yaelle Goldschlag, Yashesh Gaur, Yasmine Babaci, Yi Wen, Yiwen Song. Yuchen Zhang, Yue Li. Yuning Mao, Zacharie Delpierre Coudert, Zheng Yan, Zhengxing Chen, Zoe Papakipos, Aaditya Singh, Aaron Grattafiori, Abha Jain, Adam Kelsey, Adam Shajnfeld, Adithya Gangidi, Adolfo Victoria, Ahuva Goldstand, Ajay Menon, Ajay Sharma, Alex Boesenberg, Alex Vaughan, Alexci Baevski, Allie Feinstein, Amanda Kallet, Amit Sangani, Anam Yunus, Andrei Lupu, Andres Alvarado, Andrew Caples, Andrew Gu, Andrew Ho, Andrew Poulton, Andrew Ryan, Ankit Ramchandani, Annie Franco. Aparajita Saraf, Arkabandhu Chowdhury, Ashley Gabriel. Ashwin Bharambe, Assaf Eisenman, Azadeh Yazdan, Beau James, Ben Maurer. Benjamin Leonhardi. Bernie Huang. Beth Loyd. Beto De Paola, Bhargavi Paranjape, Bing Liu, Bo Wu, Boyu Ni, Braden Hancock, Bram Wasti, Brandon Spence, Brani Stojkovic, Brian Gamido, Britt Montalvo, Carl Parker, Carly Burton, Catalina Mcjia, Changhan Wang. Changkyu Kim, Chao Zhou, Chester Hu, Ching-Hsiang Chu, Chris Cai, Chris Tindal, Christoph Feichtenhofer, Damon Civin, Dana Beaty, Daniel Kreymer, Daniel Li. Danny Wyatt, David Adkins, David Xu, Davide Testuggine, Delia David, Devi Parikh, Diana Liskovich, Didem Foss, Dingkang Wang, Duc Le. Dustin Holland, Edward Dowling, Eissa Jamil, Elaine Montgomery, Eleonora Presani, Emily Hahn, Emily Wood, Erik Brinkman, Esteban Arcaute, Evan Dunbar, Evan Smothers, Fei Sun, Felix Kreuk, Feng Tian, Firat Ozgenel, Francesco Caggioni, Francisco Guzmán, Frank Kanayet, Frank Seide, Gabriela Medina Florez, Gabriella Schwarz, Gada Badeer, Georgia Swee, Gil Halpern, Govind Thattai, Grant Herman, Grigory Sizov, Guangyi, Zhang, Guna Lakshminarayanan, Hamid Shojanazeri. Han Zou, Hannah Wang. Hanwen Zha, Haroun Habeeb, Harrison Rudolph. Helen Suk, Henry Aspegren, Hunter Goldman, Igor Molybog, Igor Tufanow, Irina-Flena Veliche, Itai Gat, Jake Weissman, James Geboski, James Kohli, Japhet Asher, Jean- Baptiste Gaya, Jeff Marcus, Jeff Tang, Jennifer Chan, Jenny Zhen, Jeremy Reizenstein, Jeremy Tehoul, Jessica
Rantala-Yeary,
13

---

Published as a conference paper at ICLR 2025
Zhong, Jian Jin, Jingyi Yang. Joe Cummings, Jon Carvill, Jon Shepard, Jonathan McPhie, Jonathan Torres, Josh Ginsburg, Junjie Wang, Kai Wu, Kam Hou U, Karan Saxena, Karthik Prasad, Kartikay Khandelwal, Katayoun Zand, Kathy Matosich, Kaushik Veeraraghavan, Kelly Michelena, Keqian Li, Kun Huang, Kunal Chawla, Kushal Lakhotia, Kyle Huang, Lailin Chen, Lakshya Garg. Lavender A, Leandro Silva, Lee Bell, Lei Zhang. Liangpeng Guo, Licheng Yu, Liron Moshkovich, Luca Wehrstedt, Madian Khabsa, Manav Avalani, Manish Bhatt, Maria Tsimpoukelli, Martynas Mankus, Matan Hasson, Matthew Lennie, Matthias Reso, Maxim Groshev, Maxim Naumov, Maya Lathi, Meghan Keneally, Michael L. Seltzer, Michal Valko, Michelle Restrepo, Mihir Patel, Mik Vyatskov, Mikayel Samvelyan. Mike Clark, Mike Macey, Mike Wang. Miquel Jubert Hermoso, Mo Metanat, Mohammad Rastegari, Munish Bansal, Nandhini Santhanam, Natascha Parks, Natasha White. Navyata Bawa. Nayan Singhal, Nick Egebo, Nicolas Usunier, Nikolay Pavlovich Laptev, Ning Dong Ning Zhang, Norman Cheng, Oleg Chernoguz, Olivia Hart, Omkar Salpekar, Ozlem Kalinli, Parkin Kent, Parth Parekh, Paul Saab, Pavan Balaji, Pedro Ritner, Philip Bontrager, Pierre Roux, Piotr Dollar, Polina Zvyagina, Prashant Ratanchandani, Pritish Yuvraj, Qian Liang, Rachad Alao, Rachel Rodriguez, Rafi Ayub, Raghotham Murthy, Raghu Nayani, Rahul Mitra, Raymond Li, Rebekkah Hogan, Robin Battey, Rocky Wang. Rohan Maheswari, Russ Howes. Ruty Rinott, Sai Jayesh Bondu, Samyak Datta, Sara Chugh, Sara Hunt, Sargun Dhillon, Sasha Sidorov, Satadru Pan, Saurabh Verma, Seiji Yamamoto, Sharadh Ramaswamy, Shaun Lindsay, Shaun Lindsay, Sheng Feng, Shenghao Lin, Shengxin Cindy Zha, Shiva Shankar, Shugiang Zhang, Shuqiang Zhang, Sinong Wang. Sneha Agarwal, Soji Sajuyigbe, Soumith Chintala. Stephanic Max. Stephen Chen, Steve Kehoc, Steve Satterfield. Sudarshan Govindaprasad, Sumit Gupta, Sungmin Cho, Sunny Virk, Suraj Subramanian, Sy Choudhury, Sydney Goldman, Tal Remez, Tamar Glaser, Tamara Best, Thilo Kohler, Thomas Rohinson, Tianhe L.,, Tianjun Zhang. Tim Matthews, Timothy Chou, Tzook Shaked, Varun Vontimitta, Victoria Ajayi, Victoria Montanez, Vijai Mohan, Vinay Satish Kumar, Vishal Mangla, Vlad Ionescu, Vlad Poenaru. Vlad Tiberiu Mihailescu. Vladimir Ivanov, Wei Li. Wenchen Wang. Wenwen Jiang, Wes Bouaziz, Will Constable. Xiaocheng Tang, Xiaofang Wang. Xiaojian Wu, Xiaolan Wang. Xide Xia, Xilun Wu, Xinbo Gao, Yanjun Chen, Ye Hu, Ye Jia, Ye Qi. Yenda Li, Yilin Zhang, Ying Zhang, Yossi Adi, Youngiin Nam, Yu, Wang, Yuchen Hao, Yundi Qian, Yuzi Hc, Zach Rait, Zachary DeVito, Zef Rosnbrick, Zhaoduo Wen, Zhenyu Yang, and Zhiwci Zhao. The llama 3 herd of models. arXiv preprint arXiv: 2407.2 21733, 2024. URL https: //tpxiv.////xxvooo http///wwww..iaaa///ooooo//oo/// hatgg///b//////// 21783v1.
Stefan Elfwing, F. Uchibe, and K. Doya. Sigmoid-weighted linear units for neural network function approxima- tion in reinforcement learning. Neural Networks, 2017. doi: 00.1016f.neumet.211.1.1
Mahan Fathi, Jonathan Pilault, Orhan Firat, Christopher Pal, Pierre-Lacc Bacon, and Ross Goroshin. Block-state transformers. NEURIPS, 2023. URL
Daniel Y Fu, Tri Dao, Khaled Kamal Saab, Armin W Thomas, Atri Rudra, and Christopher Re. Hungry hungry hippos: Towards language modeling with state space models. In The Eleventh International Conference on Learning Representations, 2023. URL https: http://pppppeeeee/eeeppeppp/p wwwp:/eefformmmiieeeeoo dd-ooDeww......
Albert Gu and Tri Dao. Mamba: I. Lincar-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.0075: 2023.
Albert Gu, Karan Goel, and Christopher R'e. Efficiently modeling long sequences with structured state spaces. International Conference On Learning Representations, 2021.
Albert Gu, Ankit Gupta, Karan Goel, and Christopher Ré. On the parameterization and initialization of diagonal state space models. ARXIVORG, 2022. doi: 10.48550//rXiw.22066 1183..
Chi Han, Qifan Wang. Wenhan Xiong, Yu Chen, Heng Ji, and Sinong Wang. Lm-infinite: Simple on-the-fly length generalization for large language models. arXiv preprint arXiv: 2308.16137 2023. URL https: //arxiv. org. abe/ 16137v3.
Kaiming He, Xiangyu Zhang. Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. CVPR, 2016. URL https:://arkess http:///aaxxv.o/aaaabbao abs/ 1atss////..... 03385v1.
Yihui He, Jianing Qian, Jianren Wang, Cindy X. Le, Congrui Hetang, Qi Lyu. Wenping Wang. and Tianwei Yue. Depth-wis wise decomposition for accelerating separable convolutions in efficient convolutional neural networks. arXiv preprint arXiv: 1910.09455, 2019. URL hnttps://arxe org/abs/11n 0995533...
Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Stein- hardt. Measuring massive multitask language understanding. 9th International Conference on Learn- ing Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. URL https://opeereviiin net net/forum?? peet/forun?dddd3jjm11 1ddd7dB3mm33mmmm
Ari Holtzman, Jan Buys, Li Du, Maxwell Forbes, and Yejin Choi. The curious case of neural text degeneration. In- termational Conference on Learning Representations, 2019, URL htttps://arxvv.org///s 9977195v.......
14

---

Published as a conference paper at ICLR 2025
Luyang Huang. Shuyang Cao. Nikolaus Parulian, Heng Ji, and Lu Wang. Efficient attentions for long document summarization. Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pp. 1419-1436, 2021.
Samy Jelassi, David Brandfonbrener, Sham M. Kakade, and Eran Malach. Repeat after me: Transformers are better than state space models at copying. arXiv preprint arXiv: 2402.01032, 2024. URL https:
//arxiv.org htt//bbs/24/// 2402.0222200 000.22001030 01032v1.
Albert Q. Jiang, Alexandre Sablayrolics, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, Lélio Renard Lavaud, Marie- Anne Lachaux, Pierre Stock, Teven Le Scao, Thibaut Lavril, Thomas Wang. Timothée Lacroix, and William El Sayed. Mistral 7b, arXiv preprint arXiv; 2310.06825, 2023. URL wtt.tttts///aatte www. hatgg//b///330/ 2.10..///////////e 068225v1
Hongye Jin, Xiaotian Han, Jingfeng Yang, Zhimeng Jiang, Zirui Liu, Chia-Yuan Chang, Huiyuan Chen, and Xia Hu. Llm maybe longlm: Self-extend IIm context window without tuning. arXiv preprint arXiv: 2401.01325, 2024. URL //atps///xxiiv org/ abs/2481. 01 325v1.
Tobias Katsch. Gateloop: Fully data-controlled linear recurrence for sequence modeling. arXiv preprint arXiv: 2311.01927, 2023. URL
2311.01927v1.
Nikita Kitaev, Fukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. arXiv preprint
arXiv:2001.0445// 2020.
Yuanzhi Li, Sébastien Bubeck, Ronen Eldan, Allic Del Giomo, Suriya Gunasekar, and Yin Tat Lec, Textbooks are all you need ii: phi-1.5 technical report. arXiv preprint arXiv: 2309.05463, 2023. URL https: wwwp:///ggbb//299999 //arxiv.org//////// http//bb/223000 2309../////eeee 05463v1.
Opher Lieher, Barak Lenz Hofit Bata, Gal Cohen, Jhonathan Osin, Itay Dalmedigos, Frez Safahi, Shaked Meirom, Yonatan Belinkov, Shai Shalev-Shwartz, Omri Ahend, Raz Alon, Tomer Asida, Amir Bergman, Roman Glozman, Michael Gokhman, Avashalom Manevich, Nir Ratner, Noam Rozen, Erez Shwartz, Mor Zusman, and Yoav Shoham. Jamba: A hybrid transformer-mant language model. arXiv preprint arXiv: 2403. 19887, 2024. URL https://arxiv.. or http/bb///443 19887v1.
Stephanie Lin, Jacob Hilton, and Owain Evans. TruthfulQA: Measuring how models mimic human falsehoods. In Smaranda Muresan, Preslav Nakov, and Aline Villavicencio (eds.), Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 3214-3252, Dublin, Ireland, may 2022. Association for Computational Linguistics. doi: 10.18653/1//0222aacl 2ong.229. URL bttps://achle https://aclaanthologg org/2022. acl acl--aa llng. 229.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations, 2018
Xuezhe Ma. Chunting Zhou, Xiang Kong, Junxian He, Liangke Gui, Graham Neubig, Jonathan May, and Luke Zettlemoyer. Mega: Moving average equipped gated attention. In The Eleventh International Conference on Learning Representations, 2023. URL. https: / //tpp//eveeeeeeee//////ep. www.rreeeeeeeweefrunneeeee wett:/orrrmiddigeee3eieee.
Xuezhe Ma, Xiaomeng Yang, Wenhan Xiong, Beidi Chen, Lili Yu, Hao Zhang, Jonathan May, Luke Zettlemoyer, Omer Levy, and Chunting Zhou. Megalodon: Efficient 11m pretraining and inference with unlimited context length. arXiv preprint arXiv: 2404,08801, 2024. URL https: http://awwvvi.r.r//e www.orgg////////ooo///oo/se.eess 2404. 08801v1.
Eric Martin and Chris Cundy. Parallelizing linear recurrent neural nets over sequence length. In 6th International Conference on Learning Representations, ICLR 2018, Vancouver, BC, Conada, April 30- May 3, 2018, Confer- ence Track Proceedings. OpenReview.net, 2018. URL https: /ttp////ppeeeeeeeee//eeeeeee net hoo//formmmmt weeee//orrmmidd////oer Hdywwwuwul--H------
Harsh Mchta, Ankit Gupta, Ashok Cutkosky, and Behnam Neyshabur. Long range language modeling via gated state spaces. In The Eleventh International Conference Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023. URL ht //ppnr/vie/eeeiieew net met//Forummid-555M11Yvvaaaa
William Merrill, Jackson Petty, and Ashish Sahharwal. The illusion of state in state-space models. arXiv preprint arXiv: 2404.08819, 2024. URL htttp:://arxiv.orrt
0804,,88119 08819v1.
MetaAl Introducing meta lama 3: The most capable openly available Ilm to date, 2024. URL: https: //ai ctt.tt..ccoooo/////o.oo.ooo watp:///mtaaaaaaaaaaaaaaaaaaaaaa 1lama-3/. 3/.
Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can suit of armor conduct electricity? new dataset for open book question answering. Conference on Empirical Methods in Natural Language Processing, 2018. doi: (0.18653/11/D1.. 1260. URL https: //axxiv.avvv///////////a orgg/bbs/8899 /0w.1/88.....8..
15

---

Published as a conference paper at ICLR 2025
Amirkeivan Mohtashami and Martin Jaggi. Landmark attention: Random-access infinite context length for transformers. arXiv preprint arXiv: 2305. 16300, 2023. URL https://arxiv.. wwtp:/iiioooooooeeeee att//b///2//0 2305. 16300v2.
Tsendsuren Munkhdalai, Manaal Faruqui, and Siddharth Gopal. Leave context behind: Efficient infinite context transformers with infini-attention. arXiv preprint arXiv: 2404.07143, 2024, URL https://arxiv,
org/abs/24944 87143v1.
OpenAl. Gpt-4 technical report. PREPRINT, 2023, URL https::/arrivv.rr//abbb///////// 08774v4.
Antonio Orvieto. Samuel L. Smith. Albert Gu, Anushan Fernando. Caglar Gulcehre. Razvan Pascanu, and Soham De. Resurrecting recurrent neural networks for long sequences. International Conference on Machine Learning, 2023, doi: (0.48550/rXXiv22000 06349. URL petps://arxtee wwwpeiiii.orggaaeeeee www./rrg///////////////oe /2003. 06349v1.
Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Q. N. Pham, R. Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and R. Fernández. The lambada dataset: Word prediction requiring broad discourse context. Annual Meeting of the Association for Computational Linguistics, 2016. doi: 18.18653/v1/PI11 1144.
Jongho Park, Jacseung Park, Zheyang Xiong, Nayoung Lee, Jacwoong Cho, Samet Oymak, Kangwook Lec, and Dimitris Papailiopoulos. Can mamba learn how learn? a comparative study on in-context learning tasks. arXiv preprint arXiv: 2402.04248, 2024. URL hps: / /twxiv..iiiiii///ii///ee/eaaa .og. .og. 220/2222./22222////// 04248v1.
Michael Poli. Stefano Massaroli, Eric Q. Nguyen, Daniel Y. Fu, Tri Dao, Baccus. Bengio, Stefano Ermon, and Christopher Ré. Hyena hierarchy: Towards larger convolutional language models. International Conference On Machine Learning, 2023. doi: 10.485550/rXxiv22322 10866. URL. httpps://arxiv or abs/ 2302. 10866v3.
Ofir Press, Noah A. Smith, and M. Lewis. Train short, test long: Attention with lincar biases enables input length extrapolation. International Conference On Learning Representations, 2021. URL https://arxiv org/abs/22 /2 18. 12409v2.
Zhen Qin, Songlin Yang. and Yiran Zhong. Hiierachically gated recurrent neural network for sequence modeling. Neural Information Processing Syxtems, 2023. doi: 104485500/xXi22331.482,, URL https:
Zhen Qin, Songlin Yang, Weixuan Sun, Xuyang Shen, Dong Li, Weigao Sun, and Yiran Zhong. Hgrn2: Gated linear mns with state expansion. arXiv preprint arXiv: 2404.07904, 2024. URL. httpht/ss//:tiiiii wwtpx/ii.oog//////////eeee//so 2404.0 07984v1.
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and llya Sutskever. Language models unsupervised multitask learners. arXiv preprint, 2019. URL https://api. semanticscholar.com org// CorpusID: 160025533.
Pranav Rajpurkar. Jian Zhang, Konstantin Lopyrev, and Percy Liang. Squad: 100,000+ questions for machine comprehension text. EMNLP. 2016. URL httpp:://arxiis wrr.iioooooooeeeooo orgg/rr/////////// /b000/6660.00200...
David Rein, Betty Li Hou, Asa Cooper Stickland, Jackson Petty, Richard Yuanzhe Pang. Julien Dirani, Julian Michael, and Samuel R. Bowman. Gpqa: A graduate-level google-proof qa henchmark. arXiv preprint arXiv: 2311. 12022, 2023. URL https:///rxiix hww.voogg/aaae wtwg/br///////////oo/eeee 2211. 12022v1.
Liliang Ren, Yang Liu, Shuohang Wang, Yichong Xu, Chenguang Zhu, and ChengXiang Zhai Sparse modular activation for efficient sequence modeling. NEURIPS, 2023. URL htttps:///akiii htt.aaiiii..or/////2//3/// 2306 111 11197v1.
Aurko Roy, M. Saffar, Ashish Vaswani, and David Grangier. Efficient content- sparse attention with routing transformers. International Conference On Topology, Algebra And Categories In Logic, 2020. doi: 10.1162/tacl a 00353.
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial winograd schema challenge scale. Communications of the ACM, 64(9)99-106, 2021. URL https: //arxiv. wwwpiiiiorrg/bb.../ 1977. 10641v2.
Maarten Sap, Hannah Rashkin, Derck Chen, Ronan LeBras, and Yejin Choi. Socialiqu: Commonsense reasoning about social interactions. arXiv preprint Xiviv: 1904.09728, 2019. URL htttp:://arkiv.org/abbsss 09728v3.
Imanol Schlag. Kazuki Irie, and Jürgen Schmidhuber. Linear transformers are secretly fast weight programmers. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, ICMI. 2021, 18-24 July 202 Virtual Event, volume 139 of Proceedings of Machine Learning Research, pp. 9355-9366. PMLR, 2021. URL htw.::tpp:/ppeee http:://proceeiiggs... mlr wwwpp/esssvvvve 39/ schlaglllllla htmml.
16

---

Published as a conference paper at ICLR 2025
Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, and Tri Dao. Flashattention-:: Fast and accurate attention with asynchrony and low-precision. arXiv preprint arXiv: 2407,08608, 2024, URL https://arxxiv. orrg org.ab// at////b/////0 240.2/80.0.... 08608v2.
Uri Shaham, Maor Ivgi, Avia Efrat, Jonathan Berant, and Omer Levy. Zeroscrolls: A zero-shot benchmark for long text understanding. Conference on Empirical Methods in Natural Language Processing, 2023, doi: 2L.48550arXii.235.... 14196. URL https: http://arvvvvo//ee/eee/oe arggabbs2335.... 14196v2.
Noam Shazeer. Fast transformer decoding: One write-head is all you need. arXiv preprint arXiv: 1911.02150.
2019. URL https://arxiv... orgg ht///b///111111 /bbs/1911.0210v1 02150v1.
Noam Shazeer. Glu variants improve transformer. arXiv preprint arXiv: 2002.05202 2020. URL https: www.iii.org/abbe.. 2002.0222.....
Jimmy T.H. Smith, Andrew Warrington, and Scott Linderman. Simplified state space layers for sequence modeling. In The Eleventh International Conference on Learning Representations, 2023. URL https: //openreview. net net/forum um? wwmiidddiiHwaAAKe.
Daria Soboleva, Faisal Al- Al-Khateeb, Robert Myers, Jacob R Steeves, Joel Hestness, and Nolan Dey. Slimpajama: A 627b token cleaned and deduplicated version of redpajama, 2023, URL: https:/ https::///wwwcc//// wwwebrrr..nneeeeee blog/slimpajer siipsaammaaaaaaaaa a -27b- token- ceaennn-eeeeedddeeee and- deduplicated wersteedvveiiioo----e of-rad
Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. Roformer: Enhanced
transformer with rotary position embedding. arXiv preprint arXiv: 2104.09864, 2021.
Yutao Sun, Dong, Shaohan Huang. Shuming Ma, Yuqing Xia, Jilong Xue, Jianyong Wang. and Furu Wei. Retentive network: A successor to transformer for large language models. arXiv preprint arXiv:2307.086: 2023.
Gemma Team. Gemma: Open models based on gemini research and technology. arXiv preprint arXiv: 2403.08295, 2024. URL httphttss///aiii.ooott arg/////24033..... 221///22333000 08299v1..
Jamba Team, Barak Lenz, Alan Arazi, Amir Bergman, Avshalom Manevich, Barak Peleg, Ben Aviram, Chen Almagor, Clara Fridman, Dan Padnos, Daniel Gissin, Daniel Jamnai, Dor Muhlgay, Dor Zimberg, Edden M Gerber, Elad Dolev, Eran Krakovsky, Frez Safahi, Erez Schwartz, Gal Cohen, Gal Shachaf, Haim Rozenblum, Hofit Bata, Ido Blass, Inbal Magar, Itay Dalmedigos, Jhonathan Osin, Julie Fadlon, Maria Rozman, Matan Danos, Michael Gokhman, Mor Zusman. Naama Gidron, Nir Ratner, Noam Gat, Noam Rozen, Oded Fried, Ohad Leshno. Omer Antverg. Omri Abend, Opher Lieber, Or Dagan, Orit Cohavi, Raz Alon, Ro'i Belson, Roi Cohen, Rom Gilad, Roman Glozman, Shahar Lev, Shaked Meirom, Tal Delbari, Tal Ness, Tomer Asida, Tom Ben Gal, Tom Braude, Uriya Pumerantz, Ychoshua Cohen, Yonatan Belinkow, Yuval Globerson, Yuval Peleg Levy, and Yoav Shoham. Jamba-1.5: Hybrid transformer-mamt models at scale. arXiv preprint arXiv: 2408, 12570, 2024. URL https: http://.aavvvvaaa org/ www.oogg///////eeee/ooo 2408/ 12570v1.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajwwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cristian Canton Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshom, Saghar Hosseini, Rui Hou, Hakan Inan. Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Maric-Ame Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nic, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiaoqing Ellen Tan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu. Zheng Yan, Iliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurelien Rodriguez. Robert Stojnic, Sergey Edunov, and Thomas Scialom. Llama Open foundation and fine-tuned chat models. arXiv preprint arXiv: 2307.09288. 2023. URL httpstttss::/ii...tt// wrt./b//////////ooooooooo 0928882..
Szymon Tworkowski, Konrad Staniszewski, Mikotaj Pacek, Yuhuai Wa, Henryk Michalewski, and Piotr Mitos. Focused transformer: Contrastive training for context scaling. NEURIPS, 2023. URL https:
///aa/avvrrg/abs////.//. 03170v2.
Dusan Varis and Ondfej Bojar. Sequence length a domain: Length- based overfitting in transformer models. In Marie-Francine Moens, Xuanjing Huang. Lucia Specia, and Scott Wen-tau Yih (eds.), Proceedings of the 2021 Conference Empirical Methods in Natural Languuge Processing, pp. 8246-8257, Online and Punta Cana, Dominican Republic, nov 2021. Association for Computational Linguisties. doi: 10. 1/221 emnlp- main.m0....50. URL https: http://///caaathllog/ooo
2/rg/221...2 .em.1p.mmlppmmmm main. 650.
Ashish Vaswani, Noam M. Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. NIPS, 2017.
17

---

Published as a conference paper at ICLR 2025
Alex Wang. Richard Yuanzhe Pang. Angelica Chen, Jason Phang. and Samuel R. Bowman. Squality: Building long-document summarization dataset the hard way. Conference on Empirical Methods in Natural Language Processing, 2022. doi: 10.48550/arX1v22205. 1465, URL
11465v1.
Yubo Wang. Xueguang Ma, Ge Zhang, Yuansheng Ni, Abhramil Chandra, Shiguang Guo, Weiming Ren, Aaran Arulraj, Xuan He, Ziyan Jiang, Tianle Li, Max Ku, Kai Wang. Alex Zhuang, Rongqi Fan, Xiang Yue, and Wenhu Chen. Mmlu-pro: A more robust and challenging multi-task language understanding benchmark. arXiv preprint aiv: 2406,01574, 2024. URL atttps://arxtii wogg abs//2/5..22.. 01574v4.
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, E. Chi, F. Xia, Quoc Le, and Denny Zhou. Chain-of-thought prompting elicits reasoning in large language models. Neural Information Processing Systems, 2022. URL https: //tp///aaivvv///e/////eeeeeaaava www.iv..oo/oo.oeee ht/g//bbb/221... 11903v6.
Kaiyue Wen, Xingyu Dang, and Kaifeng Lyu. Runs are not transformers (yet): The key bottleneck on in-context retrieval. arXiv preprint Xrviv 2402. 18510, 2024, URL https://arxiv.. wwt.:/rr///////////// 18510v1
Yuhuai Wu, Markus N. Rabe, DeLesley S. Hutchins. and Christian Szegedy. Memorizing transformers. International Conference On Learning Representations, 2022. doi: 10448550/a2Xi.2223..890) URL
https://arxxv.. hww.v.org/bbsss/// abs/ 2203. 08913v1.
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks, arXiv preprint arXiv: 2309.1745.. 2023. URL https://arxiv. wttp:r/wwww..rr/////////////// 17453v1.
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net, 2024. URL https: / /opn//eeeeeeeeeeeeeeeeeeeee net um #SG=GGG75512
Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Yan Liu. On layer normalization in the transformer architecture. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pp. 10524-10533, PMLR, 2020. URL http: //proceedings mlr wwtr.pres//vv. 19/ 9/wp//ggibb../////////////o..b.e httm..
Wenhan Xiong. Jingyu Liu, Igor Molybog, Hejia Zhang, Prajjwal Bhargava, Rui Hou, Louis Martin, Rashi Rungta, Karthik Abinav Sankararaman. Barlas Oguz, Madian Khabsa, Han Fang, Yashar Mehdad, Sharan Narang, Kshitiz Malik, Angela Fan, Shruti Bhosale, Sergey Edunov, Mike Lewis, Sinong Wang, and Hao Ma. Effective long-context scaling of foundation models. arXiv preprint arXiv: 2309. 16039, 2023.
Songlin Yang and Yu Zhang. Fla: A triton-based library for hardware-efffiient implementations of linear attention mechanism, January 2024. URL https: /gi thub. com/ sus wwwosgloollggooooo in wwt.//fashhhaaa
Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim. Gated linear attention transformers with hardware-efficient training. arXiv preprint arXiv:2312.06635, 2023.
Songlin Yang, Bailin Wang, Yu Zhang, Yikang Shen, and Yoon Kim. Parallelizing linear transformers with the delta rule over sequence length. arXiv preprint arXiv: 2406.06484 2024. URL httpp::/arxiv..orrr wwwpsiivooorg/oobbe abs/ 2406. 06484v1.
Manzil Zaheer, Guru Guruganesh, Kumnr Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang. et al. Big bird: Transformers for longer sequences. Advances in neural information processing systems, 33:17283 17297, 2020. URL https://arxiv: orrg/ abs/2007. 14062v2.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can a machine really finish your sentence? Annual Meeting of the Association for Computational Linguistics, 2019. doi: 10. 1L.18665/1/P1191473 URL https::/s//: www.aawwww.oooooooooooooooeoooor h/tg//bs/19..// 07830v1.
Biao Zhang and Rico Sennrich. Root mean square layer normalization. Neural Information Processing Systems, 2019. doi: 10.5167/UZH- 177483. URL https: http://wawww.waaaa/////// 1910. 07467v1.
Edward Ayers Zhangir Azerbayev and Bartosz Piotrowski. Proof-pile, 2022.
URL: htL::htts://giithbbbcom
thanggrrazerbbaee//roooo
Hao Zheng, Zhanlei Yang, Wenju Liu, Jizhong Liang, and Yanpeng Li. Improving deep neural networks using softplus units. 2015 International Joint Conference on Neural Networks (IJCNN). pp. 1-4, 2015. doi:
10.1109/1JCNN.20155 2015.72804599 URL https: / h/ieex///eeeloorreee... ieee. wtrg//occment/22880
Simiao Zuo, Xiaodong Liu, Jian Jiao, Denis Charles, Eren Manavoglu, Tuo Zhao, and Jianfeng Gao. Efficient long sequence modeling via state space augmented transformer. arXiv preprint arxiv: 2212.08136, 2022. URL httpsst//////i.....tm/ wrw. abs/221 0818883366 36v1.
18

---

Published as a conference paper at ICLR 2025
A RELATED WORKS
Hybrid Recurrent Models Many recent works (Park et al.. 2024; Jelassi et al., 2024; Akyürek et al., 2024) point out the lack of retrieval ability of linear SSMs, and propose hybridization of SSMs with the Attention mechanism. However, the history of SSM/RNN- Attentioon hybridization can be directly dated back to the birth of the Attention mechanism (Bahdanau et al., 2014) which proposed as soft feature alignment technique for recurrent models to cope better with long sequences. The revitalization of the fact that linear recurrent models are sequentially parallelizable (Martin & Cundy, 2018; Gu et al., 2021) has catalyzed a contemporary renaissance in hybrid recurrent architectures. SPADE (Zuo et al., 2022), GSS (Mehta al., 2023), MEGA (Ma et al., 2023), Block State transformers (Fathi et al., 2023) and Megalodon (Ma et al., 2024) combine SSMs with chunked attention, while H3 (Dao et al., 2022b), Mambaformer (Park et al., 2024) and Jamba (Lieber et al., 2024; Team et al., 2024) propose to hybridize with quadratic self-attention. Our works focus particularly the wall-time efficiency and the length extrapolatability of the hybrid SSM-Attention models, and propose to interleave SSMs with Sliding Window Attention (SWA), which has both linear computation complexity and the conslation-invariann property over the sequence length. Infini- Attention (Munkhdalai et al., 2024) is a recently proposed method that implements an intra-layer hybridization (Wu et al., 2022) between SWA and Linear Attention with the delta rule (Schlag et al., 2021). While the preliminary results look promising, its performance in the setting of large-scale pre-training from scratch remains questionable. The most similar work to ours is Griffin (De et al., 2024), which interleaves the Real-Gated Linear Recurrent Unit (RG-LRU) with Sliding Window Attention (SWA). However, Samba hybridizes SWA with Mamba instead of RG-LRU and shows that this simple hybrid architecture can provide substantially better performance over state-of-the-an Transformer architectures across scales, while Griffin and its follow-up work RecurrentGemma (Botev et al., 2024) only show comparable or worse results than Transformers. The original Mamba paper (Gu& & Dao, 2023) also explores hybridizing pure Mamba models with full attention or MLP layers, but it does not consider the wall-time efficiency of these hybridization and only achieves marginally better performance than the pure Mamba model. In contrast, we are the first to show that interleaving Mamba with both SWA and MLP can substantially outperform modern Transformers (and Mamba) at a scale up to 3.8B parameters, while achieving comparable training speed and better length extrapolation ability under the perplexity metrics.
Efficient Sparse Attention Previous works have proposed sparsifying self-attention (Vaswani et al., 2017) with a static attention pattern (Child et al., 2019; Zaheer et al., 2020; Beltagy et al., 2020) or a dynamic learnable pattern (Roy et al., 2020; Kitaev al., 2020; Ren et al., 2023) to model long sequences with subquadratic complexity over the sequence length. However, due to the lack of hardware-aware efficient implementation, its actual wall-time training efficiency is often worse than the dense attention optimized with FlashAttention (Dao et al., 2022a; Dao, 2023; Shah et al., 2024). In this work, we choose Sliding Window Attention, a simple static sparse attention pattern, because can easily leverage the highly optimized FlashAttention kernels to enjoy an actual training speed-up over its dense self-attention counterpart.
Length Extrapolation Many previous works have focused on extending the context length of pretrained Transformers to improve their performance on long-context tasks. Methods such as LM- Infinite (Han et al., 2023), StreamingLLMM (Xiao al., 2024), and LongLoRA (Chen et al., 2023b) achieve linear complexity for length extrapolation, but they can only stabilize perplexity beyond the training sequence length rather than significantly improve it. In contrast, demonstrate that pre- training Transformers with Sliding Window Attention from scratch enables natural improvements in perplexity beyond the training sequence length. Other approaches, including LLaMA-2-Long (Xiong et al.. 2023), LongLLaMA (Tworkowski et al., 2023), PI (Chen et al., 2023a), LongRoPE (Ding et al., 2024) and Self-Extend (Jin al., 2024), attempt to extend the full attention through modifying position embedding or continual training strategies, but they typically retain quadratic complexity in the attention mechanism with additional computation or memory I/O overhead, therefore they do not scale well to very long sequences. Although these methods achieve improved perplexity on a sequence length that is multiple times longer than the training sequence length, their perplexity still explodes if the sequence is extremely long. Our method achieves both linear complexity and superior extrapolation performance compared to zero-shot length extrapolation methods, such as Self-Extend, under the perplexity metric. However, acknowledge that, in terms of zero-shot
19

---

Published as a conference paper at ICLR 2025
retrieval performance, method still lags behind these approaches. This underscores trade-off between perplexity and retrieval performance in length extrapolation, which we plan to explore and address in future work.
B ADDITIONAL EVALUATION RESULTS
In Table 7, we conduct comprehensive evaluations on a diverse subset of benchmarks to assess SAMBA 3.8B base model's performance across all the domains mentioned in Section ensure thorough examination of the model's capabilities. We also report the performance of the Trans- former++ (TFM++) model, which uses the same architecture, pre-training recipe as Phi3-mini, for fair comparison. The details of the generation configurations are included in Appendix G. We compare with several strong baselines, including Llama (Touvron et al., 2023), Mistral (Jiang et al., 2023), Mamba (Gu & Dao, 2023), Gemma (Team, 2024), Recurrent-Gemma (R-Gemma) (Botev et al., 2024), Llama (MetaAI, 2024) and TFM++. As shown Table SAMBA achieves highest average score on all benchmarks, demonstrating its superior performance in handling various language comprehension tasks. Notably, SAMBA excels in the GSM8K benchmark, achieving an absolute 18.1% higher accuracy than TFM+ trained the same dataset. This shows the surprising complementary effect of combining SSM with the attention mechanism. We conjecture that when combined with attention, Mamba, arithmetic operation through its recurrent states than on doing the retrieval operation which can be easily learned by the sliding window attention.
input-dependent SSM, can focus more performing the
Table 7: Downstream performance comparison of the SAMBA 3.8B base model with other pretrained base language models without instruction tuning ARC-C and HellaSwag normalized accuracy. MMI.U and GSM8K are measured in 5-shot, while others are in zero-shot. We report the MC2 score for TruthfulQA, maj I for GSM8K, and pass@ for HumanEval. Measured by ours. The fair comparison should only be considered between TFM+ and Samba.
measured with character-
Hella- Swag
ARC-
Wino- Gran.
Truth.
GSM
Hum.
Avg.
Model
Size
Tokens
MMLU
QA
8K
Llama
6.7B
2T
45.3
77.2
45.9
69.2
38.8
14.6
12.8
43.4
13B
2T
54.8
80.7
49.4
72.8
37.4
28.7
18.3
48.9
Mistral
7.2B
60.1
81.3
55.5
75.3
42.2
35.4
30.5
53.6
Mamba
2.8B
600B
26.2
71.0
41.7
65.9
34.4"
3.6"
7.3*
35.7
Gemma
2.5B
3T
42.3
71.4
42.1
65.4
33.1
17.7
22.0
42.0
8.5B
6T
64.3
81.2
53.2
72.3
44.8
46.4
32.3
56.4
R-Gemma
2.7B
2T
38.4
71.0
42.3
67.8
35.1
13.4
21.3
41.3
Llama
8.0B
15T+
66.6
79.2"
53.2*
72.6*
43.9
45,8
28.7"
55.8
TFM+
3.8B
3.2T
67.2
76.6
53.8
72.6
47.3
51.5
51.8
60.1
SAMBA
3.8B
3.2T
71.2
77.4
55.7
77.1
43.4
69.6
54.9
64.2
As shown in Table we can see that post- tost-trained hybrid models can achieve superior performance compared to industry-standarr Transformer bassed LLMs such as Llama-3. Innttruct 8B and Llama- 3.2- Instructt 3B, and SSM-based I.MMs such FalconMamba Recent progress on hybrid L.MMs, including Jamba 1.5 (Team et al., 2024) and our own work on SAMBA, shows significant improvement over earlier approaches like R-Gemma (Botev et al., 2024), which hybridizes attention with linear recurrent models but is trained smaller data scales. SAMBA delivers comparable performance to Jamba- 1.5aa-155.. 1.5- while using around 3x fewer active parameters and 13x fewer total parameters, due advanced text-book data synthesis technique (Abdin et al., 2024). Additionally, SAMBA outperforms the Phi3 architecture, which is trained on the same data and optimization setting, further highlighting the superiority of our hybrid architecture over modern Transformer models.
20

---

Published as a conference paper at ICLR 2025
Table 8: Post-trained models quality on representative benchmarks under the chat mode. The fair comparison should only be considered between SAMBA and Phi3 as we control the training recipes and datasets to be the same. Best results are in bold, second best underlined.
SAMBA (June)
Phi3 (June)
R Gemma
FalconMambu
Jamba- 1.S-Mini
Llamu-3.2-16 3.2
Llama--3 3,1
Category
3.3B
3.8B
9B
7B
12B/52B
SB
MMLU
(5-shot)
MMLU
MMLU-Pro
17.8
(U-shol, CcT)
ARC-C
(10-shot)
Reasening
GPQA
29.5
4.7
32.3
26.3
(0-chot, CoT)
GSMSK
77.4
(8shot, CcT)
HumanEval
(0-shot)
Code
MBPP
71.7
70.0
42.0
75.8
67.2
69.4
(3-chat)
Averago
66.1
64.4
35,8
63.5
58.5
62.1
60000
Tokensis
30000
3.73*
Samba 1.78
20000
Llama-3 1.68
3
Mistral 1.60
10000
8K
16K
32K
Prompt Length
Figure 6: Prompt processing throughput of different models with around 1.7B parameters.
Mistral 1.63
3
Lo5
1
0
100
200
300
400
500
Steps
Figure 7: Training loss curves of Samba 1.7B and Mistral 1.6B models during 500 steps of instruction tuning on Passkey Retrieval with 4K sequence length. We plot the loss curves for both models using the simple moving average of window size 10.
21

---

Published as a conference paper at ICLR 2025
1.0
Pac-kkay Rairival Accunayy
0.8
0.6
0.4
Samba 1:76
0.2
Histral 1.68
50
150
200
250
300 Steps
350
400
450
500
Figure Overall passkey retrieval accuracy on the 256K document length of Samba 1.7B and Mistral 1.6B models during 500 steps of instruction tuning.
C ADDITIONAL EXPERIMENT DETAILS
We perform instruction tuning for both Mistral 1.6B and Samba 1.7B Passkey Retrieval using document length 4096, where we generated the data on the fly through randomly sampling a 5-digit integer passkey value and a location/depth between zero and the document length insert the passkey. The model is then asked to generate the passkey given the full document. We train both models using batch size 2048, 250 warm-up steps with a peak learning rate of 1e-4, and 0.1 weight decay with AdamW (Loshchilov & Hutter, 2018 8) optimizer. In both cases, the loss converges quickly in 100-200 steps. During the evaluation, we measure the overall average accuracies of the passkey retrieval the document length of [4k, 8k, 16k, 32k, 64k, 128k, 256k], for each length we evaluate at 1 different depths of the document (from 0, 0.1, 0.2, to 1.0). addition, for each location of the passkey evaluate the model with five different passkeys measure accuracy. As
(depth) in the document, seen Figure 8, the average passkey retrieval accuracy for Samba 1.7B almost reaches 100% in around 1 150 steps, while the accuracy for Mistral 1.6B remains low, demonstrating the extrapolation ability of the Samba architecture.
D ADDITIONAL ANALYSES
How to train models with Sliding Window Attention (SWA)? Since SWA has linear complexity with respect to the sequence length, it seems alluring trade off the batch size to have longer training sequence length without substantially decreasing the training throughput. However, as shown in Table 9, when the sequence length is increased, the validation perplexity also increases in all context lengths due to smaller batch sizes (Varis & Bojar, 2021), and the optimal ratio of sequence length/wing whngo//w size observed is 2, resulting in a training length of 4096.
Table 9: Perplexity on SlimPajama of -lama-2-SWA 438M models trained on different context sizes and batch sizes. We fix the sliding window size 2048 and the training tokens per step as 2M.
Training Speed
Validation Context Length
Batch Size
Sequence Length
(x 10 tokens/s)
2048
4096
8192
16384
1024
2048 (Full Attention)
10.4
11.59
38.12
156.18
357.32
512
4096
9.88
11.87
11.16
10.69
10.61
256
8192
9,66
11.98
11.26
10.79
10.69
128
16384
9.48
12.37
11.63
11.12
11.02
64
32768
9.29
12.94
12.46
11.96
11.86
httpp::/wwggigffacc.cc/ttteeeeee mamba- 2a mamba-7b-inntruct instruct
22

---

Published as a conference paper at ICLR 2025
Table 10: Perplexity on the SlimPajama validation set of different linear recurrent and sliding window attention models with Short Convolution (SC) modules added separately to query, key and value representations. For hybrid models, SC is applied only to linear attention layers. The training speed is measured on 8 A10 GPUs.
Training Speed
Validation Context Length
Architecture
Size
x100 tokens/s)
4096
8192
16384
Llama-2 2-SWA
438M
4.96
11.12
10.66
10.57
SC
438M
4.69
10.83
10.39
10.31
Sliding GLA
438M
4.94
10.43
10.00
9.92
SC
438M
4.44
10.39
9.96
9.87
Sliding RetNet
446M
4.32
10.38
9.96
9.87
SC
446M
3.80
10.25
9.82
9.74
Fair comparison between Mamba and other linear recurrent models? We can notice that the Short Convolution (SC) operator in Equation (I) is independent to the design of other parts of Mamba and can be applied other linear recurrent models. As shown in Table 10, we explore the effect of SC on model performance through enhancing Llama-2-1 WAA Sliding GLA, and Sliding RetNet with SC. Surprisingly, besides boosting the performance of RetNet, adding SC can also significantly improve the SWA's performance, while the effect on GLA is less prominent. We think this is because GLLA already has the fine- ineiggrined decays at the channel level, so the depthwise convolution doesn't 't add much of the useful inductive bias for better modeling power. Notably, even with the SC enhancer, Sliding GLA and Sliding RetNet still fall short than the original Samba 42 IM's performance shown in Table 3. This further justifies our choice of using Mamba for hybridization. We also find that adding SC to both the SWA and the linear attention layers in hybrid models produces negative results, and we leave it future work to understand the surprising effectiveness of SC in language modeling.
E DETAILS OF ENTROPY MEASUREMENT
Given a causal attention probability matrix A E Rxxxxn Aija = 0 < k, with h number of heads and sequence length of n, and the generation length 0 < 1< n, we calculate the average attention entropy per decoding step as follows,
Ha 1-
Σ
Aigk (Aggk)))).
For the selective gate A Raxd used by S6 in Equation (2) of the Mamba layers, first normalize it to be the simplex [0, mxd ie.,
A
E (0, 1]xx..
-= A;
The average selection entropy of S6 throughout the entire sequence is then calculated as
de
H.
A'g log (1y))
F DETAILS OF DOWNSTREAM LONG-CONTEXT EVALUATION
We use the GovReport (Huang et al., 2021) and the SQUALITY (Wang et al., 2022) datasets from the ZeroSCROLLS (Shaham et al., 2023) benchmark to evaluate models' long-context summarization capability in the real world. After tokenizing with the Phi3-mini-4k tokenizer, the average document length for the GovReport dataset is 11.533 tokens, with a median of 10,332, minimum of 1,493, and a maximum of 40,592 tokens. For the SQuALITY dataset, the average sequence length is 7,974 tokens, with median of 8, 145,, minimum of 5,457, and maximum of 10,757 tokens. For evaluation, we use greedy decoding for both tasks. A maximum generation length of 450 tokens is applied for GovReport and 600 for SQUALITY.
23

---

Published as a conference paper at ICLR 2025
G IMPLEMENTATION DETAILS
Table 11: Detailed ppper-parameters of the basclines models trained on the Phi2 dataset with 230B
tokens.
Architecture
Llama-3
Mistral
Mamba
Mamba-SWA-MLP
Mamba-MLL
Parameters
1.6B
1,6B
1.8B
1.6B
1.9B
Batch size
2048
2048
2048
2048
2048
Learning rate
0.0006
0.0006
0.0006
0.0006
0.0006
Weight decay
0.1
0.1
0.1
0.1
0.1
Gradient clipping
1.0
1.0
1.0
1.0
1.0
Sequence length
4096
4096
4096
4096
4096
Sliding window size,
2048
2048
Number of layers, N
54
48
Model width, dm.
2048
2048
2048
2048
2048
MLP intermediate size, dp
8196
8196
8196
8196
Number of query heads
32
32
32
32
Number of KV heads
Number of Attention Layers
24
18
Number of Mamba Layers
0
0
64
18
24
Vocabulary size
50304
50304
50304
50304
50304
For the GLA layer in the Sliding GLA architecture, we use the number of heads dm dm/384, a key expansion ratio of 0.5, and a value expansion ratio of 1. For the RetNet layer we use number of head that is half of the number attention query heads, key expansion ratio of and value expansion ratio of 2. The GLA and RetNet implementations are from the Flash Linear Attention (Yang & Zhang, 2024) repository". We use the FlashAttention bassoddd-dddddd-e implementation for Self-Extend extrapolation" The Mamba 432M model has model width of 1024 and the Mamba 1.3B model has model width of 2048. All models trained on SlimPajama have the same training configurations and the MLP intermediate size Samba, unless otherwise specified. The training infrastructure SlimPajama is based on a modified version of the TinyL codebase
Table 12:
Detailed lyper-parametent of the SAMBA models trained different scales. We only show optimization settings for the first training phase 3.8B model.
Total Parameters
421M
1.3B
1.7B
3.8B
Dataset
SlimPajama
SlimPajama
Phi-2
Phi-3
Batch size
512
512
2048
2048
Learning rate
0.0004
0.0004
0.0006
0.0006
Total training tokens
20B
100B
230B
3.2T
Weight decay
0.1
0.1
0.1
0.1
Gradient clipping
1.0
1.0
1.0
1.0
Sequence length
4096
4096
4096
4096
Sliding window size, w
2048
2048
2048
Number of layers,
24
36
48
64
Model width, dm
1536
2304
2048
2816
MLP intermediate size,
4096
6144
8196
9984
Number query heads
32
11
Number of key-value heads
12
1
Vocabulary size
32000
32000
50304
32064
In the generation configurations for the downstream tasks, we use greedy decoding for GSM8K, and Nucleus Sampling (Holtzman et al., 2019) with a temperature of r 0.2 and top-p = 0.95 for HumanEyal. For MBPP and SQuAD, we set 0.01 and top-p 0.95.
hhttps://githuu.coo//// htttess//iittuu..cc////aaa//////
24

---

Published as a conference paper at ICLR 2025
H LIMITATIONS & BROADER IMPACT
Although Samba demonstrates promising memory retrieval performance through instruction tuning. its pre-trained base model has retrieval performance similar to that of the SWA-based model, as shown in Figure 8. This opens up future direction on further improving the Samba's retrieval ability without compromising its efficiency and extrapolation ability. In addition, the hybridization strategy of Samba is not consistently better than other alternatives in all tasks. As shown in Table 2, Mamba-SWA-MLA. shows improved performance on tasks such as WinoGrande, SIQA, and GSM8K. This gives us the potential to invest in a more sophisticated approach to perform input-dependent dynamic combinations of SWA-based and SSM-based models (Ren al., 2023). With the improved short-context performance and the long-term memorization ability of linear complexity LLMs such as Samba, cost-effective applications can be developed for personalized learning and automated tutoring. Samba also be used for emotional accompaniment. The efficiency of the Samba architecture can save inference energy costs for models deployed on the edges, resulting in greener and more sustainable Al applications.
25


Reviews:
[{"review_id": "j5O6ZNCPjc", "reviewer": "Reviewer_A57R", "summary": "This paper proposes a hybrid architecture mixing Mamba layer with Sliding Window Attention.", "questions": "- In the original Mamba paper, they already experiment with mixed Mamba/Attention. Can you contextualize the your contribution with respect to this original architecture?\n\n- What is the evaluation setup for the summarization in Appendix B? Also, please bold the best performance on SQUALITY.\n\n- I’m having difficulty identifying the key factors that contribute to the optimal architectural design. What elements are beneficial here? Is it the MLP, Mamba, or both, and to what extent? I’d be interested in seeing results for a SWA > MLP > SWA > MLP architecture as an example (where some SWA layers are substituted with MLP).\n\n- In the length extrapolation experiments, can you clarify whether the models can extrapolate up to 256k tokens in Passkey Retrieval, but struggle with generalization in Phonebook performance?\n\n- In several sections, you mention that Full Attention performance deteriorates beyond a specific context size. This seems logical, couldn't these attention be converted to sliding attention especially l. 444?\n\n- In addition to the training curves shown in Fig. 7 of Appendix D, could you also provide the validation curves?\n\n- In Section 4, I don’t fully understand the relevance of the question: \"How to train models with Sliding Window Attention (SWA)?\"\n\n- This question might benefit from rephrasing, since your response is exploratory rather than definitive, focusing on the distribution of attention rather than giving a clear answer.\n\n- Line 051: The phrase \"limited context extrapolation ability in perplexity metrics\" is unclear. Are you suggesting that the extrapolation metrics increase perplexity when the models exceed their initial training context size?\n\n- Line 164 implies that SWA was specifically designed to address Mamba’s limitations, this is misleading.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The proposed architecture shows great results on common benchmark. \n- It seems also to be better at long contexts handling.\n- Shows that hybrid models can be competitive alternatives to transformers.", "weaknesses": "- The main contribution of combining attention with other linear mechanisms is not novel, and, as noted in the paper, a lot of alternatives exist.\n- A comprehensive benchmarking against existing alternatives is lacking. Comparisons are only made to their proposed variants and Sliding Window Attention in fair setups. A thorough comparison with other models listed in Appendix A (such as MEGA, adapted to Mamba) would strengthen the findings. Additionally, selected architectures are evaluated on a very small scale and only measured by perplexity. While some models achieve lower perplexity, this alone may not suffice to establish superiority (e.g., in H3 by Dao et al., 2022b, lower perplexity is reported against transformer baselines).\n- Results on common benchmarks are somewhat misleading. The paper aims to showcase the architecture’s strengths, yet comparisons are often made against models trained on different data distributions, which weakens the robustness of the conclusions.\n- Conclusions on long-context handling remain vague, although this should be a key advantage over transformers. It would be helpful to include dataset statistics (average, median, min, max lengths) to clarify context length relevance.\n- The only substantial long-context experiment, the summarization task, should be included in the main paper, with clearer discussion and analysis.\n- Section 4, “Analysis,” could benefit from clearer motivation. Some explored dimensions may appear intuitive (e.g., l. 444, where SWA is shown to outperform full attention on larger sequence lengths than those used in training), which might limit the novelty of the findings. Other questions seems a bit unrelated to the paper topics (see Questions).\n- Length extrapolation, a key aspect of the paper, is barely motivated or discussed in relation to prior work.\n- The paper overall feels somewhat unstructured and difficult to follow. Tables present different baselines inconsistently, and messages regarding architectural advantages are interleaved with comments on training data quality (l. 229). The evaluation setup lacks consistency (performance is sometimes assessed on real benchmarks, other times by perplexity), and the rationale behind baseline choices or research questions is insufficiently explained.", "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": "XYTwQ3OUod", "reviewer": "Reviewer_4dj2", "summary": "Authors investigate the utility of attention-state-space hybrids over attention-only models (e.g. sliding window attention (SWA)) and state-space-only models (e.g. Mamba) by training language models in the 400M-4B parameter range. They propose a model block Mamba->MLP->SWA->MLP and perform ablation studies in a language modeling (LM) setup. Pretrained models are evaluated 0-shot on 15 established downstream benchmarks.\n\n**Experiments**\n\n1. Authors perform an ablation study on the Phi-2 training corpus (230B tokens) and report modest improvements over the Mamba->SWA->MLP baseline. On reading comprehension tasks such as SQuAD, the addition of SWA to Mamba leads to significant improvement over the Mamba-only baseline.\n\n2. Authors explore replacing Mamba in the hybrid block with other efficient alternatives such as Gated Linear Attention and RetNet. Models are trained on SlimPajama corpus with 4k length and evaluated on lengths up to 16k, with the Mamba-based block reporting the lowest perplexity.\n\n3. Authors study the length extrapolation abilities of models trained with 4k length. On the ProofPile test set, the hybrid block shows improved and slowly decreasing perplexity with increasing context length, and is lower than that of the SWA-only baseline.\n\n4. Authors investigate the model's performance on memory-intensive tasks. After finetuning for 500 steps with 4k input length, the hybrid block achieves perfect performance on PassKey Retrieval for lengths up to 256k. On the PhoneBook task of retrieving phone numbers from name-number pairs, models finetuned for 100 steps with 4k input length show superior performance compared to full attention when evaluated on lengths up to 8k.\n\n5. Analysis of attention scores reveals that middle layers have lower entropy compared to top/bottom layers, suggesting specialization in precise information retrieval.\n\n6. The hybrid model trained on Phi-3's 3.2T token corpus matches parameter-matched Phi-3's performance on downstream tasks.", "questions": "1. In models with SSM layers, was RoPE still necessary in SWA? (This doesn't affect my evaluation)\n\n2. Are the Phi-2 and Phi-3 training sets open-sourced? Reproducibility is challenging without access to these datasets.\n\n3. Please add Mamba->SWA->MLP baseline to Table 3.\n\n4. Consider including Mistral results on PhoneBook.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "1. The investigation of SSM-attention hybrids is well motivated and timely. The empirical findings and model performance should interest a broad audience. The trained checkpoints would be a useful resource for the research community.\n\n2. Experiments are comprehensive, showing significant improvements in length extrapolation tasks and efficiency on long inputs.", "weaknesses": "1. Limited novelty - adding an MLP in the model block is an incremental contribution. The performance gap between Mamba->MLP->SWA->MLP and Mamba->SWA->MLP is modest.\n\n2. The models' strong performance relies on proprietary datasets (Phi-2 and Phi-3), making fair comparisons difficult.\n\n3. The relationship with previous SSM+SWA hybrid works (Griffin/Hawk) needs clearer positioning.", "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": "x8a1coIDbV", "reviewer": "Reviewer_MPi8", "summary": "The paper describe the hybrid Language model which combined Mamba layers and Sliding Windows Attention layers. Authors trained a number of models upto 3.8B size and showed that model performed well both on multiple standard LLM benchmarks and on some long context benchmarks (e.g. Passkey retrieval ). They also evaluated model on  Phonebook to show model limitation. Samba has better throughput comparing to regular Transformers both for  sequences with long prompt, and on generation long context.", "questions": "It would be interesting to see if Samba is able to handle \n a) more complex Needle-in-the-Haistack from RULER benchmark https://arxiv.org/abs/2404.06654 \n b) complex multi--hop  reasoning tasks from HELMET benchmark https://arxiv.org/pdf/2410.02694", "limitations": "", "rating": 8, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Authors build a good LM which can both process long prompt and generate very long context  reliably. The simple idea - to combine mamba with SWA -  is very well executed, Solid experimental part, good ablation study. The paper is well written.", "weaknesses": "The main weakness of  Samba, is that similar to other hybrid models which use local attention it solves long context problem only partially. Experiments on Phonebooks show that the accuracy of the model rapidly goes down after 4K (Fig. 5). \n\nI would also suggest to move the following observation from Appendix G  to main part of the paper, e.g. to Analysis or Conclusion section:  \n\"Although Samba demonstrates promising memory retrieval performance through instruction tuning, its pre-trained base model has retrieval performance similar to that of the SWA-based model ... \"\n \nA few minor comments:\n1. Inconsistent location of captions for figures and tables (sometimes above , sometimes below)\\. Figure 3 is located inside text. \n2. The implementation details: RoPE base is missing", "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)