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