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
Primary area: foundation or frontier models, including LLMs
Keywords: Large Language Models;Length Extrapolation;Efficiency;Hybrid State Space Models

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.

Decision: accept (poster)
Reviewer scores: 6, 6, 8

Reviews and discussion:
<thread>
creation date:
1729971283905

thread_type:
official_review

<thread_object>
creator: reviewer_mpi8

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

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

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

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
</content>
</thread_object>

</thread>
<thread>
creation date:
1730563325542

thread_type:
official_review

<thread_object>
creator: reviewer_4dj2

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

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.

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.

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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731082963057

thread_type:
official_review

<thread_object>
creator: reviewer_a57r

<content>
summary:
This paper proposes a hybrid architecture mixing Mamba layer with Sliding Window Attention.

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.

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.

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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732852754778

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response

comment:
We thank the reviewers for their constructive feedback. We have revised the **PDF** and highlighted our revisions in purple. Major revisions are summarized in the following sections. 

## Updated Perplexity and Downstream Results

We train the following four alternative architectures with both around 400M and 1.3B parameters on the SlimPajama dataset to provide more comprehensive and reproducible benchmarking and ablation studies:

- Mega-S6: We replace all MD-EMA modules in the Mega architecture [1] with the ShortConv+S6 combinations from Mamba to adapt Mega to the modern Mamba architecture. Rotary position embedding, RMSNorm and Softmax attention are also adopted. We set the intermediate dimension of the Mega-S6 layer to be $d_m$ so that it has a roughly $5d_m^2$ number of parameters. This represents a classical baseline that conducts intra-layer SSM-Attention hybridization sequentially.

- Mamba-SWA-MLP: The Mamba->SWA->MLP architecture as in Figure 1 of our Paper.

- MLP2-SWA-MLP: We replace all Mamba layers in the Samba architecture to SwiGLU layers with $6d_m^2$ number of parameters.

- Samba-NoPE: We remove the rotary relative position embedding in Samba.

In the Table 3 of the revised **PDF** (also partially shown below), we can first notice that Mega-S6 shows significantly worse performances than other inter-layer hybridization strategies. Comparing Mamba-SWA-MLP with Samba, we can see that Samba has both slightly better perplexity scores and higher training throughput. Mamba-SWA-MLP trades off the MLP layers with more I/O intensive Mamba and Attention layers, leading to slower training speed. This also indicates that Mamba-SWA-MLP will have slower decoding speed than Samba due to larger total cache size resulting from more SSMs and Attention layers. We can observe that replacing Mamba with MLP in speeds up the training but harms perplexity significantly, indicating the importance of the Mamba layers in the Samba architecture. Interestingly, although we use SWA in Samba, Samba-NoPE has exploded perplexity beyond its training length without RoPE.

| **Architecture**    | **Size** | **Layers** | **Training Speed** | **Val. @ 4k** | **Val. @ 8k** |  **Val. @ 16k** |
|---|--|--|--|-|-|-|
| *20B training tokens on 8×A100 GPUs* |   |   |  |  |  |  |
| Mega-S6  | 422M | 24 | 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     | 24      | 4.46      | **10.06**  | **9.65**| **9.57** |
| *100B training tokens on 64×H100 GPUs* |     |   |    |    |  |   |
| Mega-S6      | 1.3B     | 36     | 17.9     | 9.01 | 8.81 | 8.68  |
| Mamba-SWA-MLP | 1.3B     | 36    | 23.5   | 7.37  | 7.16 | 7.00  |
| MLP2-SWA-MLP    | 1.3B     | 36     | **26.6**   | 7.81     | 7.58 | 7.42  |
| Samba-NoPE      | 1.3B     | 36     | 25.2  | 7.33    | 20.40| 326.17|
| Samba | 1.3B     | 36    | 25.2    | **7.32**  | **7.11**| **6.96** |

In the Table 4 of the revised **PDF** (also shown below), we evaluate all our ~1.3B scale models pre-trained with 100B tokens from SlimPajama on five typical commonsense reasoning tasks 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 a large margin. Similar to our perplexity evaluation, Samba and Samba-NoPE have similar average accuracies, whereas Mamba-SWA-MLP falls slightly behind. We observe that different architectures excel at different tasks. Mamba-SWA-MLP performs best on ARC-Easy, while Samba and Samba-NoPE achieve superior results on LAMBADA. Hybrid models based on Mamba generally outperform hybrid linear attention models and pure softmax-attention models on HellaSwag.
 


| **Architecture**      | **Size** | **ARC-Easy (acc ↑)** | **HellaSwag (acc_norm ↑)** | **WinoGrande (acc ↑)** | **PIQA (acc ↑)** | **LAMBADA (acc ↑)** | **Avg.** |
|--|--|--|--|--|--|-|--|
| LLaMA-2     | 1.3B    | 55.09    | 52.32  | 53.35    | 71.11   | 48.52    | 56.08    |
| LLaMA-2-SWA  | 1.3B    | 56.65  | 52.59   | 54.93      | 71.60  | 47.56   | 56.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-MLP   | 1.3B    | **59.64**    | 54.50  | 55.25  | **72.42**  | 49.12 | 58.19    |
| MLP2-SWA-MLP | 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** |
</content>
</thread_object>

</thread>


Metareview:
metareview:
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.