Mezcla de Mamba: Mejorando los Modelos de Espacio de Estados Multi-Modales con Esparseidad Consciente de la Modalidad

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

January 27, 2025

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

cs.AI

Resumen

Los Modelos de Espacio de Estados (SSMs) han surgido como alternativas eficientes a los Transformadores para modelado secuencial, pero su incapacidad para aprovechar las características específicas de la modalidad limita su rendimiento en el preentrenamiento multi-modal. Aquí, proponemos Mezcla-de-Mamba, una arquitectura SSM novedosa que introduce dispersión consciente de la modalidad a través de la parametrización específica de la modalidad del bloque Mamba. Basándonos en Mezcla-de-Transformadores (W. Liang et al. arXiv:2411.04996; 2024), extendemos los beneficios de la dispersión consciente de la modalidad a los SSMs mientras preservamos su eficiencia computacional. Evaluamos Mezcla-de-Mamba en tres configuraciones de preentrenamiento multi-modal: Transfusión (tokens de texto entrelazados e imágenes continuas con pérdida de difusión), Camaleón (tokens de texto entrelazados e imágenes discretas) y un marco de tres modalidades extendido que incorpora habla. Mezcla-de-Mamba alcanza consistentemente los mismos valores de pérdida en pasos de entrenamiento más tempranos con costos computacionales significativamente reducidos. En la configuración de Transfusión, Mezcla-de-Mamba logra una pérdida de imagen equivalente utilizando solo el 34.76% de los FLOPs de entrenamiento en la escala de 1.4B. En la configuración de Camaleón, Mezcla-de-Mamba alcanza una pérdida de imagen similar con solo el 42.50% de los FLOPs en la escala de 1.4B, y una pérdida de texto similar con solo el 65.40% de los FLOPs. En la configuración de tres modalidades, MoM iguala la pérdida de habla al 24.80% de los FLOPs en la escala de 1.4B. Nuestro estudio de ablación destaca los efectos sinérgicos de desacoplar componentes de proyección, donde el desacoplamiento conjunto produce mayores beneficios que las modificaciones individuales. Estos resultados establecen la dispersión consciente de la modalidad como un principio de diseño versátil y efectivo, extendiendo su impacto desde los Transformadores a los SSMs y estableciendo nuevos puntos de referencia en el preentrenamiento multi-modal. Nuestro código está disponible en 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