Samba: Simple Hybrid State Space Models for Efficient Unlimited Context Language Modeling

## Abstract

Abstract Efficiently modeling sequences with infinite context length has been a long-standing problem. Past works suffer from either the quadratic computation complexity or the limited extrapolation ability on 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 memories with the attention mechanism.
We scale Samba up to 3.8B parameters with 3.2T training tokens and show that Samba substantially outperforms the state-of-the-art models based on pure attention or SSMs on a wide range of benchmarks. When trained on 4K length sequences, Samba can be efficiently extrapolated to 256K context length with perfect memory recall and show improved token predictions up to 1M context length. As a linear-time sequence model, Samba enjoys a 3.73 × 3.73\times 3.73 × higher throughput compared to Transformers with grouped-query attention when processing user prompts of 128K length, and 3.64 × 3.64\times 3.64 × speedup when generating 64K tokens with unlimited streaming. A sample implementation of Samba is publicly available in https://github.com/microsoft/Samba .

## 1 Introduction

Figure: (a) Perplexity on the test set of Proof-Pile
Refer to caption: x1.png

Attention-based models [, ] have dominated the neural architectures of Large Language Models (LLMs) [, , , ] due to their ability to capture complex long-term dependencies and the efficient parallelization for large-scale training [].
Recently, State Space Models (SSMs) [, , , ] have emerged as a promising alternative, offering linear computation complexity and the potential for better extrapolation to longer sequences than seen during training. Specifically, Mamba[], a 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 [].
However, SSMs struggle with memory recall due to their Markovian nature [], and experimental results on information retrieval-related tasks [, , ], have further shown that SSMs are not as competitive as their attention-based counterparts.

Previous works [, , , ] have explored different approaches to hybridize SSMs and the attention mechanism, but none of them achieve unlimited-length extrapolation with linear-time complexity. The existing length generalization techniques [, , ] developed for the attention mechanism suffer from quadratic computation complexity or limited context extrapolation ability.
In this paper, we introduce Samba, a simple neural architecture that harmonizes the strengths of both the SSM and the attention-based models, while achieving an unlimited sequence length extrapolation with linear time complexity.
Samba combines SSMs with attention through layer-wise interleaving Mamba [], SwiGLU [], and Sliding Window Attention (SWA) []. Mamba layers capture the time-dependent semantics and provide a backbone for efficient decoding, while SWA fills in the gap modeling complex, non-Markovian dependencies.

We scale Samba with 421M, 1.3B, 1.7B and up to 3.8B parameters. In particular, the largest 3.8B base model pre-trained with 3.2T tokens achieves a 71.2 score for MMLU [], 54.9 for HumanEval [], and 69.6 for GSM8K [], substantially outperforming strong open source language models up to 8B parameters, as detailed in Table 1. Despite being pre-trained in the 4K sequence length, Samba can be extrapolated to 1M length in zero shot with improved perplexity on Proof-Pile [] while still maintaining the linear decoding time complexity with unlimited token streaming, as shown in Figure 1.
We show that when instruction-tuned in a 4K context length with only 500 steps, Samba can be extrapolated to a 256K context length with perfect memory recall in Passkey Retrieval []. In contrast, the fine-tuned SWA-based model simply cannot recall memories beyond 4K length. We further demonstrate that the instruction-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.
Finally, we conduct rigorous and comprehensive analyzes and ablation studies, encompassing up to 1.7 billion parameters, to validate the architectural design of Samba. These meticulous investigations not only justify our architectural designs but also elucidate the potential mechanisms underpinning the remarkable effectiveness of this simple hybrid approach.

## 2 Methodology

We explore different hybridization strategies consisting of the layers of Mamba, Sliding Window Attention (SWA), and Multi-Layer Perceptron [, ]. We conceptualize 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 [] and GLA [] 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 2, we explore three kinds of layerwise hybridization strategies on the 1.7B scale: Samba, Mamba-SWA-MLP, 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-MLP, and Mamba, while Mamba-SWA-MLP has 54 layers, so each model has approximately 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.

