Samba: Einfache hybride Zustandsraummodelle für effizientes unbegrenztes Kontext-Sprachmodellieren

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

June 11, 2024

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

cs.AI

Zusammenfassung

Die effiziente Modellierung von Sequenzen mit unendlicher Kontextlänge war ein langjähriges Problem. Frühere Arbeiten leiden entweder unter der quadratischen Rechenkomplexität oder der begrenzten Extrapolationsfähigkeit bei der Längenverallgemeinerung. In dieser Arbeit präsentieren wir Samba, eine einfache Hybridarchitektur, die Mamba, ein selektives Zustandsraummodell (SSM), mit Sliding-Window-Attention (SWA) schichtweise kombiniert. Samba komprimiert eine gegebene Sequenz selektiv in rekurrente versteckte Zustände, während es gleichzeitig die Fähigkeit behält, Erinnerungen präzise mit dem Aufmerksamkeitsmechanismus abzurufen. Wir skalieren Samba auf 3,8 Mrd. Parameter mit 3,2 Bio. Trainings-Token und zeigen, dass Samba auf einer Vielzahl von Benchmarks wesentlich besser abschneidet als die State-of-the-Art-Modelle, die auf reiner Aufmerksamkeit oder SSMs basieren. Wenn auf 4K-langen Sequenzen trainiert, kann Samba effizient auf eine Kontextlänge von 256K extrapolieren, mit perfektem Erinnerungsrückruf und verbesserten Token-Vorhersagen bis zu einer Kontextlänge von 1M. Als Sequenzmodell mit linearer Laufzeit hat Samba eine 3,73-fach höhere Durchsatzrate im Vergleich zu Transformatoren mit gruppierten Query-Attention, wenn Benutzeranfragen von 128K Länge verarbeitet werden, und eine 3,64-fache Beschleunigung beim Generieren von 64K Tokens mit unbegrenztem Streaming. Eine beispielhafte Implementierung von Samba ist öffentlich verfügbar unter 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.