OmniMamba: Efficient and Unified Multimodal Understanding and Generation via State Space Models¶
Conference: ECCV2026
Paper: ECCV Paper
Code: https://github.com/hustvl/OmniMamba
Area: Multimodal VLM
Keywords: state space models, unified understanding and generation, decoupled vocabularies, task-specific LoRA, data efficiency
TL;DR¶
OmniMamba shares an autoregressive Mamba-2 backbone while decoupling visual encoding, output vocabularies, and task adapters, achieving 5.50 FID-30K on COCO with approximately 2M image-text pairs as reported for unified-model training, a data accounting boundary that excludes inherited pretrained components.
Background & Motivation¶
A unified vision-language model should both answer questions about images and generate images from text, but these tasks impose different requirements on visual representations. Understanding needs object semantics, relationships, and details useful for answering questions; generation also needs representations that reconstruct spatial structure and texture. Show-o combines text autoregression with discrete diffusion, JanusFlow combines autoregression with rectified flow, and Emu3 uses next-token prediction throughout. These approaches show that task unification does not prescribe a single generation paradigm, although most still use Transformer backbones. As visual tokens and output sequences grow, attention computation and caching become costly; tightly coupled representations and vocabularies may also burden learning with limited data.
The paper does not merely replace a Transformer with a faster module: it asks how a Mamba-2 model with relatively limited foundational capabilities can learn both tasks from small-scale multimodal training. With a shared image-text output vocabulary, the model must learn not to emit text tokens inside an image sequence in addition to learning the content itself. If joint pretraining requires a fixed task ratio, imbalanced understanding and generation datasets can further constrain their effective use. Computational efficiency and data efficiency therefore require separate interventions: the state space backbone addresses the former, while explicit task boundaries and staged training address the latter. The small-data setting adapts an existing language model, visual encoders, and image tokenizer; it does not train every component from scratch.
OmniMamba retains a shared backbone and a common cross-entropy objective without requiring every interface to be shared. Understanding follows a continuous-visual-feature-to-text path, generation follows a text-to-discrete-image-code path, and shared layers retain task-specific low-rank updates. This turns some modality rules that would otherwise require repeated training examples into structural constraints. Core Idea: unify next-token prediction and backbone computation while decoupling visual representations, adaptation parameters, and output spaces to reduce task interference under limited multimodal training data.
Method¶
Overall Architecture¶
The model uses Mamba-2-1.3B with 48 layers for multimodal understanding (MMU) and text-to-image generation (T2I). MMU takes an image and a text instruction and produces text; T2I takes a text condition and produces discrete codes that can be decoded into an image. The paths pass through task-decoupled visual encoding, the shared backbone containing task-specific LoRA, and decoupled output vocabularies. The generation encoder converts real images into supervision tokens during training; inference does not require an additional target image as input. Generated discrete image codes are reconstructed into pixels by the corresponding image decoder, whereas the understanding branch directly emits text tokens.
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
Input["Image and text conditions"] --> Enc["Task-decoupled visual encoding"]
Enc -->|"Continuous understanding features / training image codes"| Backbone["Task-specific LoRA<br/>Shared Mamba-2 backbone"]
Text["T2I inference text condition"] --> Backbone
Backbone --> Vocab["Decoupled output vocabularies"]
Vocab -->|"MMU: text tokens"| Answer["Text answer"]
Vocab -->|"T2I: image tokens"| Decoder["Image decoder โ image"]
The diagram shows representation and prediction paths; two-stage training determines which parameters are updated rather than adding inference stages. Mamba-2 uses state space sequence computation without explicitly constructing all pairwise token attention relationships, providing the basis for long-sequence efficiency. Generation nevertheless remains token-by-token autoregression rather than predicting all image tokens in parallel. Consequently, faster computation with a linear backbone does not by itself establish stronger long-context semantic understanding.
Key Designs¶
1. Task-decoupled visual encoding: use representations suited to semantics and reconstruction
For understanding, the authors follow Prismatic VLMs by combining continuous visual features from DINOv2 and SigLIP. DINOv2 supplies spatial information, SigLIP supplies language-related semantics, and a projection layer connects the visual representation to the language backbone. The main understanding experiments use input resolution 384, with the aim of supplying visual conditions useful for answering questions. Understanding is not forced to first compress the image into discrete reconstruction-oriented codes, so a single compression objective need not serve every semantic requirement.
For generation, a LlamaGen-trained VQVAE encodes real images into discrete image tokens; the main generation resolution is 256. These tokens provide next-token supervision, and inference produces a sequence in the same code space for the decoder. Continuous understanding features and discrete generation codes are not interchangeable: they serve conditional perception and image reconstruction, respectively. The tokenizer was pretrained on ImageNet and further fine-tuned on 50M LAION-COCO images and 10M internal high-aesthetic-quality examples. Thus, the claim that generation training uses only 83K COCO images describes the adaptation stage of the generation model, not the entire cost of learning its visual representation.
2. Task-specific LoRA: adapt input projections before shared state space computation
The input linear projection of every Mamba-2 layer receives separate LoRA modules for MMU and T2I. Understanding activates only MMU LoRA; generation activates only T2I LoRA, with the other path disabled. The low-rank update works alongside the original input projection, allowing both tasks to share most parameters while adjusting the features entering each layer. The authors hypothesize that explicitly task-parameterized input projections help the state space model adapt to different task requirements with less data.
This is not a mixture-of-experts system that learns content-dependent routing for each token; a known task selects its corresponding adapter. Routing separates part of the parameter updates for understanding and generation, reducing the risk of competing objectives repeatedly pulling the same lightweight adaptation parameters in different directions. The LoRA rank is 8, and the paper reports only a 0.65% parameter increase. This increase describes adapter size, not a guarantee that only 0.65% of model parameters are updated throughout training; the second-stage update scope is discussed separately. The no-LoRA comparison in Table 5 gives this architectural choice direct evidence beyond the proposed explanation involving task interference.
3. Decoupled output vocabularies: turn output modality into an interface constraint
A large shared vocabulary requires the model to decide both what comes next and whether it belongs to text or images. OmniMamba uses separate text vocabularies and image codebooks, selecting the corresponding output head and vocabulary by task. MMU predicts text tokens and T2I predicts image tokens, excluding the wrong modality at the output interface. The authors observe that the unified-vocabulary configuration sometimes emits text-related tokens during visual generation, requiring additional post-processing. Two vocabularies do not automatically ensure semantic correctness in images, but they reduce the burden of learning and correcting these invalid outputs.
Both paths still use next-token prediction, differing in their conditioning representations and permitted target sets.
Unification therefore applies to sequence computation and the training objective, without requiring images and text to share every discrete index.
Input formats distinguish tasks with [MMU] and [T2I], and mark image and text boundaries with [SOI], [EOI], [SOT], and [EOT].
Understanding places image conditions before text, while generation places text conditions before the predicted image codes.
The extracted full T2I format line is corrupted, so this note explains only the verifiable ordering and markers rather than reconstructing a purported exact string.
A Worked Example¶
The following illustrates the mechanism rather than introducing a quantitative experiment: provide a street-scene image for description, then separately provide a street-scene description for image generation. For description, DINOv2 and SigLIP process the image, and projected continuous features become visual conditions for Mamba-2. The model activates MMU LoRA and predicts an answer through the text vocabulary; it does not need to reconstruct the input pixels. For generation, text enters the shared backbone, which instead uses T2I LoRA and the image output vocabulary. Each prediction is an image code rather than a natural-language word, and the completed code sequence goes to the image decoder. During training, the VQVAE encodes a real image into target codes; at inference these targets do not exist and must be generated progressively. Sharing a language backbone across these calls does not establish support for arbitrary interleaved long image-text conversations or image editing.
Loss & Training¶
Stage 1 performs separate task-specific pretraining with the core Mamba-2 frozen, preventing differently sized datasets from competing over backbone updates at initialization. MMU trains the visual projection and MMU LoRA with the T2I path disabled; T2I trains its LoRA and image prediction components with the MMU path disabled. Section 3.5 alternates between the terms "image head" and "image decoder," leaving it unclear whether the latter specifically denotes the VQVAE pixel decoder. This note therefore does not assert that the pretrained image decoder must be updated. Stage 2 jointly fine-tunes the tasks while retaining independent LoRA paths, summing their losses for a unified backward pass. Section 3.5 first states that visual encoders are frozen and other modules are trained, but later describes sharing through a "frozen backbone"; the backbone's actual frozen state requires implementation verification.
Section 3.6 specifies standard next-token cross-entropy for both tasks, without a diffusion denoising objective. The following is standard notation for that verbal description, not a numbered equation from the paper; \(y_i\) is a target token for the selected task and \(c\) is its input condition.
Section 4.1 lists 676K understanding pretraining pairs, comprising 118K from COCO and 558K from LLaVA-1.5; the rationale in Section 3.5 instead gives 665K, an inconsistent count. Understanding fine-tuning uses 665K LLaVA-1.5 multi-turn conversations, 220K LVIS-Instruct-4V examples, and 400K LRV-Instruct examples; generation uses 83K MS-COCO 2014 images. These are different sample counts across stages and datasets, so their sum should not be treated as a unique image-text pair count, nor should the authors' "fewer than 2M" claim be read as total pretraining across every component. All stages use AdamW with momentum parameters 0.9 and 0.95, cosine scheduling with warm-up, weight decay 0, and gradient clipping at 1.0. Training uses BF16 on NVIDIA A800 GPUs, whereas speed measurements use FP16 on an NVIDIA 4090; these are different operating conditions.
Key Experimental Results¶
Main Results¶
The following excerpt is from Table 1, page 10. All entries are unified models with 1.3B language backbones; higher is better, but resolution, pretrained backbone, and data differ, so this is not a controlled architecture-only comparison.
| Model | Understanding resolution | POPE | MME-P | VQAv2 test | GQA | MMMU |
|---|---|---|---|---|---|---|
| Show-o | 512 | 80.0 | 1097.2 | 69.4 | 58.0 | 26.7 |
| JanusFlow | 384 | 88.0 | 1333.1 | 79.8 | 60.3 | 29.3 |
| OmniMamba | 384 | 86.3 | 1290.6 | 77.7 | 60.8 | 30.6 |
OmniMamba exceeds Show-o on every listed metric but does not uniformly exceed JanusFlow: GQA and MMMU are higher, while POPE, MME-P, and VQAv2 are lower. The next excerpt is from Table 2, page 11; FID-30K uses 30K randomly selected text prompts from the COCO validation set, and lower is better. "Generation training images" refers to the task data reported in the table, not all pretraining dependencies; OmniMamba and U-ViT train on COCO's training set, while other methods use different protocols.
| Model | Parameters | Generation training images | FID-30K |
|---|---|---|---|
| Show-o | 1.3B | 35M | 9.24 |
| U-ViT | 45M | 83K COCO | 5.95 |
| OmniMamba | 1.3B | 83K COCO | 5.50 |
Thus, 5.50 supports strong in-distribution COCO generation, not a general claim of superior open-domain prompt following, aesthetics, or face generation. Show-o's 35M generation images in Table 2 and the 2B unified-model data discussed in the text use different accounting boundaries and should not be conflated.
Ablation Study¶
This excerpt comes from Table 5, page 13. Section 4.6 describes a CLIP understanding encoder, resolution 224, and identical reduced training steps for ablations; only the 130M and 370M rows are selected here, without mixing in the main experiment's 1.3B row.
| Experiment | Backbone | Decoupled encoders | Decoupled vocabularies | Task LoRA | POPE | MME | GQA | FID-30K |
|---|---|---|---|---|---|---|---|---|
| 1 | Mamba2-130M | No | Yes | Yes | 65.6 | 598 | 35.7 | 31.4 |
| 2 | Mamba2-130M | Yes | Yes | Yes | 80.6 | 930 | 46.1 | 23.7 |
| 5 | Mamba2-370M | Yes | No | Yes | 80.8 | 1036 | 53.6 | 19.1 |
| 6 | Mamba2-370M | Yes | Yes | No | 81.2 | 1003 | 54.0 | 14.4 |
| 7 | Mamba2-370M | Yes | Yes | Yes | 81.9 | 1100 | 55.3 | 10.3 |
Experiments 1 โ 2 support decoupled encoders; 5 โ 7 support separate vocabularies; 6 โ 7 support task-specific LoRA. Adding task-specific LoRA reduces FID-30K from 14.4 to 10.3, a decrease of 4.1, while increasing MME from 1003 to 1100. The vocabulary-ablation paragraph cites Experiments 1 and 3, but those rows do not toggle vocabulary sharing; this note follows the table fields and compares Experiments 5 and 7, preserving the discrepancy.
Key Findings¶
- Table 4, page 12: on one NVIDIA 4090 with FP16, batch size 16, and resolution 256, OmniMamba reaches 5.68 image/s, versus 0.81 for Show-o and 1.02 for JanusFlow.
- Table 3, page 11: at sequence length 16384, OmniMamba reaches 248.1 token/s, versus 57.40 token/s for SmolLM2-1.7B using FlashAttention-2 and KV cache.
- Section 4.5 reports a 119.2-fold speedup over Show-o at sequence length 16k and 63.0% memory savings at 23k; these are different measurement points, and Show-o does not use FlashAttention-2 in that comparison.
- The long-sequence throughput test removes the output-length limit without jointly establishing semantic correctness at extreme context lengths; speed, memory, and task quality require separate evaluation.
Highlights & Insights¶
- Data efficiency can come from reducing the rules a model must learn, not just its parameter count. Separate vocabularies exclude the wrong modality by construction, an interface design transferable to other discrete multimodal systems.
- A shared backbone can coexist with task-specific input projections. LoRA adjusts features before state space computation while duplicating fewer parameters than combining two complete models.
- Separate pretraining followed by unified fine-tuning avoids task competition from randomly initialized interfaces. This is particularly relevant when dataset sizes are asymmetric but reliable pretrained components already exist.
Limitations & Future Work¶
- Section 6 acknowledges that COCO's limited coverage constrains image quality, particularly for faces; lower FID does not imply reliable generation of every visual subject.
- The Mamba-2 foundation model was trained on sequences of at most 2048 tokens. The paper demonstrates the computational feasibility of long outputs, not adequately trained long-context understanding or CoT capabilities.
- Total component-level data costs substantially exceed the 83K generation-adaptation images. Assessing data savings requires separate accounting for the language model, visual encoders, tokenizer, and unified model.
- No independently controlled table for the two-stage training strategy is verifiable here; conflicting descriptions of freezing, pretraining sample counts, and vocabulary-ablation indices also complicate reproduction.
- Further comparisons should match pretraining dependencies and training budgets and use broader generation prompts, separating system throughput gains from representation and data advantages.
Related Work & Insights¶
- Relation to Janus / JanusFlow: follows task-decoupled visual encoding, but uses a Mamba-2 sequence backbone and purely autoregressive image-token prediction instead of rectified flow.
- Relation to Show-o / Emu3: does not require a shared vocabulary for unification. The ablations favor explicit output-modality constraints under limited training budgets, without proving that shared vocabularies are inherently worse at scale.
- Relation to Cobra / Prismatic VLMs: the former informs Mamba-based multimodal understanding and the data recipe, while the latter motivates DINOv2 and SigLIP feature fusion; OmniMamba extends understanding toward unified generation.
- Relation to LlamaGen: reuses its image tokenizer to simplify learning the generation interface and inherits its pretraining dependencies; this is not an experiment in learning a visual codebook from scratch.
- Classification: the central task is multimodal understanding and generation in one model, so Multimodal VLM covers the contribution better than Image Generation alone.
Rating¶
- Novelty: 4/5. Combines a linear state space backbone with explicit task decoupling, with novelty primarily in unified system design.
- Experimental Thoroughness: 3/5. Covers understanding, generation, speed, and architectural ablations, but data accounting and controlled training-strategy comparisons remain incomplete.
- Writing Quality: 3/5. The main argument is clear, but conflicting freezing descriptions and ablation indices affect reproducibility.
- Value: 4/5. Offers practical design choices for low-budget unified models built on pretrained components, without establishing broad long-sequence competence.