Samba: Eenvoudige hybride toestandsruimtemodellen voor efficiënte taalmodellering met onbeperkte context

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

June 11, 2024

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

cs.AI

Samenvatting

Het efficiënt modelleren van sequenties met een oneindige contextlengte is een lang bestaand probleem. Eerdere werken lijden onder ofwel de kwadratische rekencomplexiteit ofwel de beperkte extrapolatiecapaciteit bij lengtegeneralisatie. In dit werk presenteren we Samba, een eenvoudige hybride architectuur die laagsgewijs Mamba, een selectief State Space Model (SSM), combineert met Sliding Window Attention (SWA). Samba comprimeert selectief een gegeven sequentie in recurrente verborgen toestanden, terwijl het nog steeds de mogelijkheid behoudt om herinneringen precies op te roepen met het aandachtmechanisme. We schalen Samba op tot 3,8B parameters met 3,2T trainings-tokens en laten zien dat Samba aanzienlijk beter presteert dan de state-of-the-art modellen gebaseerd op pure aandacht of SSM's op een breed scala aan benchmarks. Wanneer getraind op sequenties van 4K lengte, kan Samba efficiënt worden geëxtrapoleerd naar een contextlengte van 256K met perfect geheugenrecall en verbeterde tokenvoorspellingen tot 1M contextlengte. Als een lineair-tijdsequentiemodel geniet Samba van een 3,73x hogere doorvoersnelheid in vergelijking met Transformers met grouped-query attention bij het verwerken van gebruikersprompts van 128K lengte, en een 3,64x versnelling bij het genereren van 64K tokens met onbeperkte streaming. Een voorbeeldimplementatie van Samba is publiekelijk beschikbaar op 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.