Samba: Modelos Simples de Espaço de Estados Híbridos para Modelagem de Linguagem Eficiente com 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

Resumo

Modelar sequências com comprimento de contexto infinito de forma eficiente tem sido um problema de longa data. Trabalhos anteriores sofrem com a complexidade computacional quadrática ou com a capacidade limitada de extrapolação na generalização de comprimento. Neste trabalho, apresentamos o Samba, uma arquitetura híbrida simples que combina camada por camada o Mamba, um Modelo de Espaço de Estados (SSM) seletivo, com a Atenção de Janela Deslizante (SWA). O Samba comprime seletivamente uma sequência dada em estados ocultos recorrentes, mantendo ainda a capacidade de recordar memórias com precisão por meio do mecanismo de atenção. Escalonamos o Samba até 3,8 bilhões de parâmetros com 3,2 trilhões de tokens de treinamento e mostramos que o Samba supera substancialmente os modelos state-of-the-art baseados em atenção pura ou SSMs em uma ampla gama de benchmarks. Quando treinado em sequências de comprimento 4K, o Samba pode ser extrapolado eficientemente para um comprimento de contexto de 256K com recall de memória perfeito e mostra melhorias na previsão de tokens até um comprimento de contexto de 1M. Como um modelo de sequência de tempo linear, o Samba desfruta de um throughput 3,73x maior em comparação com Transformers com atenção de consulta agrupada ao processar prompts de usuário com comprimento de 128K, e uma aceleração de 3,64x ao gerar 64K tokens com streaming ilimitado. Uma implementação de amostra do Samba está disponível publicamente em 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.