PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Liliang Ren, Yang Liu, Yadong Lu, Yelong Shen, Chen Liang, Weizhu Chen

arXiv:2406.07522v3cs.CLcs.LG

TL;DR

Efficient language modeling over unlimited context remains difficult because prior approaches face quadratic computation or weak length extrapolation. Samba addresses this with a hybrid of Mamba and Sliding Window Attention, achieving strong benchmark performance, long-context extrapolation, and efficient processing. Its retrieval ability and hybridization strategy nevertheless remain uneven across settings and tasks.

  • Problem

    Efficiently modeling sequences with unlimited context is challenging because prior approaches suffer from quadratic computation or limited length extrapolation.

  • Method

    Samba layer-wise combines Mamba, Sliding Window Attention, and SwiGLU to capture recurrent structure, retrieve memories, and transform information.

  • Results

    Samba substantially outperforms pure attention-based and SSM-based models across diverse benchmarks while providing long-context extrapolation and throughput gains.

  • Takeaways & Limitations

    Samba supports practical long-context language modeling by combining strong downstream performance with extrapolation to 256K retrieval contexts and efficient processing.

  • Takeaways & Limitations

    The pretrained base model has retrieval performance similar to SWA, and Samba is not consistently better than alternative hybrids across all tasks.

Abstract

from arXiv · show

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.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 on open source data is publicly available at https://github.com/microsoft/Samba.

1 INTRODUCTION

Samba combines selective state-space modeling with sliding-window attention to target efficient, long-context language modeling. The architecture is designed to retain recurrent sequence information while precisely retrieving recent memories, and it is evaluated against existing hybrid and Transformer approaches.

  • Motivation: Attention-based models capture complex long-term dependencies and parallelize efficiently, whereas SSMs offer linear computation and potential length extrapolation.
  • Motivation: Existing SSM-attention hybrids had not demonstrated significantly better language-modeling performance than state-of-the-art Transformers.
  • Architecture: Samba layer-wise interleaves Mamba, SwiGLU, and Sliding Window Attention to combine recurrent sequence modeling with non-recurrent dependency modeling.Mamba provides a backbone for efficient decoding, while SWA models complex dependencies that are not recurrent.
  • Motivation: Samba is presented as a linear-time architecture intended to address quadratic computation and insufficient context extrapolation in prior approaches.

2 METHODOLOGY

Samba combines Mamba, Sliding Window Attention, and MLP layers to retain recurrent sequence structure, retrieve recent information precisely, and recall factual knowledge. Its Mamba component uses input-dependent selective state updates, while SWA preserves linear complexity through a fixed sliding window.

  • Hybrid architecture: Samba layer-wise hybridizes Mamba, Sliding Window Attention, and MLPs to capture recurrent structures, retrieve memories precisely, and recall factual knowledge.The architecture assigns these roles respectively to Mamba, SWA, and MLP components.
  • Hybrid architecture: Samba explores layer-wise arrangements including Samba, Mamba-SWA-MLP, and Mamba-MLP while holding other configurations constant for approximately 1.7B-parameter comparisons.The models differ in layer-level arrangement, with N = 48 layers for Samba, Mamba-MLP, and Mamba, and 54 for Mamba-SWA-MLP.
  • Mamba layer: Mamba uses input-dependent gating and recurrent states to select sequence elements and memorize relevant information over the long term.Its selective SSM processes inputs through recurrent inference in an expanded state space and uses a final gating mechanism.
  • Mamba layer: The Mamba input pipeline expands representations, applies Short Convolution for smoothing, and computes selective gates through low-rank projection and Softplus.The convolution uses kernel size k = 4 for hardware-aware efficiency, while the gate range is initialized to ∆∈[0.001, 0.1].
  • Sliding Window Attention: Sliding Window Attention uses a window size w = 2048 to access recent context while keeping computational complexity linear in sequence length.RoPE is applied within the sliding window, allowing direct access to middle- and short-term context.

3 EXPERIMENTS AND RESULTS

