UMM-Reflection: Unified Multimodal Models Self-Correct Images Without External Verifiers, 3x Repair Success

Core Results

Repair success rate tripled: On initially erroneous images, reflection SFT fixed only 20.6%; after RL, 64.9%.

GenEval +12 points (0.72→0.84): Spatial position jumped from 0.47 to 0.89 (+42), color binding +14, counting +10. All three models' initial image accuracy was nearly identical (70–73%), so gains come entirely from reflection rounds.

Transfer to unseen benchmarks: WISE 0.63→0.74, T2I-CompBench++ 0.50→0.55, OneIG 0.79→0.83.

Direct generator RL baseline: Same RL budget on plain text-to-image only reached GenEval 0.76 and WISE below base; UMM-Reflection beats it by 19.8 points on WISE.

No external critic at inference: The same model performs inspection, diagnosis, editing, and stopping decisions.

Background: Unified Models Can See and Draw, But Can't Self-Correct

As text prompts grow complex, single-pass generation often yields errors—wrong spatial relations, counts, or color bindings. Traditional text-to-image models cannot detect or fix these. Unified multimodal models (UMMs) combine visual understanding and image generation in one network, theoretically enabling a self-correction loop akin to LLMs: check, diagnose, revise.

Prior work follows two paths: (1) RL on single-pass generation (e.g., text planning first, or using model's own understanding as reward) but never revisiting the generated image; (2) multi-step reasoning/editing on intermediate images, mostly via supervised imitation (SFT) on synthetic trajectories. Some use RL for multi-round regeneration, yet rounds are independent and don't inherit context.

Key Observation: SFT Learns the Format, Not Reliability

Researchers first built 29,529 multi-turn reflection trajectories and ran reflection SFT on BAGEL. Results show SFT learns more than format: 95% of trajectories follow the reflection protocol, and sampling 16 trajectories per erroneous training image yields at least one correct repair path for 78% of images—effective edits exist within the model's capability.

But reliability is low. On a single trajectory, SFT fixes only 20.6% of initial errors; three reflection rounds add merely ~2 points on GenEval and plateau after round one. As Figure 1 (top-left) shows, SFT edits often loop: move tie, move dog, move tie back, still wrong. This reveals a missing outcome-driven signal: whether a reflection helps is only known after the triggered image is generated, possibly several rounds later.

Method: Trajectory-Level RL in a Single Model

The authors term the UMM's own multi-turn "inspect–diagnose–edit" behavior native reflection . Because diagnosis text and the edited image come from the same parameters (text head writes diagnosis, flow head re-renders per diagnosis), a single outcome reward can update both along the same trajectory. Training has two stages.

Stage 1: Reflection SFT Cold Start

Each round, the model observes the original request, full image-text history, and current image, then outputs structured reflection (THINKING / ACTION / EDIT) and decides continue or stop. Training data is synthesized by external models to avoid distilling BAGEL's own failure modes: GPT-5.5 acts as critic, Qwen-Image and Qwen-Image-Edit generate and edit. Data covers one-shot success, natural repair, and multi-step planning trajectories; prompts do not overlap any evaluation benchmark.

Stage 2: Full-Trajectory Reinforcement Learning

Two core designs:

Shared initial image sampling: For each prompt, generate one initial image, detach it from the computation graph, then sample 16 complete reflection trajectories from that identical image. This makes within-group advantage comparison reflect reflection strategy quality, not initial image luck.

One trajectory, one advantage: Each trajectory receives a single outcome-driven advantage value that simultaneously updates all reflection tokens and flow generation steps. This avoids exponential rollouts for per-step credit assignment and eliminates the need for a separate value model.

Reward: A frozen verifier gives each image a graded 0–1 score, so partial repairs also yield non-zero progress signals. Trajectory reward combines final image quality, per-round improvement magnitude, and a penalty for "breaking what was correct." The verifier is used only during training; at inference the model relies entirely on itself for reflection and stopping.

Experiments: Gains Concentrate in Reflection Rounds and Transfer to Unseen Benchmarks

All comparisons use the same BAGEL backbone, prompts, resolution, and scorer. RL runs on 3,000 training prompts from GenEval's six task categories; evaluation covers all 553 official GenEval prompts plus unseen WISE, OneIG-Bench, and T2I-CompBench++.

On GenEval, UMM-Reflection beats reflection SFT by 12 points. The most composition-dependent metric, spatial position, leaps from 0.47 to 0.89; color binding +14, counting +10. Base, SFT, and RL initial image accuracy are nearly identical (70–73%), and the initial generator receives no RL gradients, confirming gains come from reflection rounds, not a better first image. Per-prompt paired statistics: RL wins on 109 prompts, loses on 38 (McNemar p < 1e−8); comparing only initial images shows no significant difference.

On three benchmarks never seen during training, the same checkpoint improves over reflection SFT: WISE +11.0, T2I-CompBench++ +4.6, OneIG +3.5. Notably, on OneIG, plain SFT even slightly underperforms Base.

From a test-time scaling view, conditional repair rate (fraction of initially wrong images eventually fixed) rises from 20.6% to 64.9%. Cost: inference may generate up to 4 images, so authors also compare under equal 4-image budget.

Qualitative Examples

Figure 4 shows multi-turn reflection trajectories on four benchmarks. Each row: initial image → rounds 1, 2, 3 edits. The model progressively corrects spatial relations, colors, materials, counts, and world-knowledge errors.

Ablation: Training Only Generator or Only One Head Is Insufficient

What if the same RL budget directly optimizes the generator? Direct text-to-image RL (1,000 steps) lifts single-pass GenEval from 71 to 76 but barely transfers: WISE 54 (Base 55) vs. UMM-Reflection's 84 and 74. Full ablation in Table 3.

Key contrasts: Reflection SFT alone adds only 2 points; training only flow head or only text head reaches 73 and 78 respectively; joint full-trajectory optimization hits 84. Under equal 4-image budget, picking best of 4 samples from direct RL model reaches only 80; forcing SFT to edit every round yields only 72—proving gains are not from "just generating more images" but from learning which edits actually work.

What RL Actually Changes: Action, Not Perception

To understand the source of gains, researchers trained linear probes on the model's internal representations to read out "is this image correct?" across all 553 trajectories for Base, SFT, and RL checkpoints. The probe AUC on the understanding branch stays 0.80–0.82 across all three—RL barely changes the model's internal correctness judgment.

What changes is where images move in that internal space. For initially wrong images, after three reflection rounds the fraction landing in the "correct region" rises from 34% to 62% under RL, but only from 36% to 42% under SFT. In other words, the model already "knows" if its image is right; RL teaches it to act on that knowledge: among the many edits SFT can already produce, RL selects paths that truly push erroneous images toward the correct region.

Conclusion

This work shows that the "SFT cold-start + full-trajectory RL" paradigm, widely validated in LLM post-training, transfers to native multi-turn visual self-correction in unified multimodal models: SFT teaches reflection format and meaningful edits; RL makes those edits reliable. At inference, no external critic or verifier is needed—the model completes inspection, diagnosis, editing, and stopping within a single context.

Caveats: Multi-turn reflection adds inference overhead (up to 4 images); RL reward comes from a GenEval-style compositional alignment verifier, and training prompts cover only GenEval's six task categories—broader reward dimensions remain for future work. The team has open-sourced all training/evaluation code, RL and SFT model weights, and SFT data to provide a reproducible starting point for reflection and RL research on unified models.

Paper: https://arxiv.org/abs/2609.35767 (Learning Native Reflection in Unified Models with Interleaved Reinforcement Learning) Project page: https://waltstephen.github.io/UMM-Reflection/ Code: https://github.com/waltstephen/UMM-Reflection Models & data: https://huggingface.co/collections/YijiaFan/umm-reflection-6ab95afe909092518d70a158