Samba: 효율적인 무제한 컨텍스트 언어 모델링을 위한 단순 하이브리드 상태 공간 모델

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

June 11, 2024

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

cs.AI

초록

무한한 컨텍스트 길이를 가진 시퀀스를 효율적으로 모델링하는 것은 오랫동안 해결되지 않은 문제였습니다. 기존 연구들은 이차 계산 복잡성이나 길이 일반화에서의 제한된 외삽 능력으로 인해 어려움을 겪어 왔습니다. 본 연구에서는 선택적 상태 공간 모델(SSM)인 Mamba와 슬라이딩 윈도우 어텐션(SWA)을 계층적으로 결합한 간단한 하이브리드 아키텍처인 Samba를 제안합니다. Samba는 주어진 시퀀스를 반복적인 은닉 상태로 선택적으로 압축하면서도 어텐션 메커니즘을 통해 정확한 메모리 회수를 유지할 수 있습니다. 우리는 Samba를 3.8B 파라미터와 3.2T 학습 토큰으로 확장하여, 다양한 벤치마크에서 순수 어텐션 또는 SSM 기반의 최신 모델들을 크게 능가하는 성능을 보여줍니다. 4K 길이의 시퀀스로 학습된 Samba는 256K 컨텍스트 길이로 효율적으로 외삽될 수 있으며 완벽한 메모리 회수를 보여주고, 최대 1M 컨텍스트 길이까지 토큰 예측이 개선됩니다. 선형 시간 시퀀스 모델인 Samba는 128K 길이의 사용자 프롬프트를 처리할 때 그룹화된 쿼리 어텐션을 사용하는 트랜스포머에 비해 3.73배 높은 처리량을 가지며, 64K 토큰을 무제한 스트리밍으로 생성할 때 3.64배의 속도 향상을 보입니다. Samba의 샘플 구현은 https://github.com/microsoft/Samba에서 공개되어 있습니다.

English

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.73x higher throughput compared to Transformers with grouped-query attention when processing user prompts of 128K length, and 3.64x speedup when generating 64K tokens with unlimited streaming. A sample implementation of Samba is publicly available in https://github.com/microsoft/Samba.