Mamba: Linear-Time Sequence Modeling with Selective State Spaces

Albert Gu; Tri Dao

Статьи · arXiv 2312.00752

Mamba: Linear-Time Sequence Modeling with Selective State Spaces

Albert Gu, Tri Dao

Опубликовано 2023-12-01cs.LGcs.AI arXiv PDF

Читать с переводом и ИИ

Кратко в трёх строках

🎯 Трансформеры масштабируются квадратично по длине последовательности и требуют хранения объемного KV-кэша при авторегрессионном инференсе, в то время как прежние субквадратичные модели (например, LTI-модели пространства состояний) уступают в дискретных задачах, требующих анализа контекста, таких как языковое моделирование.
🔧 Параметры SSM сделаны зависимыми от входа (селективными) для динамической фильтрации или распространения информации, что объединено с аппаратно-оптимизированным параллельным ассоциативным сканированием с перевычислением промежуточных состояний во избежание материализации расширенного скрытого состояния в HBM памяти GPU.
📈 Mamba-3B сравнивается или превосходит вдвое большие трансформеры в задачах downstream zero-shot (например, средняя точность 63.3% против 59.1% у Pythia-2.8B и 61.7% у Pythia-6.9B), обеспечивая при этом до 5 раз более высокую пропускную способность при инференсе.

Аннотация

Фундаментальные модели, лежащие в основе большинства передовых приложений глубокого обучения, почти повсеместно опираются на архитектуру Transformer и ее ключевой модуль внимания. Для решения проблемы вычислительной неэффективности трансформеров на длинных последовательностях было разработано множество архитектур с субквадратичной сложностью — таких как линейное внимание, стробированные свертки, рекуррентные модели и структурированные модели пространства состояний (SSM), — однако они уступали вниманию в таких важных модальностях, как естественный язык. Мы определили, что их ключевая слабость заключается в неспособности выполнять рассуждения на основе контента, и предложили несколько улучшений. Во-первых, вычисление параметров SSM как функций от входа устраняет их ограничения при работе с дискретными модальностями, позволяя модели выборочно сохранять или забывать информацию вдоль последовательности в зависимости от текущего токена. Во-вторых, хотя данное изменение исключает использование эффективных сверток, мы разработали аппаратно-оптимизированный параллельный алгоритм в рекуррентном режиме. Мы объединили эти селективные SSM в упрощенную сквозную нейросетевую архитектуру без блоков внимания и MLP (Mamba). Mamba демонстрирует быстрый инференс (в 5 $\times$ более высокая пропускная способность по сравнению с трансформерами) и линейную масштабируемость по длине последовательности, а ее качество на реальных данных растет вплоть до длины контекста в миллион токенов. В качестве универсальной базовой модели последовательностей Mamba достигает передовых результатов (state-of-the-art) в различных модальностях, включая текст, аудио и геномику. В задачах языкового моделирования наша модель Mamba-3B превосходит трансформеры аналогичного размера и не уступает моделям вдвое большего размера как при предобучении, так и в downstream-оценке.

