Смесь Мамба: Улучшение многомодальных моделей пространства состояний с осведомленной о модальности разреженностью

Mixture-of-Mamba: Enhancing Multi-Modal State-Space Models with Modality-Aware Sparsity

January 27, 2025

Авторы: Weixin Liang, Junhong Shen, Genghan Zhang, Ning Dong, Luke Zettlemoyer, Lili Yu

cs.AI

Аннотация

Модели пространства состояний (State Space Models, SSM) стали эффективной альтернативой трансформерам для последовательного моделирования, однако их неспособность использовать модально-специфические особенности ограничивает их производительность в мульти-модальном предварительном обучении. Здесь мы предлагаем Mixture-of-Mamba, новую архитектуру SSM, которая вводит модально-осознанную разреженность через модально-специфическую параметризацию блока Mamba. На основе Mixture-of-Transformers (W. Liang и др., arXiv:2411.04996; 2024) мы расширяем преимущества модально-осознанной разреженности на SSM, сохраняя их вычислительную эффективность. Мы оцениваем Mixture-of-Mamba в трех настройках мульти-модального предварительного обучения: Transfusion (чередующиеся текстовые и непрерывные изображения с потерей диффузии), Chameleon (чередующиеся текстовые и дискретные изображения) и расширенная трехмодальная структура, включающая речь. Mixture-of-Mamba последовательно достигает тех же значений потерь на более ранних этапах обучения с значительным снижением вычислительных затрат. В настройке Transfusion Mixture-of-Mamba достигает эквивалентной потери изображения, используя лишь 34.76% от вычислительных операций обучения на масштабе 1.4 млрд. В настройке Chameleon Mixture-of-Mamba достигает схожей потери изображения всего с 42.50% от вычислительных операций на масштабе 1.4 млрд и схожей потери текста всего с 65.40% от вычислительных операций. В трехмодальной настройке MoM соответствует потере речи на 24.80% от вычислительных операций на масштабе 1.4 млрд. Наше исследование абляции подчеркивает синергетические эффекты разделения компонентов проекции, где совместное разделение приносит большую пользу, чем индивидуальные модификации. Эти результаты устанавливают модально-осознанную разреженность как универсальный и эффективный принцип проектирования, расширяя ее влияние от трансформеров к SSM и устанавливая новые стандарты в мульти-модальном предварительном обучении. Наш код доступен по ссылке https://github.com/Weixin-Liang/Mixture-of-Mamba

English

State Space Models (SSMs) have emerged as efficient alternatives to Transformers for sequential modeling, but their inability to leverage modality-specific features limits their performance in multi-modal pretraining. Here, we propose Mixture-of-Mamba, a novel SSM architecture that introduces modality-aware sparsity through modality-specific parameterization of the Mamba block. Building on Mixture-of-Transformers (W. Liang et al. arXiv:2411.04996; 2024), we extend the benefits of modality-aware sparsity to SSMs while preserving their computational efficiency. We evaluate Mixture-of-Mamba across three multi-modal pretraining settings: Transfusion (interleaved text and continuous image tokens with diffusion loss), Chameleon (interleaved text and discrete image tokens), and an extended three-modality framework incorporating speech. Mixture-of-Mamba consistently reaches the same loss values at earlier training steps with significantly reduced computational costs. In the Transfusion setting, Mixture-of-Mamba achieves equivalent image loss using only 34.76% of the training FLOPs at the 1.4B scale. In the Chameleon setting, Mixture-of-Mamba reaches similar image loss with just 42.50% of the FLOPs at the 1.4B scale, and similar text loss with just 65.40% of the FLOPs. In the three-modality setting, MoM matches speech loss at 24.80% of the FLOPs at the 1.4B scale. Our ablation study highlights the synergistic effects of decoupling projection components, where joint decoupling yields greater gains than individual modifications. These results establish modality-aware sparsity as a versatile and effective design principle, extending its impact from Transformers to SSMs and setting new benchmarks in multi-modal pretraining. Our code can be accessed at https://github.com/Weixin-Liang/Mixture-of-Mamba