Mistura-de-Mamba: Aprimorando Modelos de Espaço de Estado Multi-Modal com Esparsidade Consciente da Modalidade

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

Resumo

Os Modelos de Espaço de Estados (SSMs) surgiram como alternativas eficientes aos Transformadores para modelagem sequencial, mas sua incapacidade de aproveitar características específicas de modalidade limita seu desempenho na pré-treinamento multi-modal. Aqui, propomos Mixture-of-Mamba, uma arquitetura SSM inovadora que introduz esparsidade consciente de modalidade por meio da parametrização específica da modalidade do bloco Mamba. Baseando-se em Mixture-of-Transformers (W. Liang et al. arXiv:2411.04996; 2024), estendemos os benefícios da esparsidade consciente de modalidade para SSMs, preservando sua eficiência computacional. Avaliamos Mixture-of-Mamba em três cenários de pré-treinamento multi-modal: Transfusão (tokens de texto intercalados e imagens contínuas com perda de difusão), Camaleão (tokens de texto intercalados e imagens discretas) e um framework estendido de três modalidades incorporando fala. Mixture-of-Mamba consistentemente atinge os mesmos valores de perda em etapas de treinamento mais precoces com custos computacionais significativamente reduzidos. No cenário de Transfusão, Mixture-of-Mamba alcança perda de imagem equivalente usando apenas 34,76% dos FLOPs de treinamento na escala de 1,4B. No cenário de Camaleão, Mixture-of-Mamba atinge perda de imagem similar com apenas 42,50% dos FLOPs na escala de 1,4B, e perda de texto similar com apenas 65,40% dos FLOPs. No cenário de três modalidades, MoM iguala a perda de fala em 24,80% dos FLOPs na escala de 1,4B. Nosso estudo de ablação destaca os efeitos sinérgicos do desacoplamento de componentes de projeção, onde o desacoplamento conjunto gera ganhos maiores do que modificações individuais. Esses resultados estabelecem a esparsidade consciente de modalidade como um princípio de design versátil e eficaz, estendendo seu impacto dos Transformadores para SSMs e estabelecendo novos benchmarks no pré-treinamento multi-modal. Nosso código pode ser acessado em 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