Samba is evaluated across downstream tasks, perplexity, length extrapolation, and long-context retrieval, generally outperforming attention-only, SSM-only, and alternative hybrid architectures while retaining linear-time processing.

  • Downstream evaluation: Samba-3.8B-IT is compared with Phi-3-mini-4K on long-context and short-context tasks using task-specific accuracy, pass@1, and ROUGE-L metrics.The comparison includes MMLU, GSM8K, HumanEval, GovReport, and SQuALITY.
  • Downstream evaluation: Across 15 downstream benchmarks, Samba achieves the best average performance and outperforms pure attention-based and SSM-based models on most tasks.The evaluated tasks include commonsense reasoning, language understanding, TruthfulQA, and code generation.
  • Perplexity and architecture comparisons: Samba consistently outperforms the evaluated architectures in perplexity across context lengths and model sizes, with 4,096-token training and 2,048-token SWA windows.The comparison includes Llama-2-SWA, Sliding RetNet, Sliding GLA, MLP2-SWA-MLP, and Samba-NoPE.
  • Perplexity and architecture comparisons: Removing Mamba layers speeds training but significantly harms perplexity, while removing RoPE causes Samba-NoPE perplexity to explode beyond its training length.RetNet extrapolates well at 438M scale but shows increasing perplexity at 16K length at 1.4B scale.
  • Downstream architecture analysis: Samba has the best average accuracy on five commonsense tasks, although different hybrid architectures excel on different individual tasks.Mamba-SWA-MLP performs best on ARC-Easy, while Samba and Samba-NoPE perform best on LAMBADA.
  • Efficient length extrapolation: After 4K-length training, Samba improves prediction up to 1M tokens and achieves 3.64× faster decoding than Llama-3 for 64K-token generation.The evaluation uses the Proof-Pile test set and includes an SE-Llama-3 1.6B zero-shot extrapolation baseline.
  • Long-context understanding: With 4K-length supervised fine-tuning for 500 steps, Samba 1.7B extrapolates Passkey Retrieval to 256K context and outperforms Mistral 1.6B.The comparison is between Samba 1.7B and a model based solely on Sliding Window Attention.

4 ANALYSIS

The analysis examines why Samba combines Mamba with attention, focusing on extrapolation, parameter allocation, and specialization. Results indicate that Sliding Window Attention and recurrent processing complement each other for efficient long-context modeling.

  • Extrapolation and efficiency: Samba’s full-attention hybrids show exploding extrapolation perplexity at 16K, while intermediate full-attention placements still fail beyond training lengths.Samba also has higher training throughput than Mamba-MLP alternatives at sequence length 4096.
  • Attention allocation: Samba achieves improved validation perplexity with one key-value head, and its optimal number of query heads is 2× smaller than SWA’s.The authors interpret this as evidence that Samba can use fewer attention heads because Mamba provides recurrent compression.
  • Layer specialization: Samba’s attention entropy varies more across layers, with higher entropy in upper and lower layers and lower entropy in middle layers.The authors associate this pattern with specialized retrieval in middle layers and global-information integration in upper and lower layers.
  • Layer specialization: Compared with Mamba-MLP, Samba has higher middle-layer input-selection entropy, allowing Mamba layers to focus more on recurrent structure than precise retrieval.The authors suggest this specialization may benefit downstream performance.

5 CONCLUSION

Samba is presented as a hybrid architecture for efficient language modeling with unlimited context length. It outperforms pure attention and SSM models across benchmarks, improves long-context processing efficiency, and supports memory recall at up to 256K contexts with minimal fine-tuning.

  • Benchmark performance: Samba substantially outperforms state-of-the-art pure attention-based and SSM-based models across common-sense reasoning, language understanding, mathematics, and coding benchmarks.The conclusion attributes the architecture’s scope to combining attention mechanisms with SSMs.
  • Efficiency: Samba achieves substantial prompt-processing and decoding-throughput speedups compared with the state-of-the-art Transformer architecture.The conclusion presents efficiency as a central property of the hybrid design.
  • Long-context recall: Samba extrapolates memory recall to 256K contexts through minimal fine-tuning and shows usefulness in downstream long-context summarization tasks.The conclusion also reports analyses of training configurations and attention-SSM combinations.

A RELATED WORKS

Related work combines SSMs or recurrent models with attention to address retrieval and long-sequence modeling. Existing approaches include sparse attention, hybrid recurrent architectures, and length-extrapolation methods, but the cited literature leaves efficiency and extrapolation trade-offs unresolved.

  • Hybrid recurrent models: Hybrid recurrent models combine SSMs or linear recurrent layers with attention because linear SSMs have limited retrieval ability.The literature includes chunked-attention and quadratic-self-attention hybrids.
  • Efficient sparse attention: Efficient sparse-attention methods use static or dynamic patterns to obtain subquadratic sequence complexity, but their wall-time training efficiency is often worse without hardware-aware implementations.The passage contrasts asymptotic sparsity with practical training efficiency.
  • Length extrapolation: Length-extrapolation methods can stabilize perplexity beyond training lengths with linear complexity, but the cited passage states they do not significantly improve it.The passage frames these methods as approaches for extending pretrained Transformer context lengths.
  • Length extrapolation: Samba’s zero-shot retrieval performance still lags behind the compared approaches, exposing a trade-off between perplexity and retrieval performance in length extrapolation.The authors identify this trade-off as a topic for future work.

B ADDITIONAL EVALUATION RESULTS

