ICML 2026
We introduce UniDFlow, a discrete flow-matching model that unifies image understanding, generation, and reasoning-guided editing. Separate text and image LoRA adapters preserve specialized capabilities, while a lightweight Mixture-of-LoRA router composes them during diffusion. Multimodal Reference-based DPO (mRef-DPO) uses a visual reference and structured reflection traces to align edits with instructions, geometry, and physical transformations.
One model supports visual understanding, text-to-image generation, and edits guided by reasoning.
We propose UniDFlow, a unified discrete flow-matching framework for multimodal understanding, generation, and editing. UniDFlow decouples understanding and generation via task-specific low-rank adapters, avoiding objective interference and representation entanglement, while a novel reference-based multimodal preference alignment optimizes relative outcomes under identical conditioning, improving faithfulness and controllability without large-scale retraining. UniDFlow achieves SOTA performance across eight benchmarks and exhibits strong zero-shot generalization to tasks including inpainting, in-context image generation, reference-based editing, and compositional generation, despite no explicit task-specific training.
UniDFlow learns understanding and generation through separate adapters on a frozen backbone, then combines them for reasoning-guided editing. The three training stages isolate the two objectives before aligning their joint behavior.
Text LoRA adapters learn to answer visual instructions using discrete flow matching. The backbone remains frozen, and KL regularization preserves linguistic ability while the model develops multimodal reasoning.
This stage establishes the understanding pathway that later guides image editing.
A separate image LoRA learns conditional generation in discrete visual-token space. The text adapters remain frozen, preserving their reasoning capability while image generation is trained.
Separating the adapters prevents the understanding and generation objectives from competing for the same parameters.
A lightweight Mixture-of-LoRA (MoRA) router dynamically composes the text and image adapters during diffusion. The model can combine visual reasoning with image synthesis at each step.
Multimodal Reference-based DPO (mRef-DPO) compares outcomes under identical conditioning against a frozen reference policy and a visual reference image. Structured reflection traces train the model to reason through geometric, physical, and temporal changes before producing the edit.
UniDFlow (4B) is evaluated on general multimodal question answering and visual reasoning. Compared with BAGEL, it achieves gains of +6.9% on MME-P and +7.0% on MME-S.
| Model | Parameters | MME-P | MME-S | MMBench | MMMU | MM-Vet | MathVista | MMVP |
|---|---|---|---|---|---|---|---|---|
| Qwen2.5-VL (Bai et al., 2025b) | 3B | — | 2157 | 79.1 | 53.1 | 61.8 | 62.3 | — |
| BLIP-3 (Xue et al., 2024) | 4B | — | — | 76.8 | 41.1 | — | 39.6 | — |
| DeepSeek-VL2 (Wu et al., 2024b) | 4B | — | — | 51.1 | 60.0 | 62.8 | — | — |
| Qwen3-VL (Bai et al., 2025c) | 4B | — | — | 85.1 | 64.1 | 72.5 | — | — |
| VILA-U (Wu et al., 2024a) | 7B | 1336 | — | 66.6 | 32.2 | 27.7 | — | 22.0 |
| Chameleon (Team, 2024a) | 7B | — | — | 35.7 | 28.4 | 8.3 | — | 0.0 |
| Janus-Pro (Chen et al., 2025a) | 7B | 1567 | — | 79.2 | 41.0 | 50.0 | — | — |
| TokenFlow-XL (Geyer et al., 2023) | 13B | 1546 | — | 68.9 | 38.7 | 40.7 | — | — |
| BAGEL (Deng et al., 2025) | 7B | 1687 | 2388 | 85.0 | 55.3 | 67.2 | 73.1 | 69.3 |
| OmniGen-v2 (Wu et al., 2025) | 8B | — | — | 53.1 | 61.5 | — | — | — |
| EMMA (He et al., 2025) | 4B | — | — | 85.8 | 65.1 | 73.0 | 75.8 | — |
| MammothModa-2 (Shen et al., 2025) | 4B | 1753 | 1998 | 86.6 | 71.23 | 79.4 | 81.8 | 77.5 |
| Muddit (Shi et al., 2025) | 4B | 1700 | 1832 | 82.8 | 66.6 | 76.2 | 79.1 | 74.1 |
| UniDFlow | 4B | 1803 | 2555 | 91.2 | 74.3 | 82.7 | 85.9 | 80.2 |
UniDFlow (4B) outperforms same-scale unified competitors and larger generative models on text-to-image generation.
| Model | Parameters | GenEval | DPGBench |
|---|---|---|---|
| DALL-E 3 (Betker et al., 2023) | — | 0.67 | 83.50 |
| SD3-Medium (Esser et al., 2024) | 2B | 0.74 | 80.43 |
| Qwen-Image(-RL) (et al., 2025) | 7B+20B | 0.91 | 88.32 |
| TokenFlow-XL (Geyer et al., 2023) | 14B | 0.55 | — |
| Janus-Pro-7B (Chen et al., 2025a) | 7B | 0.80 | 84.19 |
| Bagel (Deng et al., 2025) | 7B+7B | 0.88 | 87.74 |
| OmniGen2/V2 (Wu et al., 2025) | 3B+4B | 0.78 | 83.57 |
| MammothModa-2 (Shen et al., 2025) | 8B+3B+2B | 0.87 | 87.20 |
| EMMA (He et al., 2025) | 4B | 0.93 | 85.63 |
| MUDDIT (Shi et al., 2025) | 8B | 0.90 | 86.37 |
| UniDFlow | 4B | 0.95 | 91.19 |
Qualitative comparison of compositional text-to-image generation and editing. Prompts require precise grounding of attributes and spatial relations (red text). UniDFlow consistently adheres to these constraints while maintaining realistic structure and visual fidelity, outperforming prior unified baselines.
Image generation with UniDFlow.
Reasoning-driven image editing, highlighting temporal, geometric, and physical transformations handled by UniDFlow.
Image-to-text generated results with UniDFlow.
@inproceedings{susladkar2026best,
title={Best of Both Worlds: Multimodal Reasoning and Generation via Unified Discrete Flow Matching},
author={Susladkar, Onkar and Prakash, Tushar and Deshmukh, Gayatri and Nguyen, Kiet A. and Zhang, Jiaxun and Juvekar, Adheesh and Bao, Tianshu and Chai, Lin and Mittal, Sparsh and Dhillon, Inderjit S and Lourentzou, Ismini},
booktitle={Proceedings of the International Conference on Machine Learning (ICML)},
year={2026}
}