Samba: Простые гибридные модели пространства состояний для эффективного моделирования языка с неограниченным контекстом

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

June 11, 2024

Авторы: Liliang Ren, Yang Liu, Yadong Lu, Yelong Shen, Chen Liang, Weizhu Chen

cs.AI

Аннотация

Моделирование последовательностей с бесконечной длиной контекста эффективно является давней проблемой. Прошлые работы страдали либо от квадратичной вычислительной сложности, либо от ограниченной способности экстраполяции по длине. В данной работе мы представляем Samba, простую гибридную архитектуру, которая слоевым образом объединяет Mamba, селективную модель пространства состояний (SSM), с механизмом внимания со скользящим окном (SWA). Samba сжимает данную последовательность в рекуррентные скрытые состояния, сохраняя при этом способность точно воспроизводить воспоминания с помощью механизма внимания. Мы масштабируем Samba до 3.8 миллиарда параметров с 3.2 триллионами обучающих токенов и показываем, что Samba значительно превосходит современные модели на основе чистого внимания или SSM на широком спектре тестов. Обученный на последовательностях длиной 4K, Samba может быть эффективно экстраполирован до длины контекста 256K с идеальным воспроизведением памяти и улучшенными предсказаниями токенов до длины контекста 1M. Как модель последовательности линейного времени, Samba обладает пропускной способностью в 3.73 раза выше по сравнению с Трансформерами с групповым вниманием по запросам при обработке запросов пользователей длиной 128K и ускорением в 3.64 раза при генерации 64K токенов с бесконечным потоком. Образец реализации Samba доступен публично по ссылке 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.