Additional evaluations compare Samba with pretrained and post-trained language models across broad benchmarks, while figures examine throughput and Passkey Retrieval training behavior. Samba achieves the highest average score in the reported base-model comparison and favorable performance against controlled Transformer baselines.

  • Base-model evaluation: SAMBA achieves the highest average score across the diverse benchmarks in Table 7 and 18.1% higher GSM8K accuracy than TFM++ trained on the same dataset.The table evaluates language comprehension across multiple domains and specifies different zero-shot and few-shot settings.
  • Post-trained evaluation: Post-trained hybrid models outperform industry-standard Transformer and SSM-based LLMs in the reported representative benchmark comparison.The fair comparison is restricted to models evaluated under the stated comparison conditions.
  • Parameter efficiency: SAMBA delivers performance comparable to Jamba-1.5-Mini with around 3× fewer active parameters and 13× fewer total parameters.The passage attributes this comparison to the reported data-synthesis technique and parameter counts.
  • Efficiency: Figure 6 compares prompt-processing throughput for models with approximately 1.7B parameters.The supplied caption identifies the comparison but does not state the numerical outcome.
  • Passkey Retrieval: Figures 7 and 8 track instruction-tuning loss and 256K-document Passkey Retrieval accuracy for Samba 1.7B and Mistral 1.6B.Both figures cover 500 training steps, with 4K sequence length specified for the loss curves.

C ADDITIONAL EXPERIMENT DETAILS

The Passkey Retrieval experiment instruction-tunes Mistral 1.6B and Samba 1.7B on synthetically generated documents containing randomly placed five-digit passkeys. Both models use the same optimization setup.

  • Passkey Retrieval: Mistral 1.6B and Samba 1.7B are instruction-tuned on Passkey Retrieval with 4096-token documents.Passkeys are randomly sampled five-digit integers inserted at random locations or depths.
  • Data generation: The training data are generated on the fly by varying each five-digit passkey and its insertion location.The model must generate the passkey after reading the full document.
  • Optimization: Both models use batch size 2048, 250 warm-up steps, peak learning rate 1e−4, and 0.1 weight decay with AdamW.

D ADDITIONAL ANALYSES

Additional analyses examine how context length and batch size affect SWA training, how Short Convolution changes recurrent baselines, and how attention and selective-gate entropy are measured.

  • SWA training: Increasing training sequence length raises validation perplexity at all context lengths when smaller batch sizes result.Table 9 fixes the sliding-window size at 2048 and training tokens per step at 2M.
  • Short Convolution: Short Convolution improves RetNet and SWA performance, but has a less prominent effect on GLA.The authors attribute GLA’s smaller gain to its existing fine-grained channel-level decays.
  • Hybrid comparison: Even with Short Convolution, Sliding GLA and Sliding RetNet underperform the original Samba 421M model.This comparison supports using Mamba as the recurrent component in the hybrid architecture.
  • Entropy analysis: Average attention entropy is computed per decoding step from a causal attention probability matrix with h heads and sequence length n.The setup considers generation length l satisfying 0 < l < n.
  • Entropy analysis: The S6 selective gate is normalized to the simplex [0, 1]^n×d_e before selection entropy is calculated across the sequence.

F DETAILS OF DOWNSTREAM LONG-CONTEXT EVALUATION

Long-context summarization is evaluated on GovReport and SQuALITY from ZeroSCROLLS, using documents whose lengths extend well beyond standard short-context settings.

  • Evaluation datasets: GovReport and SQuALITY are used to evaluate real-world long-context summarization.Both datasets are drawn from the ZeroSCROLLS benchmark and tokenized with the Phi3-mini-4k tokenizer.
  • GovReport: GovReport documents average 11,533 tokens, with lengths ranging from 1,493 to 40,592 tokens.The median document length is 10,332 tokens.
  • SQuALITY: SQuALITY documents average 7,974 tokens after tokenization.

G IMPLEMENTATION DETAILS

Implementation details specify baseline configurations, model-scale training settings, recurrent-layer implementations, and decoding strategies for downstream tasks.

  • Baseline configurations: Baseline models are trained on the Phi2 dataset with 230B tokens using the hyperparameters listed in Table 11.
  • Layer implementations: Sliding GLA and RetNet use specified head and expansion ratios, with implementations from the Flash Linear Attention repository.Self-Extend extrapolation uses a FlashAttention-based implementation.
  • SAMBA configurations: The SAMBA models use scale-specific hyperparameters, while Table 12 reports only first-phase optimization settings for the 3.8B model.
  • Decoding: Downstream generation uses greedy decoding for GSM8K and nucleus sampling for HumanEval, MBPP, and SQuAD.HumanEval uses temperature 0.2; MBPP and SQuAD use temperature 0.01, all with top-p 0.95.

H LIMITATIONS & BROADER IMPACT

Samba’s reported benefits are accompanied by limitations in retrieval performance and consistency across task-specific hybridization strategies. The authors identify these boundaries as opportunities for improving dynamic model combinations.

  • Retrieval limitations: The pre-trained Samba base model has retrieval performance similar to the SWA-based model.Instruction tuning produces promising memory retrieval performance, but the base model does not exceed the SWA-based model in Figure 8.
  • Task dependence: Samba’s hybridization strategy is not consistently better than alternative architectures across all tasks.Mamba-SWA-MLP improves performance on WinoGrande, SIQA, and GSM8K, indicating task-dependent differences among hybrid designs.
  • Future direction: The authors suggest input-dependent dynamic combinations of SWA-based and SSM-based models as a future direction.This direction is motivated by the differing task performance of the evaluated hybridization strategies.
Loading 2406.07522v3…