OmniMamba: Эффективное и унифицированное мультимодальное понимание и генерация с использованием моделей пространства состояний

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

March 11, 2025

Авторы: Jialv Zou, Bencheng Liao, Qian Zhang, Wenyu Liu, Xinggang Wang

cs.AI

Аннотация

Последние достижения в области унифицированного мультимодального понимания и моделей визуальной генерации (или мультимодальной генерации) были ограничены их квадратичной вычислительной сложностью и зависимостью от крупномасштабных обучающих данных. Мы представляем OmniMamba — первую модель мультимодальной генерации на основе линейной архитектуры, которая генерирует как текст, так и изображения через унифицированную парадигму предсказания следующего токена. Модель полностью использует высокую вычислительную и память-эффективность Mamba-2, расширяя её возможности от генерации текста до мультимодальной генерации. Для решения проблемы неэффективности данных в существующих унифицированных моделях мы предлагаем два ключевых нововведения: (1) разделённые словари для управления генерацией, специфичной для модальности, и (2) LoRA для задач, обеспечивающую параметрически эффективную адаптацию. Кроме того, мы вводим разделённую двухэтапную стратегию обучения для смягчения дисбаланса данных между двумя задачами. Оснащённая этими методами, OmniMamba демонстрирует конкурентоспособные результаты с JanusFlow, превосходя Show-o на различных бенчмарках, несмотря на обучение всего на 2 миллионах пар изображение-текст, что в 1000 раз меньше, чем у Show-o. Примечательно, что OmniMamba выделяется выдающейся эффективностью вывода, достигая ускорения до 119,2 раз и сокращения использования памяти GPU на 63% для генерации длинных последовательностей по сравнению с моделями на основе Transformer. Код и модели доступны по адресу 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