Best of Both Worlds: Multimodal Reasoning and Generation via
Unified Discrete Flow Matching

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 is trained with a three-stage pipeline designed to bridge the gap between high-level reasoning and high-fidelity generation.

Stage I: Text Alignment (Reasoning Foundation)

The first stage focuses exclusively on multimodal understanding to establish a strong reasoning foundation.

• Objective: We align a frozen vision-language backbone to follow visual instructions using a discrete flow-matching objective.
• Mechanism: We introduce and train specialized LoRA_text adapters while keeping the base model parameters frozen.
• Stabilization: To prevent semantic drift and ensure our model retained its original linguistic intelligence, we regularize the training with a KL divergence loss.

Stage II: Vision Alignment (Generative Capability)

Stage II endows the model with the ability to synthesize images.

• Objective: We adapt the model for conditional generation within a discrete visual token space. By operating in a discrete latent space, we enable high-fidelity image synthesis that seamlessly integrates with the backbone's token-based architecture.
• Isolation: To avoid objective interference, we keep Stage I reasoning adapters frozen and introduce a separate set of LoRA_image adapters.

Stage III: Multimodal Preference Alignment (Refined Editing)

Stage III improves reasoning-based capabilities to enable instruction-based image editing and complex reasoning tasks.

• Dynamic routing (MoRA): Since understanding and generation require different specializations, we propose a lightweight Mixture-of-LoRA (MoRA) router, which dynamically composes our task-specific adapters at every step of the diffusion process.
• mRef-DPO alignment: We introduce a novel multimodal Reference-based Direct Preference Optimization (mRef-DPO) method to teach the model to distinguish between faithful edits and subtle errors by comparing outcomes against a frozen reference policy and a visual reference image.
• Structured reasoning: We incorporate reflection traces into this stage, requiring the model to "think through" the editing process before generating pixels, which significantly improves its ability to handle geometric, physical, and temporal transformations.

Table 1 reports UniDFlow’s performance on a broad suite of multimodal understanding benchmarks, covering both perception- and reasoning-oriented evaluations. For UniDFlow (4B), we observe strong results across all benchmarks, demonstrating competitive performance on both general multimodal QA and more reasoning-intensive visual tasks. We further show that UniDFlow consistently outperforms comparable unified baselines. For example, UniDFlow achieves gains over BAGEL of +6.9% on MME-P and +7.0% on MME-S, indicating improved perceptual and reasoning consistency. Table 2 shows evaluation on text-to-image generation benchmarks, where UniDFlow (4B) outperforms same-scale unified competitors and even larger generative models.