Результаты

Авторы валидируют новую архитектуру на следующих задачах:
1⃣️️️️️️ Selective Copy и Induction Heads
2⃣️️️️️️ Language model pretraining
3⃣️️️️️️ Моделирование ДНК
4⃣️️️️️️ Моделирование и генерация аудио

На Selective Copying S4, H3 и Hyena выдают качество порядка 20-60% в то время как S6 достигает почти 100% качества. На Induction Heads как трансформеры с разными позиционными энкодингами, так и прошлые state-space модели не обобщаются далеко за пределы контекста увиденного на обучении. А S6 + Mamba выдает стабильно хорошее качество на последовательностях любой длины.

Mamba + S6 достойно себя показывает на языковом моделировании при обучении на Pile, демонстрируя масштабируемость на уровне и даже чуть лучше, чем трансформер архитектуры LLama-2, со всеми последними достижениями и наработками в развитии архитектур данного семейства. H3, RetNet, Hyena, RWKV заметно отстают. Обучают модели следуя предписаниям Chincilla размером от 130M до 2.8B. Mamba заметно опережает по качеству открытые модели сопоставимого размера, а иногда и вдвое большие.

На моделировании ДНК Mamba превосходит уже заметно Трансформер и state-space модели прошлых поколений. Модели, кстати, совсем небольшие, 1.4M и 7M параметров.

На генерации аудио Mamba сильно опережает по метрикам прошлую SOTA DiffWave+ SaShiMi и иные подходы. S4+ MLP, стоит отметить, работает тоже весьма достойно.

На больших последовательностях Scan из Mamba работает куда быстрее, чем даже оптимизированный FlashAttention-2. Авторская реализация Scan в десятки раз быстрее той что в торче. Благодаря эффективности по памяти в одну GPU можно пихать гораздо больший батч, чем в трансформер. Например, 6.9B Мамбе можно скормить 128 последовательностей длины 2048 и не поперхнуться.

В ablation показывают, что полезно иметь все параметры state-space модели A, B, C, \Delta обучаемыми. Но больше всего накидывает \Delta. Комплексная параметризация A не накидывает по сравнению с нормальным распределением. Expansion внутреннего состояния улучшает качество - чем больше фактор, тем больше информации можно запихнуть в скрытое состояние. По ходу работы пользуются N=64.

Вывод

Выглядит весьма впечатляюще, крутая работа с хорошим анализом и мотивацией, серьезными инженерными достижениями и убедительными результатами. Однако, все же интересно, сможет ли Mamba родить “убийцу трансформеров”, или останется одной из многих, пусть и сильных попыток ограниченной успешности. Самая большая модель, которую использовали в бенчмарках по качеству имеет размер менее 3B параметров, что немного по современным меркам. Как признаются сами авторы, масштабирование архитектуры дальше может породить новые технические сложности и вызовы. Поживем-увидим, как говорится.