Увеличение вычислительных ресурсов во время логического вывода стало основным фактором, влияющим на производительность больших языковых моделей (LLM). Это привело к тому, что архитектурный фокус сместился в сторону повышения эффективности логического вывода наряду с качеством модели.
Хотя архитектуры на основе трансформеров остаются стандартом, их квадратичная вычислительная сложность и линейные требования к памяти создают значительные узкие места при развёртывании. Команда исследователей из Университета Карнеги-Меллона (CMU), Принстонского университета, Together AI и Cartesia AI представила Mamba-3 — модель, которая решает эти проблемы за счёт «ориентированного на вывод» дизайна.
Основные методологические обновления Mamba-3
Mamba-3 основана на фреймворке State Space Model (SSM) и включает три основных обновления:
1. Экспоненциально-трапециевидная дискретизация.
Модели пространства состояний — это системы непрерывного времени, которые необходимо дискретизировать для обработки дискретных последовательностей. Предыдущие версии, такие как 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 к состоянию-входу Btxt в рамках основного повторения. В эмпирических испытаниях эта неявная свёртка в сочетании с обучаемыми смещениями B и C позволяет Mamba-3 эффективно функционировать без внешних коротких причинных свёрток, которые обычно требуются рекуррентным моделям.
2. Комплекснозначные модели пространства состояний и «трюк» RoPE.
Ограничением вещественных линейных моделей является их неспособность решать задачи «отслеживания состояния», такие как определение чётности битовых последовательностей. Это связано с ограничением собственных значений матрицы перехода вещественными числами, которые не могут представить «вращательную» динамику, необходимую для таких задач.
Mamba-3 использует комплекснозначные SSM для решения этой проблемы. Исследовательская группа установила теоретическую эквивалентность между дискретизированными комплексными SSM и вещественными SSM, которые используют зависящие от данных ротационные позиционные встраивания (RoPE) для проекций B и C.
Используя «трюк» RoPE, модель применяет агрегированные зависящие от данных вращения по временным шагам. Это позволяет Mamba-3 решать синтетические задачи, такие как чётность и модульная арифметика, где Mamba-2 и вещественные варианты не работают лучше, чем случайное угадывание.
3. Формулировка с несколькими входами и выходами (MIMO).
Чтобы решить проблему аппаратной неэффективности декодирования, связанного с памятью, Mamba-3 переходит от одновходного одновыходного (SISO) повторения к структуре MIMO.
В стандартном декодировании SSM арифметическая интенсивность составляет примерно 2,5 операций на байт, что намного ниже вычислительно-ограниченного режима современных графических процессоров, таких как H100. MIMO увеличивает ранг R входных и выходных проекций (Bt E RNR и xt E RPR), преобразуя обновление состояния из внешнего произведения в матрично-матричное умножение.
Этот сдвиг увеличивает количество операций FLOPs при декодировании до 4 раз по сравнению с Mamba-2 при фиксированном размере состояния. Поскольку дополнительные вычисления накладываются на существующий ввод-вывод памяти, необходимый для обновления состояния, MIMO улучшает качество моделирования и перплексию, сохраняя при этом аналогичную задержку декодирования по часам.
Архитектура и нормализация
Блок Mamba-3 следует макету в стиле Llama, чередуясь с блоками SwiGLU. Ключевые усовершенствования включают:
* BC/QK-нормализация: нормализация RMS применяется к проекциям B и C, что соответствует QKNorm в трансформерах. Это стабилизирует обучение и позволяет удалить пост-гейт RMSNorm, используемый в предыдущих версиях.
* Смещение для конкретных голов: к компонентам B и C после нормализации добавляются обучаемые канальные смещения для индукции поведения, подобного свёртке.
* Гибридная интеграция: при использовании в гибридных архитектурах — чередовании линейных слоёв с самовниманием — добавление предварительного шлюза, сгруппированного RMSNorm, улучшает обобщение длины в задачах поиска.
Результаты и эффективность
Оценки проводились на наборе данных FineWeb-Edu в четырёх масштабах моделей (от 180 млн до 1,5 млрд).
Производительность ниже по течению: в масштабе 1,5 млрд вариант Mamba-3 SISO превосходит Mamba-2 и Gated DeltaNet (GDN). Вариант MIMO (R=4) дополнительно улучшает среднюю точность ниже по течению на 1,2 пункта по сравнению с базовой линией SISO.
Фронтир Парето: Mamba-3 достигает сопоставимой предварительной подготовки перплексии с Mamba-2, используя при этом вдвое меньший объём состояний (например, Mamba-3 с размером состояния 64 соответствует Mamba-2 со 128).
Производительность ядра: оптимизированные ядра Triton (для предварительного заполнения) и CuTe DSL (для декодирования) гарантируют, что дополнительные математические компоненты останутся лёгкими. Ядра SISO Mamba-3 демонстрируют более низкую задержку, чем выпущенные ядра Mamba-2 и GDN при стандартных настройках BF16.
| Модель (1,5 млрд) | Средняя точность ниже по течению (%) ↑ | 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 демонстрирует, что фундаментальные корректировки точки зрения на модель пространства состояний могут преодолеть разрыв между теоретической субквадратичной эффективностью и практическими возможностями моделирования.
1. Какие основные проблемы решает модель Mamba-3 в контексте больших языковых моделей?
В статье указано, что Mamba-3 решает проблемы, связанные с высокой вычислительной сложностью и требованиями к памяти у моделей на основе трансформеров. Это достигается за счёт «ориентированного на вывод» дизайна.
2. Какие ключевые обновления включает в себя модель Mamba-3?
Основные обновления Mamba-3 включают:
* экспоненциально-трапециевидную дискретизацию для более точного приближения интеграла состояния и входа;
* использование комплекснозначных моделей пространства состояний для решения задач «отслеживания состояния»;
* переход от одновходного одновыходного (SISO) повторения к структуре MIMO для улучшения аппаратной эффективности декодирования.
3. Как Mamba-3 улучшает производительность по сравнению с предыдущими версиями?
Mamba-3 использует новые методы дискретизации и комплекснозначные модели, что позволяет ей более эффективно функционировать без внешних коротких причинных свёрток. Это приводит к улучшению качества моделирования и перплексии, сохраняя при этом аналогичную задержку декодирования по часам.
4. Какие результаты были получены при оценке Mamba-3 на наборе данных FineWeb-Edu?
В масштабе 1,5 млрд вариант Mamba-3 SISO превосходит Mamba-2 и Gated DeltaNet (GDN). Вариант MIMO (R=4) дополнительно улучшает среднюю точность ниже по течению на 1,2 пункта по сравнению с базовой линией SISO. Mamba-3 достигает сопоставимой предварительной подготовки перплексии с Mamba-2, используя при этом вдвое меньший объём состояний.
5. Какие выводы можно сделать о потенциале Mamba-3 для будущих исследований в области языковых моделей?
Mamba-3 демонстрирует, что фундаментальные корректировки точки зрения на модель пространства состояний могут преодолеть разрыв между теоретической субквадратичной эффективностью и практическими возможностями моделирования. Это указывает на потенциал Mamba-3 для дальнейших исследований и разработок в области языковых моделей.