Learning Native Reflection in Unified Models with Interleaved Reinforcement Learning

Learning Native Reflection in Unified Models with Interleaved Reinforcement Learning. By Yijia Fan and seven others, from Nanyang Technological University, Shanghai Jiao Tong University, and one other institution. Published on the twenty-eighth of September. Its code repository is worked out: about one hundred and forty source files, with a readme and a licence.

A unified multimodal model can both look at images and render them. In principle, then, it can repair its own generations: diagnose what an image gets wrong, revise it, look at the result, and diagnose again. The catch is that whether a revision helped is only known after it has been rendered. The reflection text and the image generation therefore have to be learned together, across the whole loop. Existing recipes fall short of that. Supervised fine-tuning on reflection trajectories gives a cold start but does not find the repair paths that succeed most often. Naive reinforcement learning that optimizes only the renderer, or only one head, leaves most of the gain on the table.

UMM-Reflection applies reinforcement learning to complete reflection trajectories, inside one unified model. The whole loop of critique and revision is the unit being optimized, rather than any single step within it. That matches the problem as stated: since the value of a revision is only visible once it is rendered, the thing to reward is the trajectory, diagnosis and revision together, as it plays out across rounds.

Two choices make this workable. Sibling trajectories share one initial image, so the group-relative advantage compares reflection strategies against each other on the same starting point. And a single trajectory-level advantage updates both the reflection tokens and the flow-based revisions. That avoids the combinatorial blow-up that assigning credit round by round would cause.

The contrast with other setups is in where credit goes. Single-round editing has no loop to credit, and pipelines with an external critic split the roles between models. Here credit flows across rounds, and to both roles of the same model. No verifier is needed at inference.

On BAGEL, UMM-Reflection improves GenEval by twelve point zero five points over supervised fine-tuning. The gains carry over to benchmarks that were never used in training: ten point nine seven points on WISE, three point four eight on OneIG-Bench, and four point six three on T2I-CompBench++.

The model that draws the picture and the model that critiques it are one and the same, and now they are trained as one as well. One advantage per trajectory, shared across every round of critique and revision, is what teaches the loop to run without an outside judge.