Figure: Figure 2: From left to right: Samba, Mamba-SWA-MLP, Mamba-MLP, 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 [, ] and skip connections [] are applied for each of the intermediate layers.
Refer to caption: extracted/5657078/Samba.png

#### 2.1.1 Mamba Layer

Mamba [] 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 a soft selection of the input sequence elements. Given an input sequence representation $\mathbf{X}\in\mathbb{R}^{n\times d_{m}}$, where $n$ is the length of the sequence and $d_{m}$ is the hidden size, Mamba first expands the inputs to a higher dimension $d_{e}$, *i.e.*,

$$ $\mathbf{H}=\mathbf{X}\mathbf{W}_{\text{in}}~{}\in\mathbb{R}^{n\times d_{e}}$ $$

where $\mathbf{W}_{\text{in}}\in\mathbb{R}^{d_{m}\times d_{e}}$ is a learnable projection matrix. Then a Short Convolution (SC) [] operator is applied to smooth the input signal,

$$ $\mathbf{U}=\text{SC}(\mathbf{H})=\text{SiLU}(\text{DepthwiseConv}(\mathbf{H}, \mathbf{W}_{\text{conv}}))~{}\in\mathbb{R}^{n\times d_{e}}$ (1) $$

where $\mathbf{W}_{\text{conv}}\in\mathbb{R}^{k\times d_{e}}$ and the kernel size $k$ is set to 4 for hardware-aware efficiency. The Depthwise Convolution [] is applied over the sequence dimension followed by a SiLU [] activation function. The selective gate is then calculated through a low-rank projection followed by Softplus [],

$$ $\Delta=\text{Softplus}(\mathbf{U}\mathbf{W}_{\text{r}}\mathbf{W}_{\text{q}}+ \mathbf{b})~{}\in\mathbb{R}^{n\times d_{e}}$ (2) $$

