Machinelearning

🌟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