Samba : Modèles simples d'espaces d'états hybrides pour une modélisation efficace du langage avec contexte illimité

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

Résumé

La modélisation efficace de séquences avec une longueur de contexte infinie constitue un problème de longue date. Les travaux précédents souffrent soit d'une complexité de calcul quadratique, soit d'une capacité limitée d'extrapolation en termes de généralisation à des longueurs plus importantes. Dans ce travail, nous présentons Samba, une architecture hybride simple qui combine de manière hiérarchique Mamba, un modèle d'espace d'état sélectif (SSM), avec l'attention par fenêtre glissante (SWA). Samba compresse sélectivement une séquence donnée en états cachés récurrents tout en conservant la capacité de rappeler précisément les mémoires grâce au mécanisme d'attention. Nous avons mis à l'échelle Samba jusqu'à 3,8 milliards de paramètres avec 3,2 billions de tokens d'entraînement, et montrons que Samba surpasse largement les modèles de pointe basés sur l'attention pure ou les SSM sur un large éventail de benchmarks. Lorsqu'il est entraîné sur des séquences de longueur 4K, Samba peut être extrapolé efficacement à une longueur de contexte de 256K avec un rappel de mémoire parfait et montre des prédictions de tokens améliorées jusqu'à une longueur de contexte de 1 million. En tant que modèle de séquence à temps linéaire, Samba offre un débit 3,73 fois supérieur par rapport aux Transformers avec une attention par requêtes groupées lors du traitement de prompts utilisateurs de longueur 128K, et une accélération de 3,64 fois lors de la génération de 64K tokens en streaming illimité. Une implémentation de démonstration de Samba est disponible publiquement sur 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.