where $\mathbf{W}_{\text{r}}\in\mathbb{R}^{d_{e}\times d_{r}}$, $\mathbf{W}_{\text{q}}\in\mathbb{R}^{d_{r}\times d_{e}}$ and $d_{r}$ is the low-rank dimension. $\mathbf{b}\in\mathbb{R}^{d_{e}}$ is carefully initialized so that $\Delta\in[\Delta_{\text{min}},\Delta_{\text{max}}]$ after the initialization stage. We set $[\Delta_{\text{min}},\Delta_{\text{max}}]=[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 $\mathbf{B}$ and $\mathbf{C}$ of SSM,

$$ $\mathbf{B}=\mathbf{U}\mathbf{W}_{\text{b}}~{}\in\mathbb{R}^{n\times d_{s}}$ $$

$$ $\mathbf{C}=\mathbf{U}\mathbf{W}_{\text{c}}~{}\in\mathbb{R}^{n\times d_{s}}$ $$

where $d_{s}$ is the state dimension. For each time step $1\leq t\leq n$, the recurrent inference of the Selective SSM (S6) is performed in an expanded state space $\mathbf{Z}_{t}\in\mathbb{R}^{d_{e}\times d_{s}}$, *i.e.*,

$$ $\mathbf{Z}_{t}=\exp(-\Delta_{t}\odot\exp(\mathbf{A}))\odot\mathbf{Z}_{t-1}+ \Delta_{t}\odot(\mathbf{B}_{t}\otimes\mathbf{U}_{t})~{}\in\mathbb{R}^{d_{e} \times d_{s}}$ $$

$$ $\mathbf{Y}_{t}=\mathbf{Z}_{t}\mathbf{C}_{t}+\mathbf{D}\odot\mathbf{U}_{t}~{} \in\mathbb{R}^{d_{e}}$ $$

where $\mathbf{Z}_{0}=\mathbf{0}$, $\odot$ means the point-wise product, $\otimes$ means the outer product and $\exp$ means the point-wise natural exponential function. $\mathbf{D}\in\mathbb{R}^{d_{e}}$ is a learnable vector initialized as $D_{i}=1$ and $\mathbf{A}\in\mathbb{R}^{d_{e}\times d_{s}}$ is a learnable matrix initialized as $A_{ij}=\log(j),1\leq j\leq d_{s}$, following the S4D-Real [] initialization. In practice, Mamba implements a hardware-aware parallel scan algorithm for efficient parallelizable training.
The final output is obtained through a gating mechanism similar to Gated Linear Unit [, ],

$$ $\mathbf{O}=\mathbf{Y}\odot\text{SiLU}(\mathbf{X}\mathbf{W}_{\text{g}})\mathbf{ W}_{\text{out}}\in\mathbb{R}^{n\times d_{m}}$ $$

where $\mathbf{W}_{g}\in\mathbb{R}^{d_{m}\times d_{e}}$ and $\mathbf{W}_{\text{out}}\in\mathbb{R}^{d_{e}\times d_{m}}$ are learnable parameters. In this work, we set $d_{e}=2d_{m}$, $d_{r}=d_{m}/16$, and $d_{s}=16$. The Mamba layer in Samba is expected to capture the time-dependent 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.

#### 2.1.2 Sliding Window Attention (SWA) Layer

The Sliding Window Attention [] layer is designed to address the limitations of the Mamba layer in capturing non-Markovian dependencies in sequences. Our SWA layer operates on a window size $w=2048$ that slides over the input sequence, ensuring that the computational complexity remains linear with respect to the sequence length. The RoPE [] relative positions are applied within the sliding window. 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 [] for the efficient implementation of self-attention throughout this work. We also choose the 2048 sliding window size for efficiency consideration; FlashAttention 2 has the same training speed as Mamba’s
selective parallel scan at the sequence length of 2048 based on the measurements in [].

#### 2.1.3 Multi-Layer Perceptron (MLP) Layer

The MLP layers in Samba serve as the architecture’s primary mechanism for nonlinear transformation and recall of factual knowledge []. We use SwiGLU [] for all the models trained in this paper and denote its intermediate hidden size as $d_{p}$. As shown in Figure 2, 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 10 of Appendix A. We also train other hybrid architectures as mentioned in Section 2.1, including the baseline Mamba, Llama-3, and Mistral architecture on a scale of around 1.7B, with detailed hyperparameters in Table 9 of Appendix A. We do comprehensive downstream evaluations on a wide range of benchmarks, focusing on four main capabilities of the models: commonsense reasoning (ARC [], PIQA [], WinoGrande [], SIQA []), language understanding (HellaSwag [], BoolQ [], OpenbookQA [], SQuAD [], MMLU []), truthfulness (TruthfulQA []) and math and coding (GSM8K [], MBPP [], HumanEval []).

### 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 [] with 3.2T tokens. We follow the same multi-phase pretraining strategy as Phi3-mini for a fair comparison. We also report the performance of the Transformer++ (TFM++ in Table 1) model, which uses the same architecture and training recipe as Phi3-mini, for a fair comparison.
In Table 1, we conduct comprehensive evaluations on a diverse subset of the benchmarks to assess Samba’s performance across all the domains mentioned above to ensure a thorough examination of the model’s capabilities. The details of the generation configurations are included in Appendix A.

**Table 1: Downstream performance comparison of Samba 3.8B with other pretrained base language models without instruction tuning. ARC-C and HellaSwag are measured with character-normalized accuracy. MMLU and GSM8K are measured in 5-shot, while others are in zero-shot. We report the MC2 score for TruthfulQA, maj@1 for GSM8K, and pass@1 for HumanEval. ^∗ Measured by ours.**
| Model | Size | Tokens | MMLU | Hella- | ARC- | Wino- | Truth. | GSM | Hum. | Avg. |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Swag | C | Gran. | QA | 8K | Eval |  |  |  |  |  |
| Llama 2 | 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 3 | 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 |

We compare with several strong baselines, including Llama 2 [], Mistral [], Mamba [], Gemma [], Recurrent-Gemma (R-Gemma) [], Llama 3 [] and TFM++. As shown in Table 1, Samba achieves the 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 on the same dataset. This shows the surprising complementary effect of combining SSM with the attention mechanism. We conjecture that when combined with attention, Mamba, as an input-dependent SSM, can focus more on performing the arithmetic operation through its recurrent states than on doing the retrieval operation which can be easily learned by the sliding window attention.

**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 in bold, second best underlined.**
| Benchmark | Llama-3 | Mistral | Mamba | Mamba-SWA- | Mamba- | Samba |
| --- | --- | --- | --- | --- | --- | --- |
| 1.6B | 1.6B | 1.8B | MLP 1.6B | MLP 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 [] dataset with 230B tokens and evaluate them in the full suite of 15 downstream benchmarks to have a 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 (ARC-Challenge), language understanding (MMLU, 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.
We can observe that replacing Mamba blocks with MLPs does not harm commonsense reasoning ability, but its performance on 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 as 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 Mamba-SWA-MLP has significantly better performance on GSM8K, potentially resulting from a closer collaboration between the Mamba and the SWA layers. The distinct downstream performances of different hybridization strategies pose interesting future work for developing task-adaptive dynamic architectures.

### 3.2 Exploration on Attention and Linear Recurrence

Since SSMs belong to a broader realm of linear recurrent models [, , , , ], there exist multiple alternatives other than Mamba when combing attention-based layers with recurrent neural networks. In addition to Mamba and Samba, we investigate the comparative analysis of the following architectures:

- •
Llama-2 [] is an attention-based Transformer architecture that utilizes full self-attention across the entire sequence.
- •
Llama-2-SWA is an attention-based architecture that replaces all full attention layers in Llama-2 with sliding window attention.
- •
Sliding RetNet replaces Mamba layers in the Samba architecture with Multi-Scale Retention  [] 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) []. GLA is a more expressive variant of linear attention with input-dependent gating.

**Table 3: 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 a fluctuation around $\pm 0.3\%$.**
| Architecture | Size | Layers | Training Speed | Validation Context Length |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| ($\times 10^{5}$ tokens/s) | 4096 | 8192 | 16384 |  |  |  |
| *20B training tokens on 8$\times$A100 GPUs* |  |  |  |  |  |  |
| Llama-2 | 438M | 24 | 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 | 438M | 24 | 4.32 | 10.38 | 9.96 | 9.87 |
| Samba | 421M | 24 | 4.46 | 10.06 | 9.65 | 9.57 |
| *100B training tokens on 64$\times$H100 GPUs* |  |  |  |  |  |  |
| Llama-2 | 1.3B | 40 | 25.9 | 7.60 | 44.32 | 249.64 |
| Llama-2-SWA | 1.3B | 40 | 26.2 | 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 |
| Samba | 1.3B | 36 | 25.2 | 7.32 | 7.11 | 6.96 |

We pre-train all models on the same SlimPajama [] dataset under both around 438M and 1.3B settings, and evaluate these models by calculating perplexity on the validation set with context length at 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 in Appendix A.
As shown in Table 3, Samba consistently outperforms all other models in different context lengths and model sizes.
The training speed of Samba is 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. We can notice that the full attention-based model cannot extrapolate beyond its context length without specific length extrapolation techniques, which motivates us to use SWA for Samba. In Section 4, we further show that even hybridizing with one full attention layer will still lead to exploding perplexity at 16k sequence length. We can also find that while RetNet can extrapolate well under the 438M scale, it has an increasing perplexity on 16K length at the 1.4B scale, which may indicate that its input-independent decay may need specific tuning at different scales to work well.

### 3.3 Efficient Length Extrapolation

Figure: Figure 3: Prompt processing throughput of different models with around 1.7B parameters.
Refer to caption: extracted/5657078/speed.jpg

We use the test split of the Proof-Pile [] dataset to evaluate the length extrapolation ability of our models at a scale of around 1.7B parameters. We follow Position Interpolation [] for data pre-processing. The sliding window approach [] is used for the perplexity evaluation with a window size of 4096.
Besides having the decoding throughput in Figure 1 for the generation efficiency metric, we also measure the prompt processing speed in Figure 3 for the models Samba 1.7B, Mistral 1.6B, Mamba 1.8B, Llama-3 1.6B and its Self-Extended [] version SE-Llama-3 1.6B with the prompt length sweeping from 1K 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 to be 128K and varying the batch size accordingly.
The throughput is measured on a single A100 GPU with the precision of bfloat16. We repeat the measurements 10 times and report the averaged results.
We can see that Samba achieves $3.73\times$ 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 1, 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.

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-tune Samba 1.7B on Passkey Retrieval with a 4K training sequence length for only 500 steps.
As presented in Figure 4, 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-tuning procedure in Figure 6 and Figure 7 of Appendix B. We find that despite the fact that both architectures can reach near-zero training loss in less than 250 steps, Samba can achieve near-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.

Figure: Figure 4: 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.
Refer to caption: x2.png

### 3.4 Long-Context Understanding

The impressive results on the synthetic passkey retrieval task encourage us to perform full-cycle instruction tuning of the Samba-3.8B model. We follow the same post-training recipe used for the Phi-3-mini series and evaluate the downstream performance of the instruction-tuned Samba-3.8B-IT (preview) on both the long-context summarization tasks (GovReport [], SQuALITY []) and the main short-context benchmarks (MMLU, GSM8K, HumanEval), as shown in Table 4. We can see that Samba has substantially better performance than Phi-3-mini-4k-instruct on both the short-context (MMLU, GSM8K, HumanEval) and long-context (GovReport) tasks, while still having the 2048 window size of its SWA layer and maintaining the linear complexity for efficient processing of long documents.

**Table 4: Downstream performance comparison between instruction-tuned Samba 3.8B 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 [] for GSM8K, 0-shot pass@1 for HumanEval, ROUGE-L for both GovReport and SQuALITY. † Results from the Phi-3 technical report [].**
| Model | MMLU | GSM8K | HumanEval | GovReport | SQuality |
| --- | --- | --- | --- | --- | --- |
| Phi-3-mini-4K-instruct † | 68.8 | 82.5 | 58.5 | 14.4 | 21.6 |
| Samba-3.8B-IT (preview) | 71.9 | 87.6 | 62.8 | 18.9 | 21.2 |

## 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 $\pm 0.3\%$. Training speed is measured on 8$\times$H100 GPUs by default. All the models in this section are trained on SlimPajama with 20B tokens and 4K sequence length, unless otherwise specified.

##### How to train models with Sliding Window Attention (SWA)?

Since SWA has linear complexity with respect to the sequence length, it seems alluring to trade off the batch size to have a longer training sequence length without substantially decreasing the training throughput. However, as shown in Table 5, when the sequence length is increased, the validation perplexity also increases in all context lengths due to smaller batch sizes, and the optimal ratio of sequence length/window size observed is 2, resulting in a training length of 4096.

**Table 5: Perplexity on SlimPajama of Llama-2-SWA 438M models trained on different context sizes and batch sizes. We fix the sliding window size as 2048 and the training tokens per step as 2M.**
| Batch Size | Sequence Length | Training Speed | Validation Context Length |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| ($\times 10^{5}$ 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 |

##### Why not hybridize with full attention?

Some previous works [, ] suggest a hybrid architecture of Mamba with full attention. However, as shown in Table 6, 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.
Samba also has much better training throughput compared to Mamba-MLP alternatives because self-attention with the FlashAttention 2 implementation is more training efficient than Mamba when the sequence length is 4096.

**Table 6: Perplexity on SlimPajama of Mamba-MLP architectures with full attention layers replacing Mamba layers at different block indices. We define a block as two consecutive layers with a Mamba/Attention layer followed by an MLP. All the models have 12 blocks in total.**
| Architecture | Size | Block Index | Training Speed | Validation Context Length |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| of Full Attention | ($\times 10^{5}$ tokens/s) | 4096 | 8192 | 16384 |  |  |
| Mamba-MLP | 449M | 11 | 7.78 | 10.29 | 10.53 | 13.66 |
| 449M | 5 | 7.78 | 10.10 | 10.05 | 12.83 |  |
| 449M | 0 | 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 7, we explore the techniques of query head grouping [, ], for both the Llama and Samba models. Surprisingly, both the Llama-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 to the contexts. We can also see that Samba has a $2\times$ smaller optimal number of query heads than the SWA model, which confirms our hypothesis that Samba can support a smaller number of attention heads.

**Table 7: Perplexity on SlimPajama of Llama-2-SWA and Samba models at the 430M scales trained with different number of Query and Key-Value heads. “KV Size” means the size of Key-Value vectors per token. Since grouped query attention will reduce the parameters for attention from $4d_{m}^{2}$ to roughly $2d_{m}^{2}$, we increase the hidden size of MLP from $8/3d_{m}$ to $3d_{m}=4608$ to have roughly the same number of total parameters as the original models.**
| Query | Key-Value | Head | KV | Model | Training Speed | Validation Context Length |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Head | Head | Dim. | Size | Size | ($\times 10^{5}$ tokens/s) | 4096 | 8192 | 16384 |
| *Llama-2-SWA Architecture* |  |  |  |  |  |  |  |  |
| 12 | 2 | 128 | 512 | 419M | 10.01 | 11.11 | 10.64 | 10.56 |
| 6 | 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 |
| 6 | 1 | 256 | 512 | 426M | 8.46 | 9.99 | 9.59 | 9.51 |
| 12 | 1 | 128 | 256 | 424M | 8.62 | 10.07 | 9.66 | 9.58 |
| 12 | 4 | 128 | 1024 | 431M | 8.57 | 10.02 | 9.62 | 9.55 |

##### Why hybrid is better?

We examine the entropy of the attention distributions for both the Samba 1.7B and the Mistral 1.6B models. As shown in Figure 5a, the Samba model has a 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-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, 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 of 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 C.

Figure: (a) Average attention entropy per decoding step
Refer to caption: x3.png

**Table 8: 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$\times$A100 GPUs.**
| Architecture | Size | Training Speed | Validation Context Length |  |  |
| --- | --- | --- | --- | --- | --- |
| ($\times 10^{5}$ tokens/s) | 4096 | 8192 | 16384 |  |  |
| Llama-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 | 438M | 4.32 | 10.38 | 9.96 | 9.87 |
| *+ SC* | 438M | 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 1 is independent to the design of other parts of Mamba and can be applied to other linear recurrent models. As shown in Table 8, we explore the effect of SC on model performance through enhancing Llama-2-SWA, 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 GLA already has the fine-grained decays at the channel level, so the depthwise convolution doesn’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 421M’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 as a future work to understand the surprising effectiveness of SC in language modeling.

## 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-based 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 to the state-of-the-art Transformer architecture. The architecture’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 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.

## Acknowledgement

We want to thank Shuohang Wang and Liyuan 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.

## Appendix A Implementation Details

**Table 9: Detailed hyper-parameters of the baselines models trained on the Phi2 dataset with 230B tokens.**
| Architecture | Llama-3 | Mistral | Mamba | Mamba-SWA-MLP | Mamba-MLP |
| --- | --- | --- | --- | --- | --- |
| 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, $w$ | - | 2048 | - | 2048 | - |
| Number of layers, $N$ | 48 | 48 | 64 | 54 | 48 |
| Model width, $d_{m}$ | 2048 | 2048 | 2048 | 2048 | 2048 |
| MLP intermediate size, $d_{p}$ | 8196 | 8196 | - | 8196 | 8196 |
| Number of query heads | 32 | 32 | - | 32 | 32 |
| Number of KV heads | 4 | 4 | - | 4 | 4 |
| Number of Attention Layers | 24 | 24 | 0 | 18 | 0 |
| 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 $d_{m}/384$, a key expansion ratio of 0.5, and a value expansion ratio of 1. For the RetNet layer we use a number of head that is half of the number of attention query heads, key expansion ratio of 1 and value expansion ratio of 2. The GLA and RetNet implementations are from the Flash Linear Attention repository(^1^11[https://github.com/sustcsonglin/flash-linear-attention](https://github.com/sustcsonglin/flash-linear-attention)) []. We use the FlashAttention-based implementation for Self-Extend extrapolation(^2^22[https://github.com/datamllab/LongLM/blob/master/self_extend_patch/Llama.py](https://github.com/datamllab/LongLM/blob/master/self_extend_patch/Llama.py)). The Mamba 432M model has a model width of 1024 and the Mamba 1.3B model has a model width of 2048. All models trained on SlimPajama have the same training configurations and the MLP intermediate size as Samba, unless otherwise specified. The training infrastructure on SlimPajama is based on a modified version of the TinyLlama codebase(^3^33[https://github.com/jzhang38/TinyLlama](https://github.com/jzhang38/TinyLlama)).

**Table 10: Detailed hyper-parameters of the Samba models trained at different scales. We only show the optimization settings for the first training phase of the 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 | 2048 |
| Number of layers, $N$ | 24 | 36 | 48 | 64 |
| Model width, $d_{m}$ | 1536 | 2304 | 2048 | 2816 |
| MLP intermediate size, $d_{p}$ | 4096 | 6144 | 8196 | 9984 |
| Number of query heads | 12 | 18 | 32 | 11 |
| Number of key-value heads | 12 | 18 | 4 | 1 |
| Vocabulary size | 32000 | 32000 | 50304 | 32064 |

In the generation configurations for the downstream tasks, we use greedy decoding for GSM8K, and Nucleus Sampling [] with a temperature of $\tau=0.2$ and $\text{top-}p=0.95$ for HumanEval. For MBPP and SQuAD, we set $\tau=0.01$ and $\text{top-}p=0.95$.

## Appendix B Additional Experiment Results

Figure: Figure 6: 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.
Refer to caption: x5.png

Figure: Figure 7: Overall passkey retrieval accuracy on the 256K document length of Samba 1.7B and Mistral 1.6B models during 500 steps of instruction tuning.
Refer to caption: x6.png

We perform instruction tuning for both Mistral 1.6B and Samba 1.7B on 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 to 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 [] 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 at the document length of [4k, 8k, 16k, 32k, 64k, 128k, 256k], for each length we evaluate at 11 different depths of the document (from 0, 0.1, 0.2, … to 1.0). In addition, for each location of the passkey (depth) in the document, we evaluate the model with five different passkeys to measure accuracy. As seen in Figure 7, the average passkey retrieval accuracy for Samba 1.7B almost reaches 100% in around 150 steps, while the accuracy for Mistral 1.6B remains low, demonstrating the extrapolation ability of the Samba architecture.

## Appendix C Details of Entropy Measurement

Given a causal attention probability matrix $A\in\mathbb{R}^{h\times n\times n},A_{ijk}=0~{}\forall j<k$, with $h$ number of heads and a sequence length of $n$, and the generation length $0<l<n$, we calculate the average attention entropy per decoding step as follows,

$$ $\mathcal{H}_{a}=-\frac{1}{l\cdot h}\sum_{i=1}^{h}\sum_{j=n-l+1}^{n}\sum_{k=1}^ {n}A_{ijk}\log(A_{ijk}).$ $$

For the selective gate $\Delta\in\mathbb{R}^{n\times d_{e}}$ used by S6 in Equation 2 of the Mamba layers, we first normalize it to be in the simplex $[0,1]^{n\times d_{e}}$, *i.e.*,

$$ $\Delta^{\prime}=\frac{\Delta}{\sum_{i=1}^{n}\Delta_{i}}~{}\in[0,1]^{n\times d_ {e}}.$ $$

The average selection entropy of S6 throughout the entire sequence is then calculated as

$$ $\mathcal{H}_{s}=-\frac{1}{d_{e}}\sum_{j=1}^{d_{e}}\sum_{i=1}^{n}\Delta^{\prime }_{ij}\log(\Delta^{\prime}_{ij}).$ $$

## Appendix D Limitations

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 7. 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-MLP 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.

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
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?

- What is the evaluation setup for the summarization in Appendix B? Also, please bold the best performance on SQUALITY.

- 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).

- 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?

- 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?

- In addition to the training curves shown in Fig. 7 of Appendix D, could you also provide the validation curves?

- In Section 4, I don’t fully understand the relevance of the question: "How to train models with Sliding Window Attention (SWA)?"

- 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.

- 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?

- Line 164 implies that SWA was specifically designed to address Mamba’s limitations, this is misleading.
soundness: 3
strengths: - The proposed architecture shows great results on common benchmark. 
- It seems also to be better at long contexts handling.
- Shows that hybrid models can be competitive alternatives to transformers.
summary: This paper proposes a hybrid architecture mixing Mamba layer with Sliding Window Attention.
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.
- 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).
- 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.
- 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.
- The only substantial long-context experiment, the summarization task, should be included in the main paper, with clearer discussion and analysis.
- 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).
- Length extrapolation, a key aspect of the paper, is barely motivated or discussed in relation to prior work.
- 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.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 4
questions: 1. In models with SSM layers, was RoPE still necessary in SWA? (This doesn't affect my evaluation)

2. Are the Phi-2 and Phi-3 training sets open-sourced? Reproducibility is challenging without access to these datasets.

3. Please add Mamba->SWA->MLP baseline to Table 3.

4. Consider including Mistral results on PhoneBook.
soundness: 4
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.

2. Experiments are comprehensive, showing significant improvements in length extrapolation tasks and efficiency on long inputs.
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.

**Experiments**

1. 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.

2. 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.

3. 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.

4. 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.

5. Analysis of attention scores reveals that middle layers have lower entropy compared to top/bottom layers, suggesting specialization in precise information retrieval.

6. The hybrid model trained on Phi-3's 3.2T token corpus matches parameter-matched Phi-3's performance on downstream tasks.
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.

2. The models' strong performance relies on proprietary datasets (Phi-2 and Phi-3), making fair comparisons difficult.

3. The relationship with previous SSM+SWA hybrid works (Griffin/Hawk) needs clearer positioning.

---

confidence: 5
contribution: 3
initial_rating: 8
presentation: 3
questions: It would be interesting to see if Samba is able to handle 
 a) more complex Needle-in-the-Haistack from RULER benchmark https://arxiv.org/abs/2404.06654 
 b) complex multi--hop  reasoning tasks from HELMET benchmark https://arxiv.org/pdf/2410.02694
soundness: 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.
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.
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). 

I would also suggest to move the following observation from Appendix G  to main part of the paper, e.g. to Analysis or Conclusion section:  
"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 ... "
 
A few minor comments:
1. Inconsistent location of captions for figures and tables (sometimes above , sometimes below)\. Figure 3 is located inside text. 
2. The implementation details: RoPE base is missing

Meta-review: The paper explores the design of state space models (SSMs) that integrate components from Mamba, sliding window attention, and MLPs. The objective is to develop efficient architectures for language modeling that achieve linear time complexity, effectively handle long contexts, and support considerable-length extrapolation. This is primarily an experimental study demonstrating high-level performance across a range of benchmarks, surpassing state-of-the-art (SOTA) SSMs and delivering substantial speedups compared to transformer baselines. Additionally, the authors provide a comprehensive analysis of the model, justifying their design choices.

The reviewers acknowledge that while the technical novelty may be limited, the proposed model delivers notable improvements over SOTA SSMs and transformers. It also offers valuable experimental insights that represent a significant contribution to the community. During the rebuttal phase, the authors included new experiments addressing the reviewers’ initial requests. I recommend acceptance.

Decision: Accept (Poster)