🌟 Mamba3

Mamba2 делала ставку на быстрое обучение. Ради этого механизм рекуррентных обновлений упрощали: матрицу переходов состояний свели к скаляру, умноженному на единичную матрицу.

Обучение ускорилось, но при декодировании GPU большую часть времени не считает, а гоняет данные между уровнями памяти. Архитектура оказалась с нюансом - вычислительные ядра простаивают.

С тех пор ландшафт изменился. RL с верифицируемыми наградами для кода и математики, агентные пайплайны - все это генерирует прорву токенов на инференсе. Команда Mamba3 задалась вопросом: как выглядела бы SSM-архитектура, если сделать ее с приоритетом на инференс, а не на обучение?

Так родились 3 главных изменения в ядре Mamba.

🟡Новая схема дискретизации.

SSM в базе - это обыкновенное дифференциальное уравнение, которое нужно перевести в дискретную рекуррентную формулу. Mamba1 и Mamba2 использовали комбинацию двух методов (ZOH и Эйлера), подобранную эмпирически.

В Mamba3 реализовали экспоненциально-трапецеидальный метод: вместо одной точки для аппроксимации интеграла берутся обе границы интервала с обучаемым коэффициентом смешивания. В результате рекуррентная формула неявно применяет свёртку к входу скрытого состояния, что делает динамику выразительнее без дополнительных компонентов.

🟡 Комплекснозначная SSM.

Ранние модели семейства S4 работали с комплексными числами, но Mamba1 от них отказалась. Из-за этого модель не справляется даже с простейшими задачами отслеживания состояний (например, определением четности последовательности).

Решение нашли во вращении в двумерном пространстве: вместо комплексных вычислений авторы разложили переход на масштабирование и поворот, а затем применили фишку из RoPE - встроили вращения в матрицы через кумулятивную сумму углов.

Переписывать ядра для поддержки комплексной арифметики не пришлось. Модель решает задачи на чётность и другие бенчмарки, недоступные предыдущим версиям.

🟡 Переход от SISO к MIMO.

В стандартной SSM каждый хэд содержит набор независимых систем (один вход - один выход). При декодировании арифметическая интенсивность составляет около 2,5 операций на байт при пороге вычислительной загруженности на H100 в районе 300.

Mamba3 расширяет матрицы, превращая внешние произведения в матричные умножения. Арифметическая интенсивность растет пропорционально. При этом размер скрытого состояния не увеличивается, а значит, латентность декодирования почти не меняется. Обучение, конечно, дорожает, но это сознательный компромисс.

Еще из архитектуры убрали короткую каузальную свёртку, присутствовавшую с первой Mamba - новая рекуррентная формула и смещения выполняют ту же функцию.

Добавили нормализацию BCNorm по аналогии с QKNorm в трансформерах, перешли на чередование SSM- и MLP-слоев.

Ядра написаны на Triton (prefill SISO), TileLang (prefill MIMO) и CuTe DSL (decode).

🟡Тесты

Mamba-3 SISO при 1,5B параметров показывает лучшую суммарную латентность prefill + decode на всех длинах последовательностей по сравнению с Mamba2, Gated DeltaNet и Llama-3.2-1B под vLLM на одном H100.

MIMO-вариант сопоставим по скорости с Mamba2, но заметно точнее. При анализе Парето-фронта Mamba-3 показывает тот же уровень качества при вдвое меньшем состоянии.


@ai_machinelearning_big_data

#AI #ML #LLM #Mamba3 #TogetherAI