Samba: Modelli Ibridi Semplici a Spazio di Stati per una Modellazione Linguistica Efficiente con Contesto Illimitato

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

June 11, 2024

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

cs.AI

Abstract

Modellare efficientemente sequenze con una lunghezza di contesto infinita è stato un problema di lunga data. I lavori precedenti soffrono o della complessità computazionale quadratica o della limitata capacità di estrapolazione nella generalizzazione della lunghezza. In questo lavoro, presentiamo Samba, una semplice architettura ibrida che combina a livello di strato Mamba, un modello di spazio degli stati (SSM) selettivo, con l'attenzione a finestra scorrevole (SWA). Samba comprime selettivamente una sequenza data in stati nascosti ricorrenti, mantenendo comunque la capacità di richiamare con precisione i ricordi attraverso il meccanismo di attenzione. Scaliamo Samba fino a 3,8 miliardi di parametri con 3,2 trilioni di token di addestramento e dimostriamo che Samba supera sostanzialmente i modelli all'avanguardia basati su attenzione pura o SSM su un'ampia gamma di benchmark. Quando addestrato su sequenze di lunghezza 4K, Samba può essere efficientemente estrapolato a una lunghezza di contesto di 256K con un richiamo perfetto della memoria e mostra miglioramenti nella previsione dei token fino a una lunghezza di contesto di 1M. Come modello di sequenza a tempo lineare, Samba beneficia di un throughput 3,73 volte superiore rispetto ai Transformer con attenzione a query raggruppate quando elabora prompt utente di lunghezza 128K, e di un'accelerazione di 3,64 volte quando genera 64K token con streaming illimitato. Un'implementazione campione di Samba è disponibile pubblicamente su 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.