Mischung aus Mamba: Verbesserung von Multi-Modalen Zustandsraummodellen durch modalitätsbewusste Sparsamkeit

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

January 27, 2025

Autoren: Weixin Liang, Junhong Shen, Genghan Zhang, Ning Dong, Luke Zettlemoyer, Lili Yu

cs.AI

Zusammenfassung

Zustandsraummodelle (SSMs) haben sich als effiziente Alternativen zu Transformatoren für die sequenzielle Modellierung herausgestellt, aber ihre Unfähigkeit, modalitätsspezifische Merkmale zu nutzen, begrenzt ihre Leistungsfähigkeit beim multimodalen Pretraining. Hier schlagen wir Mixture-of-Mamba vor, eine neuartige SSM-Architektur, die modalitätsspezifische Sparsamkeit durch die modalitätsspezifische Parametrisierung des Mamba-Blocks einführt. Aufbauend auf Mixture-of-Transformers (W. Liang et al. arXiv:2411.04996; 2024) erweitern wir die Vorteile der modalitätsspezifischen Sparsamkeit auf SSMs und bewahren gleichzeitig ihre Rechenleistung. Wir evaluieren Mixture-of-Mamba in drei multimodalen Pretraining-Szenarien: Transfusion (abwechselnde Text- und kontinuierliche Bild-Token mit Diffusionsverlust), Chameleon (abwechselnde Text- und diskrete Bild-Token) und einem erweiterten dreimodalen Framework, das Sprache integriert. Mixture-of-Mamba erreicht konsistent die gleichen Verlustwerte in früheren Trainingsschritten bei signifikant reduzierten Rechenkosten. Im Transfusion-Szenario erzielt Mixture-of-Mamba mit nur 34,76 % der Trainings-FLOPs im Maßstab von 1,4 Mrd. äquivalente Bildverluste. Im Chameleon-Szenario erreicht Mixture-of-Mamba ähnliche Bildverluste mit nur 42,50 % der FLOPs im Maßstab von 1,4 Mrd. und ähnliche Textverluste mit nur 65,40 % der FLOPs. Im dreimodalen Szenario erreicht MoM Sprachverluste bei 24,80 % der FLOPs im Maßstab von 1,4 Mrd. Unsere Ablationsstudie hebt die synergistischen Effekte der Entkopplung von Projektionskomponenten hervor, wobei die gemeinsame Entkopplung größere Gewinne bringt als einzelne Modifikationen. Diese Ergebnisse etablieren modalitätsspezifische Sparsamkeit als vielseitiges und effektives Designprinzip, das ihre Auswirkungen von Transformatoren auf SSMs ausweitet und neue Maßstäbe im multimodalen Pretraining setzt. Unser Code ist unter https://github.com/Weixin-Liang/Mixture-of-Mamba verfügbar.

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