Unified Multimodal Models Learn Self-Reflection, Tripling Repair Success

The same model writes the diagnosis and then redraws according to that diagnosis, with no external scorer needed at inference time.

A unified multimodal model can both look at images and draw images. So can it, like a human, check its own drawing after finishing, spot problems, and then go back and fix them?

Researchers from Nanyang Technological University, Shanghai Jiao Tong University, and the University of Tokyo give an affirmative answer. They propose UMM-Reflection, which uses reinforcement learning on the open-source unified model BAGEL to train the complete multi-turn "reflection—revision" trajectory: the same model writes the diagnosis and then redraws according to that diagnosis, with no external scorer needed at inference time.

After reinforcement learning, the model's success rate at repairing initially erroneous images rises from 20.6% to 64.9%, GenEval rises from 0.72 (reflection SFT) to 0.84, with the spatial position item jumping from 0.47 to 0.89; on WISE, which was never seen during training, it improves by 11 points, nearly 20 points higher than directly applying reinforcement learning to the generator. The code, model, and training data are all open-sourced.

• Paper title: Learning Native Reflection in Unified Models with Interleaved Reinforcement Learning
• Paper link: https://arxiv.org/abs/2609.35767
• Code: https://github.com/waltstephen/UMM-Reflection
• Model and data: https://huggingface.co/collections/YijiaFan/umm-reflection-6ab95afe909092518d70a158
• Authors: Yijia Fan*, Ziqi Huang*, Zhongang Cai, Yan Li, Zimo Wen, Wanqi Yin, Haiwen Diao, Ziwei Liu (*co-first authors; Ziwei Liu is the corresponding author)
• Affiliations: Nanyang Technological University, Shanghai Jiao Tong University, University of Tokyo
• Video: UMM-Reflection Chinese walkthrough (2 min 51 s).

Key results at a glance

• Repair success rate more than tripled: Among images with initial generation errors, the reflection SFT model can only fix 20.6%, while after reinforcement learning it reaches 64.9%.
• GenEval improves by 12 points (0.72→0.84): Spatial position 0.47→0.89 (+42), color attribute binding +14, counting +10. The initial image accuracy of the three models is almost the same; all the gains come from the reflection rounds.
• Transfer to benchmarks never seen during training: WISE 0.63→0.74, T2I-CompBench++ 0.50→0.55, OneIG 0.79→0.83.
• The same RL budget used to directly train the generator (direct text-to-image RL): GenEval only reaches 0.76, and WISE is even lower than Base; UMM-Reflection is 19.8 points higher than it on WISE.
• No reliance on any external critic or verifier at inference time: Looking at the image, diagnosing, revising, and stopping are all done by the same model.

1. It can see and draw, but cannot fix its own images

As text instructions become increasingly complex, a single generated image often contains errors: spatial relations reversed, wrong counts, colors bound to the wrong objects. Traditional text-to-image models have no way to notice these errors, let alone fix them. Unified multimodal models (UMMs) put visual understanding and image generation into the same network. In theory, "a model that draws can also look at images," so it may be possible to reproduce the self-correction loop common in large language models: check, diagnose, revise.

Existing work roughly follows two routes. One route applies reinforcement learning to single-pass generation, for example by first writing a textual plan and then producing the image, or by using the model's own understanding of the generated image as a reward, but the model never goes back to revise the image it has already drawn. The other route lets the model continue reasoning and editing on intermediate images, but mainly relies on supervised imitation (SFT) of synthetic trajectories. There is also work that uses reinforcement learning to train multi-round regeneration, but the rounds are independent of each other and do not inherit previous context.

The question the researchers care about is: what exactly does imitation learning teach the model, and what does it lack?

2. Key observation: SFT learns to "revise," but not to revise accurately

The researchers first constructed 29,529 multi-turn reflection trajectories and performed reflection SFT on BAGEL. The results show that what SFT learns goes far beyond format: 95% of trajectories follow the reflection protocol, and when sampling 16 trajectories each for the initially erroneous images in the training set, 78% of the images have at least one path with a correct repair. In other words, effective revisions already exist within the model's capability range.

