Mélange-de-Mamba : Amélioration des Modèles d'Espace d'États Multi-Modaux avec une Sparsité Consciente de la Modalité

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

Résumé

Les modèles d'espace d'états (State Space Models - SSM) ont émergé en tant qu'alternatives efficaces aux Transformers pour la modélisation séquentielle, mais leur incapacité à exploiter les caractéristiques spécifiques à la modalité limite leurs performances dans la pré-formation multi-modale. Ici, nous proposons Mixture-of-Mamba, une nouvelle architecture SSM qui introduit une parcimonie consciente de la modalité grâce à la paramétrisation spécifique à la modalité du bloc Mamba. En nous appuyant sur Mixture-of-Transformers (W. Liang et al. arXiv:2411.04996; 2024), nous étendons les avantages de la parcimonie consciente de la modalité aux SSM tout en préservant leur efficacité computationnelle. Nous évaluons Mixture-of-Mamba à travers trois configurations de préformation multi-modale : Transfusion (jetons de texte et d'image continue entrelacés avec perte de diffusion), Chameleon (jetons de texte et d'image discrète entrelacés), et un cadre à trois modalités étendu incorporant la parole. Mixture-of-Mamba atteint de manière cohérente les mêmes valeurs de perte à des étapes d'entraînement plus précoces avec des coûts computationnels considérablement réduits. Dans le cadre de Transfusion, Mixture-of-Mamba atteint une perte d'image équivalente en n'utilisant que 34,76 % des FLOPs d'entraînement à l'échelle de 1,4 milliard. Dans le cadre de Chameleon, Mixture-of-Mamba atteint une perte d'image similaire avec seulement 42,50 % des FLOPs à l'échelle de 1,4 milliard, et une perte de texte similaire avec seulement 65,40 % des FLOPs. Dans le cadre des trois modalités, MoM atteint une perte de parole à 24,80 % des FLOPs à l'échelle de 1,4 milliard. Notre étude d'ablation met en évidence les effets synergiques du découplage des composants de projection, où le découplage conjoint entraîne des gains plus importants que les modifications individuelles. Ces résultats établissent la parcimonie consciente de la modalité comme un principe de conception polyvalent et efficace, étendant son impact des Transformers aux SSM et établissant de nouveaux benchmarks en préformation multi-modale. Notre code est accessible sur 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