[논문 리뷰] Samba: Simple Hybrid State Space Models for Efficient Unlimited Context Language Modeling

Liliang Ren, Yang Liu외 4명|arXiv (Cornell University)|2024. 06. 11.

Topic Modeling인용 수 7

한 줄 요약

Samba는 Mamba(SSM)와 Sliding Window Attention을 공동으로 통합하여 무제한 컨텍스트 길이를 달성하고 선형 시간 복잡도로 강력한 벤치마크 성능과 높은 처리량 이득을 제공합니다.

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

연구 동기 및 목표

• 훈련 맥락 길이를 넘어 확장하는 효율적인 언어 모델링에 대한 동기를 부여한다.
• SSMs와 어텐션을 층단위 방식으로 결합하는 간단한 하이브리드 아키텍처를 제안한다.
• 롱 컨텍스트 기능을 갖춘 수십억 매개변수 모델로의 확장 가능성을 입증한다.
• 추론, 수학, 코딩 및 긴 컨텍스트 검색을 포함한 다양한 작업에서 성능을 평가한다.
• 하이브리드 설계를 이해하기 위한 훈련 및 추론 효율성 분석과 ablation 연구를 수행한다.

제안 방법

• Samba를 도입한다, Mamba(SSM)와 Sliding Window Attention(SWA)의 층별 하이브리드.
• 입력 의존 선택적 상태 공간을 통해 시간 의존적 의미를 포착하기 위해 Mamba 계층을 사용한다.
• 선형 시간 프레임워크 내에서 정밀한 메모리 검색을 가능하게 하기 위해 2048 창의 SWA를 도입한다.
• 비선형 변환을 위한 SwiGLU MLP를 활용하고 서로 다른 정보 스트림에 대해 분리된 MLP를 사용한다.
• 421M, 1.3B, 1.7B, 3.8B 매개변수 모델을 최대 3.2T 토큰에서 훈련하고 다양한 벤치마크에서 평가한다.
• 효율적인 긴 컨텍스트 외삽 테스트(최대 256K 프롬프트와 1M 컨텍스트) 및 처리량 비교를 수행한다.

(a) Perplexity on the test set of Proof-Pile

실험 결과

연구 질문

• RQ1 입력 의존 SSM과 sliding window attention의 간단한 층별 하이브드가 선형 시간 복잡도에서 무제한 컨텍스트 길이를 달성할 수 있는가?
• RQ2 Samba가 긴 컨텍스트 및 짧은 컨텍스트 벤치마크에서 순수 어텐션 및 순수 SSM 모델과 어떻게 비교되는가?
• RQ3 긴 컨텍스트 언어 모델링에서 가장 높은 성능과 효율을 얻는 아키텍처 선택은 무엇인가(예: 어텐션 대 재귀의 분배)?
• RQ4 맥락 길이를 외삽할 때 기억 회상 및 검색 기능이 어느 정도 보존되거나 향상될 수 있는가?
• RQ5 프롬프트 처리 및 생성에 대한 효율성 및 처리량 영향은 무엇인가?

주요 결과

• Samba는 최대 8B 매개변수까지의 강력한 오픈 소스 LLM과 비교하여 광범위한 벤치마크에서 평균 성능에서 최첨단을 달성한다.
• The 3.8B Samba model attains high scores on MMLU (71.2) and GSM8K (69.6) and strong HumanEval (54.9) on 4K-training-length pretraining.
• Samba는 최소한의 미세조정 후 256K 컨텍스트까지 거의 완벽한 기억 회상과 효과적인 외삽을 시연하고, Proof-Pile의 퍼플렉시티 개선을 위해 최대 1M 컨텍스트까지 가능하다.
• 처리량 이점: Samba는 128K 프롬프트에서 Llama-3 1.6B보다 프롬프트 처리 처리량이 약 3.73배 높고 스트리밍으로 64K 토큰 생성에서 3.64배 가속을 제공한다.
• Samba는 선형 시간 복잡도를 유지하며 긴 컨텍스트 작업 및 지시 학습 후 다운스트림 긴 컨텍스트 요약에서 SWA 전용 또는 순수 어텐션 벤치마크를 능가한다.
• ablation 연구는 Mamba와 SWA의 하이브리드가 모든 작업에서 최고의 성능을 보임을 시사하며, 어텐션 헤드 수를 줄여도 성능이 유지되거나 향상될 수 있다.

(b) Decoding throughput with a batch size of 16