В мире искусственного интеллекта доминирование архитектуры Transformer кажется неоспоримым, однако её фундаментальные ограничения — квадратичная сложность вычислений и огромные требования к памяти при работе с длинными текстами — заставляют исследователей искать альтернативы. В новом видео технический обозреватель Янник Килчер (Yannic Kilcher) подробно разбирает научную работу Альберта Гу (Albert Gu) и Три Дао (Tri Dao) под названием «Mamba: Linear-Time Sequence Modeling with Selective State Spaces». Автор объясняет, почему Mamba претендует на роль «убийцы Трансформеров», предлагая линейное масштабирование и высокую эффективность при сохранении качества моделирования.
🧠 Проблема выбора: Transformers против RNN и SSM 0:25
Начиная разбор, Янник Килчер противопоставляет три основных подхода к моделированию последовательностей, подчеркивая их сильные и слабые стороны .
- Трансформеры (Transformers): Работают на механизме внимания (attention), где каждый элемент последовательности может напрямую взаимодействовать с любым другим . Это позволяет модели динамически и селективно фокусироваться на важных деталях прошлого. Однако цена такой гибкости — квадратичная сложность ($L^2$, где $L$ — длина последовательности), что делает обработку очень длинных текстов чрезвычайно дорогой .
- Рекуррентные нейронные сети (RNN): Обновляют скрытое состояние (hidden state) шаг за шагом. Они эффективны по памяти ($O(1)$ для одного шага), но страдают от проблем исчезающих или взрывающихся градиентов и не могут параллелиться при обучении . По мнению Килчера, RNN слишком ограничены, так как «видят» только последнее состояние и текущий вход .
- Модели в пространстве состояний (SSM): Такие как архитектура S4, пытаются совместить лучшее из двух миров. Они ведут себя как RNN во время работы (инференса), но могут обучаться параллельно, как сверточные сети (CNN), благодаря линейной структуре переходов . Главный минус классических SSM — их «статичность»: правила перехода между состояниями фиксированы и не зависят от входных данных, что мешает им эффективно работать с естественным языком .
🐍 Mamba и концепция «Селективности» 9:40
Ключевое новшество Mamba — переход от структурированных (Structured) к селективным (Selective) моделям пространства состояний . Янник объясняет, что авторы статьи идентифицировали главную слабость предыдущих SSM: неспособность выполнять контекстно-зависимые рассуждения .
В классических SSM матрицы переходов $A$ и $B$ фиксированы для всей последовательности. Mamba же делает эти параметры зависимыми от текущего входа (input dependent) . Это позволяет модели «решать», какую информацию из прошлого сохранить в скрытом состоянии, а какую — игнорировать.
По словам Килчера, Mamba находится на золотой середине:
- Она сохраняет линейную структуру, что позволяет использовать эффективные алгоритмы вычислений.
- Она внедряет «селективность», делая модель почти такой же мощной, как Трансформеры в задачах понимания контекста .
- Она не требует Bloom-фильтров или кэша ключей/значений (KV-cache), что радикально ускоряет генерацию текста .
⚙️ Аппаратная оптимизация и параллельное сканирование 17:07
Проблема селективности в том, что она нарушает возможность использовать простую свертку (convolution) для быстрого обучения . Если параметры меняются на каждом шаге, вы не можете «промчаться» по всей последовательности одним умножением матриц.
Чтобы решить это, авторы Mamba разработали аппаратно-зависимый алгоритм (Hardware-aware Algorithm). Килчер выделяет два ключевых аспекта:
- Использование SRAM против HBM: В современных GPU есть медленная основная память (HBM) и очень быстрая, но маленькая кэш-память (SRAM) . Авторы Mamba оптимизировали передачу данных так, чтобы основные вычисления (дискретизация и рекурренция) происходили внутри SRAM, минимизируя дорогостоящие обращения к HBM . Это делает Mamba такой же быстрой в обучении, как Трансформеры с оптимизацией Flash Attention .
- Параллельное сканирование (Parallel Scan/Prefix Sum): Вместо свертки используется алгоритм префиксных сумм . Это математический трюк, позволяющий вычислить цепочку зависимых рекуррентных состояний параллельно, если операции в ней ассоциативны .
📊 Результаты и будущее архитектуры 36:10
Оценивая перспективы Mamba, Килчер отмечает, что модель показывает впечатляющие результаты в масштабировании (scaling laws) .
- Языковое моделирование: Модели Mamba размером до 1,3 млрд параметров обходят другие модели без механизма внимания и приближаются к уровню мощных конфигураций Transformer++ .
- Длинные контексты: Архитектура демонстрирует превосходство в задачах с последовательностями длиной до 1 миллиона токенов, таких как анализ ДНК или обработка аудио [11:45, 20:16].
- Производительность: Скорость генерации (throughput) на чипах A100 значительно выше, чем у Трансформеров, и этот разрыв увеличивается с ростом размера пакета (batch size) .
Тем не менее, Килчер призывает к осторожности: на момент публикации статьи эксперименты проводились на моделях до 1–3 млрд параметров, в то время как мейнстрим-модели (например, Llama) начинаются от 7 млрд . «Жюри еще не вынесло вердикт», — шутит ведущий, рассуждая о том, сможет ли Mamba полностью вытеснить Трансформеры на сверхкрупных масштабах .
💻 Разбор кода и архитектурные блоки 37:56
В завершение видео Янник проводит краткий аудит репозитория Mamba на GitHub . Он рекомендует начинать изучение с файла mamba_simple.py, так как там логика представлена в наиболее чистом виде, без сложных оптимизаций для CUDA .
Основные этапы вычисления одного блока Mamba:
- Линейная проекция: Входной токен расширяется в более высокую размерность .
- 1D-свертка: Локальное смешивание информации между соседними токенами .
- Дискретизация: Параметры непрерывной системы переводятся в дискретные веса .
- Рекурренция (SSM): Обновление скрытого состояния и вычисление выхода .
- Gating (умножение на затвор): Объединение выхода SSM с дополнительной веткой вычислений .
По мнению Килчера, эта архитектура является сильнейшим кандидатом на роль универсального «движка» для работы с последовательностями .
