OmniMamba: Compreensão e Geração Multimodal Eficiente e Unificada via Modelos de Espaço de Estados

OmniMamba: Efficient and Unified Multimodal Understanding and Generation via State Space Models

March 11, 2025

Autores: Jialv Zou, Bencheng Liao, Qian Zhang, Wenyu Liu, Xinggang Wang

cs.AI

Resumo

Os recentes avanços em modelos unificados de compreensão multimodal e geração visual (ou geração multimodal) têm sido limitados por sua complexidade computacional quadrática e dependência de grandes volumes de dados de treinamento. Apresentamos o OmniMamba, o primeiro modelo de geração multimodal baseado em arquitetura linear que gera tanto texto quanto imagens por meio de um paradigma unificado de previsão do próximo token. O modelo aproveita totalmente a alta eficiência computacional e de memória do Mamba-2, estendendo suas capacidades da geração de texto para a geração multimodal. Para abordar a ineficiência de dados dos modelos unificados existentes, propomos duas inovações principais: (1) vocabulários desacoplados para orientar a geração específica de cada modalidade, e (2) LoRA específico para tarefas, visando uma adaptação eficiente em termos de parâmetros. Além disso, introduzimos uma estratégia de treinamento em duas etapas desacopladas para mitigar o desequilíbrio de dados entre as duas tarefas. Equipado com essas técnicas, o OmniMamba alcança desempenho competitivo com o JanusFlow, superando o Show-o em diversos benchmarks, apesar de ter sido treinado com apenas 2 milhões de pares de imagem-texto, o que é 1.000 vezes menor do que o Show-o. Notavelmente, o OmniMamba se destaca com uma eficiência de inferência excepcional, alcançando uma aceleração de até 119,2 vezes e uma redução de 63% na memória GPU para geração de sequências longas em comparação com modelos baseados em Transformers. O código e os modelos estão disponíveis em https://github.com/hustvl/OmniMamba.

English

Recent advancements in unified multimodal understanding and visual generation (or multimodal generation) models have been hindered by their quadratic computational complexity and dependence on large-scale training data. We present OmniMamba, the first linear-architecture-based multimodal generation model that generates both text and images through a unified next-token prediction paradigm. The model fully leverages Mamba-2's high computational and memory efficiency, extending its capabilities from text generation to multimodal generation. To address the data inefficiency of existing unified models, we propose two key innovations: (1) decoupled vocabularies to guide modality-specific generation, and (2) task-specific LoRA for parameter-efficient adaptation. Furthermore, we introduce a decoupled two-stage training strategy to mitigate data imbalance between two tasks. Equipped with these techniques, OmniMamba achieves competitive performance with JanusFlow while surpassing Show-o across benchmarks, despite being trained on merely 2M image-text pairs, which is 1,000 times fewer than Show-o. Notably, OmniMamba stands out with outstanding inference efficiency, achieving up to a 119.2 times speedup and 63% GPU memory reduction for long-sequence generation compared to Transformer-based counterparts. Code and models are released at https://github.com/hustvl/OmniMamba