Mamba — это работа Albert Gu и Tri Dao о том, как превратить state space model в практичную альтернативу трансформеру для длинных последовательностей. Исторически важно разбирать именно исходную версию: препринт появился на arXiv 1 декабря 2023 года, а версия, опубликованная как conference paper в COLM 2024, зафиксировала тот же основной тезис — проблема прежних субквадратичных моделей была не только в вычислениях, но и в слабой контекстной избирательности. Ниже — разбор именно этой работы, а не более поздней линии Mamba-2 и Structured State Space Duality.
Почему это было важно на момент выхода статьи: к концу 2023 года у индустрии уже был понятный компромисс. Трансформеры давали качество, но дорого масштабировались по длине контекста; линейные и рекуррентные альтернативы были дешевле, но обычно проигрывали на языке. Mamba стала одной из первых работ, где этот разрыв начали закрывать не на уровне лозунга, а на уровне архитектуры, алгоритма и таблиц с результатами.
Коротко
- Главная новизна Mamba не в самом факте «линейного времени», а в selective state spaces: параметры состояния зависят от входного токена и могут по-разному помнить или забывать информацию.
- Эта избирательность ломает старый удобный режим свёртки для SSM, поэтому авторы строят hardware-aware scan, который считает рекуррентную модель параллельно и без разворачивания огромного скрытого состояния в медленную память.
- По таблице 3 самой статьи Mamba-2.8B показывает 6.22 perplexity на валидации The Pile и 63.3 среднего zero-shot score против 6.73 и 59.1 у Pythia-2.8B, а также 7.00 и 59.6 у RWKV-3B.
- По авторскому бенчмарку Figure 8 Mamba даёт примерно 4–5× более высокий throughput при генерации, чем сопоставимый Transformer, а собственная scan-реализация авторами измеряется как до 20–40× быстрее наивной.
- При этом статья не доказывает, что attention больше не нужен: языковые эксперименты здесь ограничены масштабом до 2.8B параметров и контекстом порядка тысяч токенов, тогда как миллионные длины показаны в основном на ДНК и аудио.
Контекст: зачем Mamba вообще понадобилась
До Mamba линия state space models уже существовала. S4 показала, что линейные динамические системы можно сделать полезными для очень длинных последовательностей, H3 подвела этот класс ближе к языковому моделированию, а Hyena предложила сильную субквадратичную свёрточную альтернативу attention. Но на плотных дискретных данных, прежде всего на тексте, трансформер оставался ориентиром по качеству.
Диагноз Gu и Dao был довольно точным. Прежние SSM хорошо работали там, где полезно «сглаживать» или агрегировать длинный сигнал, но хуже справлялись там, где нужно выбирать конкретные элементы контекста по содержанию. Для языка это критично: модель должна не просто помнить длинную историю, а уметь в нужный момент достать релевантный кусок и проигнорировать шум.
Именно поэтому статья строится вокруг идеи selection. Если прежний SSM похож на фиксированный фильтр, который одинаково обрабатывает все токены, то Mamba делает этот фильтр зависимым от входа. На практике это означает: один токен может «пройти» в состояние почти без потерь, другой — быть почти проигнорирован, третий — сбросить часть накопленной истории.
Метод: что именно сделали авторы
От LTI-SSM к selective SSM
Базовый structured state space model в статье описывается как линейная динамика со скрытым состоянием, где параметры перехода и чтения фиксированы во времени. Это удобно вычислительно, потому что такую систему можно считать как рекуррентно, так и через глобальную свёртку. Но именно постоянство параметров и делает модель плохо приспособленной к контентно-зависимому выбору информации.
В Mamba авторы делают зависимыми от входа три части: шаг дискретизации Δ и параметры B и C. Интуиция здесь важнее формализма. Если токен важный, модель может увеличить вклад текущего входа и ослабить влияние старого состояния; если токен несущественный, можно почти не менять накопленное состояние.
h_t = A_t h_(t-1) + B_t x_t
y_t = C_t h_t
A_t, B_t, C_t зависят от x_t через selective-параметрыАвторы отдельно показывают, что в специальном частном случае selective SSM сводится к классической gated-рекуррентности. Это важный мостик для практиков: Mamba не отбрасывает идеи RNN, а переупаковывает их в более общий и более хорошо оптимизируемый класс моделей.
Почему selectivity ломает старый трюк со свёрткой
Как только параметры зависят от входа, модель перестаёт быть линейно-инвариантной по времени. А значит, больше нельзя заранее свернуть всю последовательность одним глобальным ядром. Это центральная цена идеи Mamba: выросла выразительность, но исчез самый удобный вычислительный путь прежних SSM.
Ответ статьи — hardware-aware parallel scan. Вместо того чтобы материализовать большое расширенное состояние в HBM-памяти GPU, авторы держат критические промежуточные величины ближе к быстрым уровням памяти и фьюзят несколько стадий вычисления. Для практической архитектуры это не мелкая инженерная деталь, а обязательная часть дизайна: без неё selective SSM была бы концептуально интересной, но слишком дорогой.
Как устроен блок Mamba
Архитектурно Mamba тоже важна не только внутренним SSM-слоем. Gu и Dao упрощают типичный для языковых моделей стек, где отдельно чередуются attention-подобный блок и MLP. В Mamba эти функции сближаются в одном однородном блоке: есть линейные проекции, лёгкая локальная свёртка, selective SSM-ветка, нелинейность и остаточные связи.
- Входная последовательность сначала проецируется в более богатое представление.
- Из этого представления вычисляются selective-параметры, которые управляют тем, как обновляется состояние.
- Короткая локальная свёртка помогает обработать ближайший контекст до применения основной динамики.
- Затем запускается selective scan — рекуррентный проход, реализованный так, чтобы быть параллельным и дружелюбным к GPU.
- Выход объединяется с остаточной веткой, и блок можно повторять как обычный backbone без отдельного attention-слоя.
Если формулировать совсем кратко, Mamba — это попытка получить контентно-зависимую обработку последовательности без матрицы attention и без растущего KV-cache.
Результаты
Синтетические задачи: сначала проверка механики
Статья разумно начинает не с больших языковых таблиц, а с тестов, где провал старых SSM можно увидеть почти в чистом виде. По таблице 1 этой же работы комбинация Mamba-блока с selective SSM достигает 99.8% accuracy на selective copying; для сравнения, вариант Mamba с неселективным S4 даёт 56.4%, а Hyena внутри той же архитектурной рамки — 28.4%. Это хороший сигнал, что выигрыш идёт не от брендинга блока, а именно от механизма selection.
На задаче induction heads авторы обучают модели на длине 256 и проверяют экстраполяцию вплоть до 1,048,576 токенов. По описанию результатов в статье Mamba обобщает решение на такую длину, то есть примерно в 4000 раз дальше обучающего диапазона, тогда как другие методы не показывают сопоставимого поведения. Это не доказательство «понимания языка», но это сильная демонстрация того, что архитектура умеет управлять памятью не только на словах.
Язык: где Mamba приблизилась к Transformer
Ниже — выборка из Table 3 самой статьи. Важно читать её как авторское сравнение из одного первичного источника: токенизаторы и семейства моделей не везде совпадают, но внутри групп сравнение всё равно показательно.
| Модель | Параметры | Pile val perplexity | LAMBADA accuracy | Среднее zero-shot из Table 3 |
|---|---|---|---|---|
| Pythia-2.8B | 2.8B | 6.73 | 64.7 | 59.1 |
| RWKV-3B | 3.0B | 7.00 | 63.9 | 59.6 |
| Mamba-2.8B | 2.8B | 6.22 | 69.2 | 63.3 |
| GPT-J-6B | 6.0B | — | 68.3 | 63.0 |
Здесь и появляется главный исторический тезис Mamba. По данным одной этой таблицы, Mamba-2.8B обгоняет близкие по размеру Pythia и RWKV и примерно выходит на уровень или чуть выше более крупного GPT-J-6B по среднему zero-shot score. Это и есть содержательное основание фразы из аннотации о том, что Mamba-3B превосходит трансформеры своего размера и приближается к моделям примерно вдвое крупнее.
Ещё один важный кусок результата — Figure 4 со scaling laws на The Pile. Авторы показывают модели примерно от 125M до 1.3B параметров и утверждают, что Mamba лучше других attention-free альтернатив и первой из них приближается к сильному рецепту Transformer++. Для практиков это, возможно, даже важнее отдельных чисел из Table 3: статья утверждает не разовый удачный чекпойнт, а более устойчивую тенденцию по мере роста модели.
Скорость и длинный контекст
В эффективности нужно быть аккуратным. Все ключевые цифры здесь — авторские измерения на A100 80GB PCIe и собственной fused-реализации. Но именно в таком режиме работа и претендует на практическую ценность.
По Figure 8 статьи собственная scan-реализация Mamba быстрее FlashAttention-2 после длины последовательности около 2K и до 20–40× быстрее наивной scan-реализации в PyTorch. В том же разделе авторы сообщают о 4–5× более высоком inference throughput относительно сопоставимых Transformer-моделей, связывая это с отсутствием KV-cache и возможностью поднимать batch size при генерации.
На длинных реальных последовательностях статья тоже не ограничивается языком. В геномике авторы показывают улучшение качества вплоть до длины 2^20 на предобучении и на задаче классификации пяти видов больших человекообразных обезьян. По Table 13 той же статьи Mamba 7M достигает 81.31% accuracy на длине 2^20, тогда как HyenaDNA 1.4M на этой длине имеет 54.87%. Это уже не просто «линейная сложность на бумаге», а пример, где длинный контекст конвертируется в прикладной выигрыш.
Интерпретация
Наш комментарий
На наш взгляд, Mamba была важна по трем причинам. Первая: она сместила разговор с абстрактного «нужна модель линейного времени» к более точному «нужна модель, которая умеет выбирать, что помнить». В этом смысле selective state spaces оказались содержательнее, чем просто очередная попытка удешевить attention.
Вторая: работа изменила инженерный контракт инференса. У трансформера качество тесно связано с KV-cache, который растёт вместе с контекстом. У Mamba вместо этого фиксированное рекуррентное состояние, поэтому при длинной генерации или потоковой обработке профиль затрат выглядит иначе. Для продакшена это иногда важнее асимптотики в абстракции.
Третья: статья показала, что граница между «качественными» и «эффективными» архитектурами не была фундаментальной. Но из этого не следует, что attention устарел. Скорее Mamba открыла новый дизайн-пространство, где можно комбинировать идеи рекуррентности, селекции и аппаратно-дружественной реализации. То, что вскоре после исходной статьи авторы выпустили работу про Structured State Space Duality, лишь подтверждает: семейство было в активной фазе переосмысления, а исходная Mamba — сильный шаг, но не финальная точка.
Ограничения и критика
- Языковой масштаб в статье ограничен. Самая сильная языковая точка в исходной работе — Mamba-2.8B из Table 3. Это серьёзный масштаб для академической статьи, но не уровень поздних frontier-LLM, поэтому перенос выводов на десятки и сотни миллиардов параметров требует осторожности.
- Длинный контекст на языке здесь не доказан так же убедительно, как на ДНК и аудио. Миллионные длины в статье в основном относятся к синтетике, геномике и аудио. Для текстового моделирования центральные сравнения идут на длинах 2048 и 8192.
- Скорость сильно зависит от kernel engineering. Цифры 20–40× и 4–5× привязаны к авторской реализации и конкретному железу. Это не универсальная константа, а результат удачного совмещения архитектуры с низкоуровневой оптимизацией.
- Не все baseline-сравнения на длинном контексте полностью симметричны. Авторы сами пишут, что для RWKV и RetNet часть результатов на 8K отсутствует из-за ограничений реализаций и памяти. Это честная оговорка, но она всё равно ослабляет силу сравнений.
- Mamba решает не все свойства attention. У рекуррентной модели отличный профиль памяти и хорошая компрессия истории, но это не то же самое, что явный произвольный доступ к любому токену через матрицу attention. Где именно эта разница критична для сложных reasoning-задач, исходная статья не закрывает.
- Исторический объём важно фиксировать. Уже в 2024 году линия работ быстро развилась в Mamba-2 и SSD. Поэтому исходную Mamba лучше читать как доказательство принципа selective SSM, а не как окончательный ответ на вопрос о «замене трансформера».
Вывод
Mamba стала важной не потому, что «победила трансформер вообще», а потому что показала рабочий путь к контентно-зависимой рекуррентной модели с хорошим профилем инференса. Главный вклад работы — selective state spaces плюс реализация scan, без которой идея не имела бы практического веса. Если вы проектируете модели для длинной генерации, потоковой обработки или сред с жёсткими ограничениями памяти, исходная Mamba остаётся обязательной точкой отсчёта.
Источники
- Mamba: Linear-Time Sequence Modeling with Selective State Spaces — Albert Gu, Tri Dao, arXiv.
- Mamba: Linear-Time Sequence Modeling with Selective State Spaces (PDF) — основная версия статьи с таблицами и приложением.
- Mamba: Linear-Time Sequence Modeling with Selective State Spaces — COLM 2024, OpenReview.
- state-spaces/mamba — официальный репозиторий с кодом, чекпойнтами и командами для eval.
- Efficiently Modeling Long Sequences with Structured State Spaces — Albert Gu, Karan Goel, Christopher Ré, ICLR 2022.
- Hungry Hungry Hippos: Towards Language Modeling with State Space Models — Tri Dao и др., ICLR 2023.
- Hyena Hierarchy: Towards Larger Convolutional Language Models — Michael Poli и др., ICML 2023.
- FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning — Tri Dao, ICLR 2024.
- Pythia: A Suite for Analyzing Large Language Models Across Training and Scaling — Stella Biderman и др., arXiv.
- The Pile: An 800GB Dataset of Diverse Text for Language Modeling — Leo Gao и др., arXiv.
- Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality — Tri Dao, Albert Gu, ICML 2024.
FAQ
Чем Mamba отличается от обычных линейных attention-моделей?
Mamba не аппроксимирует attention-матрицу. Она строит рекуррентную модель со скрытым состоянием, где динамика зависит от текущего входа через selective-параметры. Это другой путь к длинному контексту и другому профилю инференса.
Заменяет ли Mamba трансформер в LLM?
Исходная статья показывает сильный результат на малом и среднем масштабе и меняет представление о том, насколько хорошими могут быть attention-free модели. Но она не доказывает универсальную замену трансформера во всех режимах и на всех масштабах.
Почему Mamba быстра при генерации?
Потому что она не хранит растущий KV-cache. На каждом шаге обновляется фиксированное скрытое состояние, поэтому стоимость декодирования не растёт вместе с уже прочитанным контекстом так же, как у стандартного attention.
Где статья сильнее всего, а где слабее?
Сильнее всего — в объяснении механизма selection, в синтетических тестах памяти и в демонстрации эффективности инференса. Слабее — в экстраполяции языковых выводов на более крупные модели и на современные frontier-настройки, которых в работе ещё нет.