Оригинальная аннотация (англ.)

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 (5 $\times$ 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.

Разбор статьи

Проблема и мотивация

Фундаментальные модели преимущественно опираются на архитектуру Transformer и механизм self-attention. Хотя плотное внимание эффективно распределяет информацию по контекстному окну, оно имеет два фундаментальных ограничения: квадратичную вычислительную сложность $O (L^{2})$ относительно длины последовательности $L$ и необходимость сохранять прошлые ключи и значения (KV-кэш), что приводит к линейному времени инференса на каждом шаге и узким местам по памяти.

Существовавшие ранее субквадратичные подходы — такие как линейное внимание, рекуррентные архитектуры и структурированные модели пространства состояний (SSM), например S4, — масштабируются по времени линейно или псевдолинейно. Тем не менее, структурированные SSM исторически уступали по качеству трансформерам на дискретных, информационно плотных модальностях, таких как текст. Авторы выявили первопричину: предыдущие SSM являются линейными стационарными (LTI). Поскольку динамика их переходов не зависит от входного содержимого, они не способны выборочно сжимать контекст, отфильтровывать нерелевантный шум или регулировать сохранение информации при различном расстоянии между токенами.

Ключевые идеи

Механизм селекции как сжатие контекста: Моделирование последовательностей требует сжатия истории в состояние. Прежние LTI SSM использовали фиксированные параметры переходов, вынуждая одинаково обрабатывать все шаги по времени. Параметризация матриц SSM на основе текущего входного токена дает модели контекстную осведомленность, необходимую для избирательной фильтрации или запоминания информации.
Преодоление ограничений несверточного режима: Зависимость параметров от входа нарушает математическую эквивалентность между линейными рекуррентными уравнениями и глобальными свертками. Без сверток наивная рекуррентность потребовала бы материализации расширенного скрытого состояния размера $(B, L, D, N)$ в высокоскоростной памяти GPU (HBM). Mamba решает эту проблему путем слияния дискретизации, параллельного ассоциативного сканирования и выходной проекции внутри быстрой встроенной памяти GPU (SRAM), минимизируя обращения к HBM.
Унифицированная архитектура: Mamba заменяет чередующиеся блоки многоголового внимания (MHA) и MLP единым однородным блоком, объединяющим стробированную структуру H3 со стробированным MLP в стиле SwiGLU.

Метод

Непрерывные модели пространства состояний отображают вход $x (t) \in R$ в $y (t) \in R$ через неявное скрытое состояние $h (t) \in R^{N}$ : $h^{'} (t) = A h (t) + B x (t), y (t) = C h (t)$

С использованием дискретизации методом фиксатора нулевого порядка (ZOH) с шагом $Δ$ непрерывная система переходит в дискретные параметры рекуррентности: $\overline{A} = exp (Δ A), \overline{B} = (Δ A)^{- 1} (exp (Δ A) - I) \cdot Δ B$ $h_{t} = \overline{A} h_{t - 1} + \overline{B} x_{t}, y_{t} = C h_{t}$

Механизм селекции (S6)

Вместо того чтобы сохранять $(Δ, B, C)$ постоянными во времени, Mamba делает их зависимыми от входа:

$B_{t} = s_{B} (x_{t}) = Linear_{N} (x_{t})$
$C_{t} = s_{C} (x_{t}) = Linear_{N} (x_{t})$
$Δ_{t} = τ_{Δ} (Parameter + s_{Δ} (x_{t}))$ , где $τ_{Δ} = softplus$ и $s_{Δ} (x) = Broadcast_{D} (Linear_{1} (x))$

Это обобщает классический механизм стробирования RNN: при $N = 1, A = - 1, B = 1$ рекуррентность упрощается непосредственно до обновления со стробированием: $h_{t} = (1 - g_{t}) h_{t - 1} + g_{t} x_{t}$ , где $g_{t} = σ (Linear (x_{t}))$ .

Аппаратно-оптимизированное расширение состояния

Для поддержания высокой пропускной способности без материализации состояния $(B, L, D, N)$ в памяти HBM:

Параметры $(Δ, A, B, C)$ загружаются из HBM напрямую во встроенную память SRAM.
Дискретизация в $(\overline{A}, \overline{B})$ и параллельное ассоциативное сканирование выполняются полностью внутри SRAM.
Скрытое состояние сворачивается с помощью $C$ до размера $(B, L, D)$ перед обратной записью в HBM.
Промежуточные состояния перевычисляются в SRAM во время обратного прохода (backward pass) вместо сохранения в HBM.

Результаты

Синтетическая экстраполяция: В задаче Induction Heads модель Mamba, обученная на последовательностях длиной 256, идеально экстраполирует результат (точность 100%) вплоть до длины $2^{20} = 1, 048, 576$ (в $4000 \times$ длиннее исходной последовательности), тогда как базовые трансформеры терпят неудачу или сталкиваются с OOM после $2^{14}$ . В задаче Selective Copying Mamba достигает точности 99.8% по сравнению с 57.0% у H3 и 18.3% у S4.
Языковое моделирование: При оценке в рамках скейлинга Chinchilla на наборе данных the Pile с числом параметров до 1.3B, Mamba не уступает или превосходит Transformer++ (с RoPE, SwiGLU, RMSNorm). В downstream-задачах zero-shot (например, LAMBADA, PIQA, HellaSwag, WinoGrande, ARC) Mamba-2.8B достигает средней точности 63.3%, опережая Pythia-2.8B (59.1%), Pythia-6.9B (61.7%) и конкурируя с моделями на 6B–7B параметров.
Геномика и аудио: В задаче классификации ДНК Great Apes на длине последовательности 1M ( $2^{20}$ ) Mamba достигает точности 71.67% (1.4M параметров) и 81.31% (7M параметров) по сравнению с 54.87% у HyenaDNA. В генерации речи на датасете SC09 Mamba снижает FID с 1.99 (SaShiMi) до 0.94 (6.1M параметров) и 0.67 (24.3M параметров).
Эффективность: Оптимизированное ядро сканирования выполняется в 20–40 раз быстрее стандартных реализаций сканирования в PyTorch и работает до 7 раз быстрее FlashAttention-2 на контексте длиной 32K. Пропускная способность инференса до 5 раз выше по сравнению с аналогичными трансформерами.

Ограничения и открытые вопросы

Качество на непрерывных сигналах: Хотя селекция на основе входа улучшает дискретные рассуждения (текст, ДНК), исследования на предобучении сырых аудиосигналов показывают, что полностью стационарные (LTI) модели (такие как стандартный S4) показывают лучшие результаты на равномерно дискретизированных непрерывных данных.
Масштаб моделей: Оцененные модели ограничены размером около 3B параметров и 300B токенов предобучения; эмпирическое поведение на масштабах 7B–70B+ требует дальнейших исследований.
Инфраструктура downstream-применения: Для трансформеров создана обширная экосистема промпт-инжиниринга, контекстных манипуляций, квантования KV-кэша и адаптивного дообучения. Адаптация моделей SSM под эти сценарии взаимодействия остается предметом активных исследований.

Ключевые термины

Structured State Space Models (SSMs): Класс глубоких моделей последовательностей, основанных на непрерывных линейных динамических системах, отображающих одномерные входы через неявное скрытое состояние.
Linear Time Invariance (LTI): Математическое свойство, при котором параметры перехода системы остаются постоянными на всех временных шагах, что позволяет производить быстрое глобальное сверточное вычисление, но ограничивает способность анализа в зависимости от контента.
Selection Mechanism: Параметризация динамики переходов SSM (Δ, B, C) как функций от входного токена, что позволяет модели динамически передавать, сбрасывать или игнорировать контекст вдоль последовательности.
S6: Обозначение для селективных структурированных моделей пространства состояний, вычисляемых с помощью ассоциативного параллельного сканирования вместо свертки.
Parallel Associative Scan: Вычислительно эффективный алгоритм префиксного вычисления, параллельно рассчитывающий линейную рекуррентность по всей длине последовательности на аппаратных ускорителях.
Hardware-Aware State Expansion: Стратегия объединения GPU-ядер (kernel fusion) и перевычисления промежуточных состояний, которая выполняет дискретизацию и сканирование в быстрой памяти SRAM, избегая передачи промежуточных состояний большого размера в HBM.
Induction Heads: Синтетический механизм ассоциативного вспоминания, проверяющий способность модели воспроизводить ранее встреченный токен вслед за предшествующим ему контекстным токеном.
Mamba Block: Упрощенный однородный блок моделирования последовательностей, объединяющий входные проекции, 1D-свертку, селективные SSM и мультипликативное стробирование без использования отдельных слоев внимания или MLP.

Вопросы

Почему селективные SSM нельзя обучать с помощью быстрых БПФ-сверток?

Быстрые глобальные свертки через БПФ (FFT) требуют, чтобы динамика системы (матрицы A, B и C) была постоянной во времени (линейная стационарность). Поскольку в Mamba параметры B, C и Δ вычисляются на основе входного токена на каждом шаге, ядро свертки меняется в каждой позиции, что нарушает сверточную эквивалентность и требует применения параллельного рекуррентного сканирования.

Как Mamba избегает чрезмерного расхода памяти при расширении рекуррентного состояния во время обучения?

Наивная рекуррентность с размерностью расширения состояния N требует материализации промежуточного состояния формы (B, L, D, N) в памяти GPU HBM. Mamba использует слияние ядер (kernel fusion), загружая в SRAM только низкоразмерные входы, производит дискретизацию и параллельное сканирование внутри быстрой памяти SRAM, возвращает в HBM только свернутый выход размера (B, L, D), а промежуточные активации перевычисляет во время обратного распространения ошибки.

Какова связь между параметром Δ в Mamba и стробированием в рекуррентных сетях?

Параметр шага Δ выступает в роли обобщенного механизма стробирования. При дискретизации методом ZOH с N=1 и A=-1 селективная рекуррентность математически сводится к стандартному вентилю RNN: большое значение Δ заставляет модель сосредоточиться на текущем входе и сбросить состояние, тогда как малое Δ сохраняет накопленное состояние и игнорирует вход.

Использует ли Mamba KV-кэш во время генерации?

Нет. Поскольку Mamba функционирует как авторегрессионная рекуррентная модель пространства состояний, при генерации она сохраняет только скрытое состояние фиксированного размера (размерностью D × N на слой), а не кэш ключей и значений для всех прошлых токенов, что обеспечивает инференс за константное время O(1) на каждом шаге.

Обзор на этой странице написан ИИ по тексту статьи и может содержать ошибки. Подробности смотрите в самой статье.

Читать с переводом и ИИ