🔬 Метод

Основной целью модификаций в Mamba-2 было повышение эффективности обучения, и для этого, в частности, упростили матрицу перехода состояний.

Однако, ограничение выразительности может ограничивать и достижимое качество.

С другой стороны, на практике не менее важно иметь и эффективный инференс, утилизирующий вычислительные возможности GPU. Как известно, инференс SSM сильно memory-bound - вычисления занимают гораздо меньше времени, чем трансфер памяти. Отсюда возникает мысль увеличить объем вычислений на один шаг, не меняя при этом количество передаваемой памяти.

В Мамбе-3 3 ключевых нововведения:
1️⃣ Более общий и выразительной вид рекурсии (экспоненциально-трапециоидальная схема)
2️⃣ Переход к комплекснозначной SSM при этом без осуществления вычислений в комплексных числах при помощи RoPE-трюка.
3️⃣ Multi-input, multi-output SSM вместо single-input, single-output. Суть в общих чертах, в том, чтобы обрабатывать состояние в несколько параллельных потоков.

Кроме того, модифицировали и сам слой Мамбы - добавили QK-norm, убрали короткую свертку (которая более не помогает, ввиду того, что схема рекурсии эффективно является сверткой). Также для MIMO-варианта добавляются MIMO-проекции.

Для реализацим используется Trtion, TileLang и CuTe DSL.

🧪 Эксперименты

Архитектуру валидируют, обучая модели размера от 180M до 1.5B токенов на Fineweb-Edu.

Mamba 3 стабильно опережает Mamba-2 и Gated Delta Net того же размера. MIMO вариант еще немного подымает качество.

На retrieval задачах (Needle in a Haystack) Mamba-2 сильно страдала с увеличением длины контекста. Mamba-3 страдает меньше и показывает себя лучше чем GDN в среднем.

Prefill/Decode у SISO варианта Mamba-3 такой же по скорости, как у Mamba-2. У MIMO префилл чуть медленее, но декодирование почти не замедляется.

💡 Выводы

Выглядит как нетривиальное развитие направления с SSM, но полезность можно будет оценить только в более боевых и прикладных сетапах. Почти наверняка 3-ую Мамбу мы увидим в гибридных transformer/ssm моделях.