The problem is reliability. Looking at a single trajectory, the SFT model can only fix 20.6% of the initially erroneous images, and three rounds of reflection bring a total gain of only about 2 points on GenEval, essentially stagnating after the first round. As shown in the upper left of Figure 1, SFT's revisions often go in circles: first move the tie, then move the dog, then move the tie back, and the final result is still wrong.

Figure 1: Comparison of reflection trajectories between SFT and RL under the same prompt and random seed. SFT can already write meaningful edits, but a single trajectory often goes in circles; whole-trajectory RL concentrates the policy on edit paths that reach the "correct image region."

This means that what is missing between "can edit" and "edits correctly" is an outcome-driven signal: whether a reflection is useful can only be known once the image generation it triggers has completed, or even after several more rounds.

3. Method: Reinforcement Learning over Whole Trajectories within a Single Model

The researchers call this multi-round "inspect—diagnose—edit" behavior performed by a unified model itself native reflection. Since the diagnostic text and the edited image come from the same set of parameters (the text head writes the diagnosis, and the flow head re-renders based on the diagnosis), a single outcome reward can update both along the same trajectory. Training is divided into two stages.

Stage 1: Reflection SFT cold start.

At each round, the model observes the original request, the full prior image-text history, and the current image, outputs a structured reflection (THINKING / ACTION / EDIT), and chooses to continue editing or stop. The training data is synthesized by external models: GPT-5.5 acts as the critic, Qwen-Image and Qwen-Image-Edit handle generation and editing, and BAGEL itself does not participate, so as to avoid distilling its own failure modes back into the supervision signal. The data covers three types of trajectories: one-shot success, natural repair, and multi-step planning, and the prompts do not overlap with any evaluation benchmark.

Stage 2: Whole-trajectory reinforcement learning.

There are two core designs:

(1) Sampling from a shared initial image. For each request, only one initial image is generated and detached from the computation graph, and then 16 complete reflection trajectories are sampled starting from this exact same image. In this way, the within-group relative advantage compares the quality of different reflection strategies rather than whose initial image got luckier.

(2) One trajectory, one advantage. Each trajectory computes only a single outcome-driven advantage value, which simultaneously updates all reflection tokens and flow generation steps within it. This avoids the exponential rollout required for per-round credit assignment and also eliminates the need to train an additional value model.

Figure 2: The reinforcement learning pipeline of UMM-Reflection. 16 trajectories share the same initial image, and each trajectory alternates between the model's own reflection and rendering (up to three rounds); a frozen verifier scores each image and each trajectory, and after within-group normalization, a trajectory-level advantage is obtained that updates both the text head and the flow head. The verifier is used only during training.

On the reward side, the frozen verifier gives each image a graded score between 0 and 1, so partial repairs also receive a non-zero progress signal. The trajectory reward considers the final image quality, the improvement magnitude at each round, and a penalty for "breaking what was right." The verifier is used only during training; at inference time, the model relies entirely on itself to perform reflection and make stopping decisions.

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

All comparisons use the same BAGEL backbone, the same prompts, resolution, and scorer. RL is trained only on 3,000 training prompts from the six GenEval task categories, while evaluation covers all 553 official GenEval prompts as well as WISE, OneIG-Bench, and T2I-CompBench++, which never appear in training.

Table 1: Main results under the same evaluation protocol (0–1 score).

On GenEval, UMM-Reflection improves by 12 points over reflection SFT, among which the position item, which most relies on correcting composition, rises from 0.47 to 0.89, and color attribute binding and counting also improve by 14 and 10 points, respectively. Notably, the initial image accuracy of Base, SFT, and RL is almost identical (70–73); the initial generation itself receives no RL gradient, so the gains come from the reflection rounds rather than from a better first image. Per-prompt paired statistics show that RL beats SFT on 109 prompts and falls behind on only 38 (McNemar test p < 1e−8); when comparing only initial images, there is no significant difference between the two.

Table 2: GenEval breakdown results (553 official prompts, 0–1 score).

