Mixture-of-Mamba: Miglioramento dei Modelli di Spazio di Stato Multi-Modali con Sparsità Consapevole della Modalità

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

January 27, 2025

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

cs.AI

Abstract

I modelli dello spazio di stato (SSM) sono emersi come alternative efficienti ai Transformers per la modellazione sequenziale, ma la loro incapacità di sfruttare le caratteristiche specifiche della modalità limita le loro prestazioni nella preformazione multimodale. Qui proponiamo Mixture-of-Mamba, una nuova architettura SSM che introduce una sparità consapevole della modalità attraverso la parametrizzazione specifica della modalità del blocco Mamba. Basandoci su Mixture-of-Transformers (W. Liang et al. arXiv:2411.04996; 2024), estendiamo i vantaggi della sparità consapevole della modalità agli SSM preservando al contempo la loro efficienza computazionale. Valutiamo Mixture-of-Mamba in tre contesti di preformazione multimodale: Transfusion (token di testo e immagine continua intercalati con perdita di diffusione), Chameleon (token di testo e immagine discreta intercalati) e un framework esteso a tre modalità che incorpora il parlato. Mixture-of-Mamba raggiunge costantemente gli stessi valori di perdita a passi di addestramento precedenti con costi computazionali significativamente ridotti. Nel contesto di Transfusion, Mixture-of-Mamba raggiunge una perdita di immagine equivalente utilizzando solo il 34,76% dei FLOPs di addestramento alla scala di 1,4 miliardi. Nel contesto di Chameleon, Mixture-of-Mamba raggiunge una perdita di immagine simile con solo il 42,50% dei FLOPs alla scala di 1,4 miliardi, e una perdita di testo simile con solo il 65,40% dei FLOPs. Nel contesto a tre modalità, MoM raggiunge una perdita di parlato al 24,80% dei FLOPs alla scala di 1,4 miliardi. Il nostro studio di ablatività evidenzia gli effetti sinergici dello sdoppiamento dei componenti di proiezione, dove lo sdoppiamento congiunto porta a guadagni maggiori rispetto alle modifiche individuali. Questi risultati stabiliscono la sparità consapevole della modalità come principio di progettazione versatile ed efficace, estendendo il suo impatto dai Transformers agli SSM e stabilendo nuovi benchmark nella preformazione multimodale. Il nostro codice è disponibile su 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