Source-linked AI summary
Best of Both Worlds: Multimodal Reasoning and Generation via Unified Discrete Flow Matching
Onkar Susladkar, Tushar Prakash, Gayatri Deshmukh, Kiet A. Nguyen, Jiaxun Zhang, Adheesh Juvekar, Tianshu Bao, Lin Chai, Sparsh Mittal, Inderjit S Dhillon, Ismini Lourentzou
TL;DR
Existing multimodal systems separate or entangle understanding and generation, leaving challenges in unified optimization, computational cost, and controllable editing. UniDFlow uses discrete flow matching with task-specific adapters and reference-based preference alignment, achieving broad benchmark performance and zero-shot generalization across multimodal generation and editing tasks.
Problem
Unified multimodal models face mismatched objectives, costly full-model retraining, and representational interference between understanding and generation.
Method
UniDFlow uses discrete flow matching with separate understanding and generation adapters, a lightweight router, and reference-based multimodal preference optimization.
Results
UniDFlow achieves state-of-the-art performance across eight benchmarks and generalizes without explicit task-specific training to inpainting, in-context image generation, reference-based editing, and compositional generation.
Takeaways & Limitations
The framework supports understanding, generation, and editing within one multimodal interface while improving faithfulness and controllability.
Takeaways & Limitations
Alignment quality depends on the diversity and representativeness of supervision signals, and misuse risks remain.
Abstract
from arXiv · showhide
We propose UniDFlow, a unified discrete flow-matching framework for multimodal understanding, generation, and editing. It 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. UniDFlpw 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.
1. Introduction
Existing unified multimodal models struggle with mismatched objectives, costly retraining, and interference between understanding and generation. UniDFlow addresses these issues with separate task-specific adapters, a lightweight router, and reference-guided preference optimization, while showing precise instruction-relevant attention during editing.
- Motivation: Unified models face unstable joint optimization, substantial compute demands, and interference between understanding and generation.These limitations arise from coupling autoregressive and diffusion objectives, full-model updates, and shared parameters.
- Approach: UniDFlow uses a pretrained vision–language model with separate understanding and generation adapters, combining them through a lightweight router.The three-stage training procedure includes understanding, generation, and joint training with reference-based preference optimization.
- Editing behavior: UniDFlow attends more precisely to instruction-relevant regions when editing both coarse objects and fine details.The cited examples include adding a T-shirt and changing a swoosh color.
- Approach: UniDFlow unifies multimodal understanding, text generation, image synthesis, and editing within one discrete diffusion framework.Its contributions also include a single discrete flow-matching objective and stable time conditioning for text and image generation.
- Results: The framework achieves state-of-the-art performance across eight benchmarks spanning understanding, generation, and editing.The contribution statement reports this broad evaluation outcome.
2. Related Work
Prior unified multimodal systems use autoregressive, hybrid, or separate-backbone designs, each reflecting trade-offs between visual-generation efficiency and multimodal understanding. Discrete diffusion instead extends diffusion-style generation to categorical multimodal tokens.
- Discrete diffusion: Discrete diffusion applies masking-based corruption in categorical spaces, motivating parallel mask-and-predict generators for efficient multimodal generation.This extends diffusion beyond continuous pixel or latent spaces.
- Autoregressive unified models: Early unified systems model text and visual tokens autoregressively, but scale poorly for high-resolution image generation.Emu and Chameleon exemplify this single-transformer approach.
- Hybrid frameworks: Hybrid frameworks combine autoregressive text modeling with diffusion-based image generation to retain language understanding while improving visual generation.Examples include EMMA, OmniGenV2, MammothModa2, and BAGEL.
- Understanding and generation separation: Multimodal understanding models commonly rely on separate diffusion backbones for image generation.LLMs and VLMs provide reasoning and visual understanding, while generation remains architecturally separated.
3. Method
UniDFlow casts multimodal understanding, conditional generation, and image editing as discrete denoising with a shared flow-matching objective. Its staged design preserves pretrained reasoning, isolates task-specific adaptation, and adds reference-anchored preference alignment for faithful, controllable outputs.
- 3.1. Preliminaries: Discrete Flow Matching maps corrupted multimodal tokens toward clean text or visual tokens, enabling unified conditioning for understanding, generation, and editing.At inference, sampling starts from noise and applies the learned flow to recover clean tokens.
- 3.2.1. Time-Step Guided RMSNorm: Time-Step Guided RMSNorm injects diffusion time through controlled scale modulation rather than directly perturbing pretrained activations.Zero-initialized modulation exactly recovers the pretrained model at initialization.
- 3.2.2. Stage I: Text Alignment: Stage I adapts the pretrained backbone for diffusion-based understanding using text alignment, LoRA adapters, and KL regularization against the original autoregressive answer distribution.The KL constraint anchors diffusion decoding to the pretrained linguistic manifold while supporting bidirectional, time-conditioned reasoning.
- 3.2.3. Stage II: Vision Alignment: Stage II trains separate generation adapters on discrete visual tokens while freezing the backbone and understanding adapters to limit interference with language and reasoning behavior.The model therefore learns conditional image generation without updating the parameters specialized for Stage I understanding.
- 3.2.4. Stage III: Reference-Based Multimodal Preference Alignment: Stage III combines DFM likelihood terms with reference-anchored multimodal preference alignment over paired preferred and rejected text, image, and reflection outcomes.A lightweight hidden-state router dynamically composes task-specific adapters, while the preference objective increases preferred-over-rejected likelihood margins relative to a frozen reference policy.
- 3.2.4. Stage III: Reference-Based Multimodal Preference Alignment: The final objective jointly trains text, editing, reference, and reflection streams while using multimodal preference alignment to promote instruction following, grounded editing, and consistent behavior.The objective combines the three output-stream DFM likelihood terms with the preference-alignment loss.
4. Experiments
UniDFlow is evaluated across multimodal understanding, generation, editing, and ablations, showing gains over unified and specialized baselines. Qualitative and zero-shot results further indicate precise prompt grounding, faithful editing, and subject preservation.
- Multimodal Understanding: +6.9% on MME-P and +7.0% on MME-S over BAGEL, indicating stronger perceptual and reasoning consistency.
- Multimodal Understanding: +13.3% on MathVista and +6.3% on MMBench over EMMA, demonstrating stronger mathematical and multi-step reasoning.
- Text-to-Image Generation: 0.95 on GENEVAL, outperforming EMMA by +2.2% and MammothModa2 by +9.2% under compositional generation constraints.UniDFlow also exceeds Qwen-Image by 4.0% on GENEVAL and 3.2% on DPGBENCH despite substantially fewer parameters.
- Text-to-Image Editing: +3.7% on GEDIT-BENCH-EN and 4.24 on IMAGEEDIT BENCH, exceeding EMMA and MammothModa2 on instruction satisfaction and perceptual quality.The largest gains occur in Extract and Remove operations, with more precise target isolation and reduced collateral degradation.
- Text-to-Image Editing: Reference-based preference alignment improves editing fidelity and controllability by selecting higher-quality edits that better satisfy user intent.
- Ablations: Larger backbones improve performance across benchmarks, while the 4B model remains competitive, supporting UniDFlow’s parameter-efficient design.
- Ablations: Removing task-specific adapters or the router degrades performance, while a single shared LoRA fails because of severe task entanglement.Performance also degrades when visual or text alignment losses are removed, and excluding reflection-based preference learning reduces editing and faithfulness metrics.
- Ablations: mRef-DPO performs best among the tested alignment methods by stabilizing cross-modal credit assignment and improving alignment and edit precision.
5. Conclusion
UniDFlow unifies multimodal understanding, text-to-image generation, and instruction-guided editing through discrete flow matching, while mRef-DPO aligns outputs relative to a frozen reference policy. Results across six benchmarks support modality-aware preference alignment for reasoning-grounded generation and precise visual editing.
- UniDFlow performs understanding, text-to-image generation, and instruction-guided editing within one discrete flow-matching framework.
- mRef-DPO jointly aligns text and image outputs relative to a frozen reference policy, improving faithfulness and controllability.
- Results across six benchmarks show consistent gains and identify modality-aware preference alignment as critical for reasoning-grounded generation and precise visual editing.
Impact Statement
Unified multimodal systems could support natural interaction, educational and design workflows, and improved human–computer interfaces while potentially reducing computational cost. However, stronger generation and editing also introduce misuse and bias risks, and alignment quality depends on supervision diversity and representativeness.
- Potential benefits: Unified multimodal systems can enhance accessibility, creativity, and productivity through natural multimodal interaction and support for educational and design workflows.They may also improve human–computer interfaces.
- Potential benefits: Parameter-efficient training can reduce computational cost compared with large-scale end-to-end retraining, potentially lowering environmental impact.
- Risks: High-quality synthesis can be misused for deceptive media manipulation, while precise editing may enable subtle alterations that are difficult to detect.
- Risks: Reference-based preference alignment aims to improve faithfulness and controllability, but its quality depends on diverse and representative supervision signals.The paper also notes that misuse risks remain and recommends safeguards including content moderation, bias evaluation, and transparency mechanisms.
A. Implementation Details
The implementation uses a three-stage curriculum progressing from understanding to visual generation and joint alignment, with stabilization techniques applied across training. Figure 10 compares inference throughput with parameter count, where higher throughput and fewer parameters are preferred.
- Training pipeline: The three-stage pipeline progresses from visual instruction following to high-fidelity visual generation and finally joint multimodal understanding and alignment.AdamW, mixed precision, and gradient clipping are used across stages to stabilize training at scale.
- Stage I: Text Alignment: Stage I teaches visual instruction following and grounds text responses in images using variable aspect ratios and resolutions for robustness to diverse inputs.Its learning-rate schedule uses warmup followed by cosine annealing, with λKL = 1.8.
- Stage II: Visual Alignment: Stage II trains visual generation with a diffusion-based objective across multiple resolutions and aspect ratios to support global structure and fine detail.A longer warmup, linear learning-rate schedule, and weight decay support optimization and generalization.
- Stage III: Preference Alignment: Stage III jointly optimizes multimodal comprehension and aligned outputs while expanding resolution and sequence-length ranges for longer-context visual reasoning.It uses cosine annealing with warmup and moderate regularization to consolidate earlier gains while maintaining stability.
- Efficiency comparison: Figure 10 plots inference throughput against parameter count, favoring higher throughput and more compact models.
B. Training Data
Training uses a three-stage data curriculum spanning supervised multimodal instruction learning, image-text pretraining, and preference-based alignment. The reported evaluations cover OCRBenchV2 reasoning subtasks and compare unified models across scales and baselines.
- Curriculum: The data curriculum transitions from supervised multimodal instruction learning to large-scale image-text pretraining and preference-based alignment for understanding, generation, and editing.
- Accounting: Reported token counts include text and discretized or embedded image tokens consumed by the multimodal sequence interface.Aggregate tokens are reported per stage.
- Stage I data: Stage I uses approximately 1.0M MMInstruct image–prompt–answer examples and trains for ≈0.6T total image-plus-text tokens.
- Stage II data: Stage II samples ≈4.5M images from LAION-5B, JourneyDB, and the jackyhate/text-to-image-2M collection for text-to-image generative pretraining.
- Stage III data: Stage III aligns understanding, image generation, and image editing using ≈3.5M curated base tasks aggregated from OpenGPT-4o-Image, AnyEdit-derived edits, and Pico-Banana-400K.
- OCRBenchV2 evaluation: Ours-14B achieves Average = 63.8 on OCRBenchV2, improving over MammothModa2 at 56.1 by +7.7 points across eight visual reasoning subtasks.Performance improves monotonically from Ours-4B to Ours-8B to Ours-14B, while Ours-4B reaches 58.4 Avg. and exceeds MammothModa2 at 56.1 and MUDDIT at 54.7.
C. Full Quantitative Results on GenEval and DPGBench
UniDFlow achieves the best global scores on GenEval and DPGBench, with improvements distributed across fine-grained attributes and compositional criteria rather than driven by a single prompt subset.
- GenEval and DPGBench: UniDFlow achieves the best global score on GenEval and DPGBench while improving consistently across entity, attribute, relation, and compositional criteria.The reported compositional dimensions include object count, color, position, and color-attribute combinations.
D. Additional Results
Additional results examine training and adaptation choices, Stage-III alignment losses, and detailed text-to-image evaluation. These analyses connect scaling and preference alignment with generation and editing behavior.
- Training and adaptation ablations: Scaling training from 0.5T to 3T tokens consistently improves TextGen, GenEval, DPGBench, and ImgEdit-Bench performance.The ablation also studies LoRA rank, with increases from 8 to 32 producing the largest reported gain fragment.
- Stage-III alignment losses: Removing LvRef-DPO weakens visual grounding, while removing LtRef-DPO degrades instruction following in image editing.The reported failures include retaining unwanted content, incomplete candle lighting, and missing or semantically incorrect target objects.
- Text-to-image evaluation: Table 9 provides detailed text-to-image generation evaluation on DPG-BENCH and GENEVAL.The table is identified as a detailed evaluation resource, but the supplied passage does not include its numerical results.
- Stage-III alignment losses: The alignment design is intended to stabilize cross-modal credit assignment for more reliable reasoning-grounded generation.
E. Extended Related Work
The extended context situates UniDFlow among diffusion, vision–language, preference-alignment, and unified multimodal systems. It also includes image-editing curation instructions and qualitative examples spanning understanding, generation, and reasoning-based editing.
- Diffusion for Visual Generation: Diffusion models improve image-generation quality and stability over GANs, while latent diffusion reduces computational cost through compressed latent spaces.
- LLMs and VLMs for Understanding: Vision–language models extend language models to visual inputs by coupling a vision encoder with a language model through lightweight projection layers.
- Preference Alignment: Preference alignment for diffusion models includes likelihood-based, score-space, and subsequent DPO-style variants.
- Unified Models for Understanding and Generation: UniDFlow introduces reference-based preference alignment across text and vision with reflection for stable and faithful generation and editing.
- Image-editing data curation: The curation pipeline produces clean multimodal editing fields and preference pairs containing one chosen output and at least three plausible rejected alternatives.
- Image-editing data curation: Editing instructions require preserving intent, respecting visible source and edited content, specifying changes and constraints, and describing only observable final outcomes.
- Qualitative comparisons: Qualitative figures cover complex reasoning edits, zero-shot multi-subject editing, text-to-image generation, image-to-text results, and complex-scene understanding.