On the three benchmarks never used in training, the same checkpoint still improves over reflection SFT by WISE +11.0, T2I-CompBench++ +4.6, and OneIG +3.5, respectively. On OneIG in particular, SFT alone is even slightly below Base.

Figure 3: GenEval accuracy as a function of the number of reflection rounds. The three rounds of SFT reflection total about +2 points and flatten after the first round; after RL, the first round alone brings about +9 points and continues to gain through the third round.

From the perspective of test-time scaling, after RL the model's conditional repair rate (the proportion of initially incorrect images that are eventually fixed) rises from 20.6% to 64.9%. The cost is that inference requires generating up to 4 images, so the researchers also ran a control under the same four-image budget.

5. Results Showcase

Below are the multi-round reflection results of UMM-Reflection on four benchmarks. In each row, from left to right, are the initial image and the edited results after three rounds of reflection; the model progressively corrects errors in spatial relations, color, material, count, and world knowledge.

Figure 4: Reflection trajectories on four benchmarks (initial image → rounds 1, 2, 3 of revision).

6. Ablation: Training Only the Generator or Only One Head Is Not Enough

A natural question is: what happens if the same reinforcement learning budget is spent directly on the generator? The researchers used the same 1,000-step RL update to directly optimize text-to-image generation (direct text-to-image RL). It raised the GenEval score of single-pass generation from 71 to 76, but this barely transferred: WISE was only 54 (Base was 55), while UMM-Reflection reached 84 and 74, respectively. The full ablation from the paper is shown in Table 3.

Table 3: GenEval ablation (0–100 points). Initial: accuracy of the first image; Final: accuracy after reflection ends; Fix rate: the proportion of initially incorrect images that are ultimately fixed.

Several comparisons point to the same conclusion: reflection SFT alone brings only 2 points; training only the flow head or only the text head reaches only 73 and 78, respectively, and only jointly optimizing the entire trajectory reaches 84. Under the same budget of four images, selecting the best image from four samples of the direct RL model reaches only 80, and forcing the SFT model to revise every round also reaches only 72, indicating that the gain does not come from "drawing a few more images," but from learning which revisions are truly effective.

7. What Exactly Does RL Change in the Model?

To understand the source of the gains, the researchers fed all 553 trajectories into the Base, SFT, and RL checkpoints, and trained linear probes inside the model to read out "whether this image is correct." The results show that RL barely changes the model's internal judgment of correctness: the probe AUC on the understanding branch remains between 0.80 and 0.82 across all three checkpoints.

What truly changes is where images go in this internal space. For initially incorrect images, the proportion that fall into the "correct region" after three rounds of reflection increases from 34% to 62% under RL, whereas under SFT it only increases from 36% to 42%. In other words, the model already "knows" whether its images are correct; what RL teaches it is to act on this: among the many revisions that SFT can already produce, it selects those paths that can truly push incorrect images into the correct region.

Figure 5: Trajectories of initially incorrect images in the model's own correctness readout space. The horizontal and vertical axes are the linear probe scores for "whether verification is passed" on the understanding branch and the generation branch, respectively; the green region represents reference images that pass verification. The left two panels are RL from round 0 to round 3, and the right two panels are SFT.

8. Conclusion

This work shows that for unified multimodal models, the paradigm of "SFT cold start + full-trajectory RL," which has been widely validated in large language model post-training, can also transfer to native multi-round visual self-correction: SFT teaches the model the format of reflection and meaningful revisions, while RL makes these revisions reliable. Since inference does not depend on any external critic or verifier, the model completes viewing, diagnosis, revision, and stopping by itself within a single context.

It should be noted that multi-round reflection incurs additional inference overhead (generating up to 4 images); the reward in the RL stage comes from a GenEval-style compositional alignment verifier, and the training prompts also cover only the six task categories of GenEval, so rewards in more dimensions remain to be explored. The research team has open-sourced all training and evaluation code, RL and SFT model weights, and SFT data, hoping to provide a reproducible starting point for research on reflection and reinforcement learning in unified models.

© THE END

For reprints, please contact this official account for authorization.

Submissions or news inquiries: [email protected]

Open in WeChat