OmniMamba: Comprensión y Generación Multimodal Eficiente y Unificada mediante Modelos de Espacio 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

Resumen

Los recientes avances en los modelos de comprensión multimodal unificada y generación visual (o generación multimodal) se han visto obstaculizados por su complejidad computacional cuadrática y su dependencia de datos de entrenamiento a gran escala. Presentamos OmniMamba, el primer modelo de generación multimodal basado en una arquitectura lineal que genera tanto texto como imágenes a través de un paradigma unificado de predicción del siguiente token. El modelo aprovecha al máximo la alta eficiencia computacional y de memoria de Mamba-2, extendiendo sus capacidades desde la generación de texto hasta la generación multimodal. Para abordar la ineficiencia de datos de los modelos unificados existentes, proponemos dos innovaciones clave: (1) vocabularios desacoplados para guiar la generación específica de cada modalidad, y (2) LoRA específico de tareas para una adaptación eficiente en parámetros. Además, introducimos una estrategia de entrenamiento en dos etapas desacopladas para mitigar el desequilibrio de datos entre las dos tareas. Equipado con estas técnicas, OmniMamba logra un rendimiento competitivo con JanusFlow mientras supera a Show-o en diversos benchmarks, a pesar de haber sido entrenado con solo 2 millones de pares de imagen-texto, lo que es 1,000 veces menos que Show-o. Cabe destacar que OmniMamba sobresale con una eficiencia de inferencia excepcional, logrando una aceleración de hasta 119.2 veces y una reducción del 63% en la memoria de GPU para la generación de secuencias largas en comparación con los modelos basados en Transformer. El código y los modelos están disponibles en 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