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

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.

한 줄 요약

Samba는 Mamba와 Sliding Window Attention을 결합한 하이브리드 아키텍처로, 3.8B 파라미터 모델에서 3.73× 높은 처리량과 256K 길이의 완벽한 메모리 회복 성능을 보인다.

핵심 기여도

Samba는 Mamba(선택적 SSM)와 SWA(슬라이딩 윈도우 어텐션)를 레이어별로 결합한 하이브리드 아키텍처로, 무제한 컨텍스트 길이를 효율적으로 모델링한다.
3.8B 파라미터 모델로 3.2T 토큰 학습 후, 128K 길이의 프롬프트에서 3.73× 높은 처리량을 Transformer 대비 달성.
4K 길이로 사전 학습된 Samba는 1M 길이에서 퍼플렉시티 개선, 256K 길이에서 Passkey Retrieval에서 완벽한 메모리 회복.
Phonebook 태스크에서 Full-attention 모델 대비 우수한 추론 성능을 보인다.

핵심 아이디어

Samba는 Mamba와 SWA를 결합한 하이브리드 아키텍처로, Mamba는 시퀀스의 반복 구조를 캡처하고 SWA는 정밀한 메모리 검색을 담당한다. Mamba는 선택적 상태 공간 모델로, 선형 계산 복잡도와 긴 시퀀스에 대한 추론 능력을 제공하며, SWA는 최근 메모리의 정확한 회복을 보장한다. 이 두 모듈은 레이어별로 교차 결합되어, 반복적 구조와 비반복적 의존성을 동시에 처리할 수 있다. Mamba-MLP, Mamba-SWA-MLP와 같은 다른 하이브리드 전략도 실험되었으나, Samba가 가장 효과적인 성능을 보였다.

기술적 접근법

**모델 아키텍처**: Samba는 Mamba, SWA, MLP 레이어를 레이어별로 결합한 하이브리드 구조를 채택.
**모델 규모**: 421M, 1.3B, 1.7B, 3.8B 파라미터 버전으로 확장.
**학습 데이터**: 3.2T 토큰의 오픈소스 데이터로 사전 학습.
**하이퍼파라미터**: 48~54개의 레이어로 구성, 모든 모델은 동일한 구성 요소를 유지하며 레이어 배치만 변경.
**추론 성능**: 128K 길이의 프롬프트에서 3.73× 높은 처리량, 64K 토큰 생성 시 3.64× 빠른 속도.

주요 결과

**Passkey Retrieval**: 256K 길이에서 완벽한 메모리 회복 (100% 정확도).
**Phonebook**: Full-attention 모델 대비 우수한 추론 성능.
**퍼플렉시티**: 4K 길이로 사전 학습된 Samba는 1M 길이에서 퍼플렉시티 개선 (256× 확장 비율).
**MMLU**: 71.9점, HumanEval: 62.8점, GSM8K: 87.6점으로 Phi-3-mini 모델 대비 훨씬 우수한 성능.

의의 및 한계

Samba는 선형 시간 복잡도를 유지하면서 무제한 길이의 시퀀스를 처리할 수 있는 첫 번째 하이브리드 모델로, 기존 Transformer와 SSM 기반 모델의 한계를 극복한다. 특히, 짧은 컨텍스트 태스크에서 SOTA 성능을 유지하면서 긴 컨텍스트에서도 우수한 추론 능력을 보인다. 그러나, Samba는 SWA 기반 모델 대비 어텐션 메커니즘의 비중을 줄이는 전략을 채택하여, 일부 복잡한 비반복적 의존성을 포착하는 데 한계가 있을 수 있다. 또한, 모델 확장 시 하이퍼파라미터 최적화가 필요하다는 점도 한계로 지적된다.

실용적 활용

Samba는 대규모 컨텍스트를 요구하는 실제 응용 분야, 예를 들어 긴 문서 요약, 대화 기록 분석, 법적 문서 처리 등에 유용하게 활용될 수 있다. 특히, 실시간 스트리밍 처리가 필요한 환경에서 높은 처리량과 메모리 효율성으로 인해 실용적 가치가 높다.