Знакомство с Mamba-3: новое поколение State Space Model с вдвое меньшим состоянием и более эффективным MIMO-декодированием
Эффективность вычислений на этапе инференса (вывода) стала главным фактором производительности больших языковых моделей (LLM). Фокус архитектурных исследований сместился в сторону оптимизации вывода. Transformer-архитектуры остаются стандартом. Однако их квадратичная вычислительная сложность и линейный рост потребления памяти создают серьёзные узкие места при развёртывании. Команда из Carnegie Mellon University (CMU), Princeton University, Together AI и Cartesia AI представила Mamba-3. Это модель с принципом «inference-first» (приоритет инференса), призванная решить эти ограничения.
Mamba-3 строится на основе State Space Model (SSM — класс моделей, представляющих последовательность данных как эволюцию скрытого состояния). Модель содержит три ключевых обновления: экспоненциально-трапецеидальную дискретизацию, комплекснозначные обновления состояния и формулировку Multi-Input Multi-Output (MIMO — многоканальный ввод и вывод).
1. Экспоненциально-трапецеидальная дискретизация
SSM работают в непрерывном времени. Чтобы обрабатывать дискретные данные (текст), их нужно преобразовать — дискретизировать. Предыдущие версии (Mamba-1 и Mamba-2) использовали эвристику первого порядка — «экспоненциально-эйлерову» дискретизацию. Mamba-3 заменяет её на экспоненциально-трапецеидальную дискретизацию. Это даёт более точную аппроксимацию второго порядка для интеграла «состояние-вход».
Дискретная рекурренция превращается из двухчленного обновления в трёхчленное:
$$h_{t}=e^{\Delta_{t}A_{t}}h_{t-1}+(1-\lambda_{t})\Delta_{t}e^{\Delta_{t}A_{t}}B_{t-1}x_{t-1}+\lambda_{t}\Delta_{t}B_{t}x_{t}$$
Формула эквивалентна свёртке шириной 2 (зависящей от входных данных) внутри рекурренции. На практике эта неявная свёртка вместе с обучаемыми смещениями B и C позволяет Mamba-3 работать без внешних коротких каузальных свёрток. Обычно такие свёртки обязательны для рекуррентных моделей.
2. Комплекснозначные SSM и «RoPE-трюк»
Вещественнозначные линейные модели не способны решать задачи отслеживания состояния. Например, определение чётности битовой последовательности. Причина: собственные значения матрицы перехода ограничены вещественными числами. Они не могут описать «вращательную» динамику, нужную для таких задач.
Mamba-3 решает это за счёт комплекснозначных SSM. Исследователи доказали теоретическую эквивалентность двух подходов. Первый — дискретизированные комплексные SSM. Второй — вещественнозначные SSM с зависимыми от данных Rotary Positional Embeddings (RoPE — метод кодирования позиции токена через поворот вектора) на проекциях B и C.
Благодаря «RoPE-трюку» модель применяет агрегированные вращения, зависящие от данных, на протяжении всех временных шагов. Это позволяет Mamba-3 решать синтетические задачи вроде Parity и Modular Arithmetic. Для сравнения: Mamba-2 и вещественнозначные варианты на этих задачах не отличаются от случайного угадывания.
3. Формулировка Multi-Input Multi-Output (MIMO)
Для борьбы с аппаратной неэффективностью memory-bound декодирования (ограниченного пропускной способностью памяти) Mamba-3 переходит от Single-Input Single-Output (SISO — один вход, один выход) рекурренции к структуре Multi-Input Multi-Output (MIMO).
При стандартном декодировании SSM арифметическая интенсивность составляет около 2.5 операций на байт. Это далеко от compute-bound режима (ограниченного вычислениями, а не памятью) современных GPU вроде H100. MIMO увеличивает ранг R входных и выходных проекций ($B_t \in \mathbb{R}^{N \times R}$ и $x_t \in \mathbb{R}^{P \times R}$). Обновление состояния превращается из внешнего произведения в матричное умножение.
Это увеличивает количество вычислительных операций (FLOPs) при декодировании в 4 раза относительно Mamba-2. Размер состояния при этом фиксирован. Дополнительные вычисления накладываются на существующий обмен с памятью (memory I/O) для обновления состояния. В итоге MIMO улучшает качество модели и перплексию (метрику предсказания, где меньше — лучше). Задержка декодирования в реальном времени остаётся аналогичной.
Архитектура и нормализация
Блок Mamba-3 следует макету в стиле Llama, чередуясь с блоками SwiGLU (типом активации). Ключевые доработки:
- Нормализация BC/QK: RMS-нормализация применяется к проекциям B и C, по аналогии с QKNorm в Transformer-ах. Это стабилизирует обучение и убирает необходимость в post-gate RMSNorm из предыдущих версий.
- Смещения для голов (head-specific biases): Обучаемые посимвольные смещения добавляются к B и C после нормализации. Это создаёт свёрточное поведение.
- Гибридная интеграция: В гибридных архитектурах (где линейные слои чередуются с self-attention) добавляется pre-gate групповая RMSNorm. Это улучшает генерализацию по длине в задачах поиска.
Результаты и эффективность
Оценка проводилась на датасете FineWeb-Edu. Было проверено четыре масштаба моделей: от 180M до 1.5B параметров.
- Качество на downstream-задачах (проверка на новых данных): На масштабе 1.5B вариант Mamba-3 SISO превосходит Mamba-2 и Gated DeltaNet (GDN). Вариант MIMO (R=4) дополнительно повышает среднюю точность на 1.2 пункта относительно SISO-базлайна.
- Pareto-фронт (оптимальный баланс качества и затрат): Mamba-3 достигает сопоставимой с Mamba-2 перплексии при препрининге (предобучении), используя вдвое меньший размер состояния. Например, Mamba-3 с состоянием 64 соответствует по качеству Mamba-2 с состоянием 128.
- Производительность вычислительных ядер: Оптимизированные ядра на Triton (для этапа prefill — первичного заполнения контекста) и CuTe DSL (для этапа decode — генерации токенов) гарантируют лёгкость дополнительных математических компонентов. SISO-ядра Mamba-3 демонстрируют меньшую задержку, чем опубликованные ядра Mamba-2 и GDN при стандартных настройках BF16.
| Модель (1.5B) | Ср. точность downstream, % ↑ | FW-Edu Ppl ↓ |
|---|---|---|
| Transformer | 55.4 | 10.51 |
| Mamba-2 | 55.7 | 10.47 |
| Mamba-3 SISO | 56.4 | 10.35 |
| Mamba-3 MIMO (R=4) | 57.6 | 10.24 |
Mamba-3 показывает, что доработка математического аппарата SSM способна преодолеть разрыв между теоретической субквадратичной (менее квадратичной) эффективностью и практическим качеством моделирования.