Samba: Modelos Híbridos Simples de Espacio de Estados para Modelado Eficiente de Lenguaje con Contexto Ilimitado

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

June 11, 2024

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

cs.AI

Resumen

Modelar secuencias de manera eficiente con una longitud de contexto infinita ha sido un problema de larga data. Trabajos anteriores adolecen de una complejidad computacional cuadrática o de una capacidad limitada de extrapolación en la generalización de longitud. En este trabajo, presentamos Samba, una arquitectura híbrida simple que combina de manera estratificada Mamba, un Modelo de Espacio de Estados (SSM) selectivo, con Atención de Ventana Deslizante (SWA). Samba comprime selectivamente una secuencia dada en estados ocultos recurrentes mientras mantiene la capacidad de recordar memorias con precisión mediante el mecanismo de atención. Escalamos Samba hasta 3.8B de parámetros con 3.2T de tokens de entrenamiento y demostramos que Samba supera sustancialmente a los modelos de última generación basados en atención pura o SSMs en una amplia gama de benchmarks. Cuando se entrena en secuencias de longitud 4K, Samba puede extrapolarse eficientemente a una longitud de contexto de 256K con un recuerdo de memoria perfecto y muestra mejoras en la predicción de tokens hasta una longitud de contexto de 1M. Como modelo de secuencia de tiempo lineal, Samba disfruta de un rendimiento 3.73 veces mayor en comparación con Transformers con atención de consulta agrupada al procesar indicaciones de usuario de 128K de longitud, y una aceleración de 3.64 veces al generar 64K tokens con transmisión ilimitada. Una implementación de muestra de Samba está disponible públicamente en 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.