Mixture-of-Mamba: Het verbeteren van multi-modale toestandsruimtemodellen met modality-bewuste spaarzaamheid.

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

January 27, 2025

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

cs.AI

Samenvatting

Staatsruimtemodellen (SSM's) zijn naar voren gekomen als efficiënte alternatieven voor Transformers voor sequentiële modellering, maar hun onvermogen om modus-specifieke kenmerken te benutten beperkt hun prestaties bij multimodale vooraftraining. Hier stellen we Mixture-of-Mamba voor, een nieuw SSM-architectuur die modaliteitsbewuste spaarzaamheid introduceert door modaliteitspecifieke parameterisatie van het Mamba-blok. Voortbouwend op Mixture-of-Transformers (W. Liang et al. arXiv:2411.04996; 2024), breiden we de voordelen van modaliteitsbewuste spaarzaamheid uit naar SSM's terwijl we hun rekenkundige efficiëntie behouden. We evalueren Mixture-of-Mamba in drie multimodale vooraftrainingsscenario's: Transfusion (geïnterlinieerde tekst en continue beeldtokens met diffusieverlies), Chameleon (geïnterlinieerde tekst en discrete beeldtokens), en een uitgebreid driedelige modusraamwerk met spraak. Mixture-of-Mamba bereikt consequent dezelfde verlieswaarden in eerdere trainingsfasen met aanzienlijk verminderde rekenkosten. In de Transfusion-instelling behaalt Mixture-of-Mamba een equivalent beeldverlies met slechts 34.76% van de trainings-FLOPs op schaal 1.4B. In de Chameleon-instelling bereikt Mixture-of-Mamba vergelijkbaar beeldverlies met slechts 42.50% van de FLOPs op schaal 1.4B, en vergelijkbaar tekstverlies met slechts 65.40% van de FLOPs. In de driedelige modusinstelling bereikt MoM spraakverlies met 24.80% van de FLOPs op schaal 1.4B. Onze ablatiestudie benadrukt de synergetische effecten van het ontkoppelen van projectiecomponenten, waarbij gezamenlijk ontkoppelen grotere voordelen oplevert dan individuele aanpassingen. Deze resultaten vestigen modaliteitsbewuste spaarzaamheid als een veelzijdig en effectief ontwerpprincipe, waarbij de impact ervan wordt uitgebreid van Transformers naar SSM's en nieuwe maatstaven worden gezet in multimodale vooraftraining. Onze code is beschikbaar op 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