Mamba: Моделирование последовательностей за линейное время с использованием избирательных пространств состояний

Mamba: Linear-Time Sequence Modeling with Selective State Spaces

December 1, 2023

Авторы: Albert Gu, Tri Dao

cs.AI

Аннотация

Фундаментальные модели, которые сейчас лежат в основе большинства впечатляющих приложений в глубоком обучении, практически повсеместно основаны на архитектуре Transformer и её ключевом модуле внимания. Многие архитектуры с субквадратичной временной сложностью, такие как линейное внимание, сверточные модели с вентиляцией, рекуррентные модели и структурированные модели пространства состояний (SSM), были разработаны для решения проблемы вычислительной неэффективности Transformers на длинных последовательностях, но они не показали таких же результатов, как механизм внимания, на важных модальностях, таких как язык. Мы выявили, что ключевая слабость таких моделей заключается в их неспособности выполнять рассуждения на основе контента, и предложили несколько улучшений. Во-первых, простое преобразование параметров SSM в функции входных данных устраняет их слабость в работе с дискретными модальностями, позволяя модели избирательно распространять или забывать информацию вдоль измерения длины последовательности в зависимости от текущего токена. Во-вторых, хотя это изменение исключает использование эффективных сверток, мы разработали аппаратно-ориентированный параллельный алгоритм в рекуррентном режиме. Мы интегрировали эти селективные SSM в упрощённую сквозную архитектуру нейронной сети без механизма внимания и даже без блоков MLP (Mamba). Mamba обеспечивает быстрое выполнение (в 5 раз выше пропускная способность по сравнению с Transformers) и линейное масштабирование по длине последовательности, а её производительность улучшается на реальных данных вплоть до последовательностей длиной в миллион элементов. В качестве общей основы для моделирования последовательностей Mamba достигает наилучших результатов на нескольких модальностях, таких как язык, аудио и геномика. В задаче языкового моделирования наша модель Mamba-3B превосходит Transformers того же размера и соответствует Transformers вдвое большего размера как на этапе предварительного обучения, так и на этапе последующей оценки.

English

Foundation models, now powering most of the exciting applications in deep learning, are almost universally based on the Transformer architecture and its core attention module. Many subquadratic-time architectures such as linear attention, gated convolution and recurrent models, and structured state space models (SSMs) have been developed to address Transformers' computational inefficiency on long sequences, but they have not performed as well as attention on important modalities such as language. We identify that a key weakness of such models is their inability to perform content-based reasoning, and make several improvements. First, simply letting the SSM parameters be functions of the input addresses their weakness with discrete modalities, allowing the model to selectively propagate or forget information along the sequence length dimension depending on the current token. Second, even though this change prevents the use of efficient convolutions, we design a hardware-aware parallel algorithm in recurrent mode. We integrate these selective SSMs into a simplified end-to-end neural network architecture without attention or even MLP blocks (Mamba). Mamba enjoys fast inference (5times higher throughput than Transformers) and linear scaling in sequence length, and its performance improves on real data up to million-length sequences. As a general sequence model backbone, Mamba achieves state-of-the-art performance across several modalities such as language, audio, and genomics. On language modeling, our Mamba-3B model outperforms Transformers of the same size and matches Transformers twice its size, both in pretraining and downstream evaluation.