Learning Native Reflection in Unified Models with Interleaved Reinforcement Learning

AI summary

Overview

Research area: Multimodal generative AI — specifically unified multimodal models that both understand and generate images, trained with reinforcement learning for self-correction.

Technical level: Advanced. The paper assumes familiarity with reinforcement learning (GRPO-style group-relative advantages, PPO-style clipped surrogates, KL penalties), flow-matching / diffusion image generators, and unified understanding-and-generation architectures.

Scope: The paper introduces UMM-Reflection, a two-stage method (supervised fine-tuning followed by whole-trajectory reinforcement learning) that trains a single unified model to inspect its own generated image, write a diagnosis, render a correction, and repeat — evaluated on four text-to-image benchmarks.

What This Paper Is About

A unified multimodal model can look at its own output and fix it, but only if it learns both what to say about a flawed image and how to render the fix. The paper's problem is that imitation learning (SFT) on reflection trajectories teaches the model to produce plausible corrections yet does not make it reliably choose the corrections that actually work. UMM-Reflection solves this by applying reinforcement learning to complete multi-round inspect–diagnose–revise trajectories, using one outcome-driven advantage per trajectory that updates both the text (reflection) head and the flow (rendering) head of the same model.

Key Contributions

1. Reinforcement learning for multi-round reflection inside one unified model. A single whole-trajectory advantage, computed over sibling trajectories that share one detached initial image, jointly optimizes the textual reflections and the flow-based revisions of the same model, avoiding per-round branching (which would need K^N rollouts, i.e. 16³ = 4,096 for three rounds).
2. A separation of "producing" from "choosing" repairs. The paper shows imitation already teaches meaningful revisions — 16 SFT rollouts contain a correct repair for 78% of failing training images — but a single SFT trajectory repairs only 20.59% of initially incorrect images; after RL this conditional repair rate rises to 64.94%.
3. Empirical gains across four benchmarks while training on one. UMM-Reflection improves GenEval by 12.05 points over SFT, with gains transferring to WISE (+10.97), OneIG-Bench (+3.48), and T2I-CompBench++ (+4.63), none of which is used in training.
4. Ablations showing neither single-head RL nor Best-of-4 selection matches the full method. Freezing either head, removing the multi-improvement reward term, or selecting among four images from a stronger single-pass renderer all fall short.

Main Findings

• In-domain accuracy. On GenEval under the paper's protocol, UMM-Reflection reaches 0.84, against 0.71 for BAGEL-Base and 0.72 for reflection SFT. The gain concentrates in families that require fixing a composition: position rises from 0.47 to 0.89 (+42.00 points over SFT), color binding by +14.00, and counting by +10.00.
• Paired prompt comparison. RL beats SFT on 109 GenEval prompts and loses on 38 (p < 10⁻⁸, McNemar). On the initial images alone the split is 48–32 (p = 0.09), which the authors read as no significant difference — the gain is made in the reflection rounds.
• Transfer to unseen benchmarks. Table 2 reports WISE 0.55 (Base) → 0.63 (SFT) → 0.74 (RL); OneIG 0.80 → 0.79 → 0.83; CompBench 0.49 → 0.50 → 0.55. SFT alone falls slightly below Base on OneIG.
• Repair versus damage. The conditional repair rate rises from 20.59% (SFT) to 64.94% (RL, 1,000 updates). Damage on initially correct images is 8.8% for UMM-Reflection versus 6.5% for reflection SFT.
• The first image does not improve. Base, SFT, and RL start from nearly identical single-round accuracy (70–73). RL assigns no loss to the detached initial render; the gains come from reflection rounds.
• Direct RL on the renderer does not transfer. Direct T2I-RL (Flow-GRPO on Base, 1,000 updates) lifts single-shot GenEval from 71 to 76 but WISE 54 versus 55 for Base, whereas UMM-Reflection reaches 84 and 74.
• Both heads must be trained. With the text head frozen, flow-only RL reaches 73 (repair 22.8%). With the flow head frozen, text-only RL reaches 78 (repair 49.4%). Joint RL reaches 84 (repair 64.9%), six points above the best single head.
• Most of the gain arrives within 500 updates. Along the same run, GenEval moves from 72 (SFT) to 75, 79, and 82 after 100, 200, and 500 updates, reaching 84 at 1,000. The repair rate goes from 21% to 61% at 500 updates and 65% at 1,000, with damage between 8% and 10%. Initial-image accuracy stays at 71–73 throughout.
• The gain is not Best-of-N selection or extra edits. At the same four-image budget, Best-of-4 selection from the stronger T2I-RL renderer reaches 80 with either selector (Base UND or UMM-Reflection UND), versus 84 for UMM-Reflection (paired McNemar p ≤ 0.04). Forcing SFT to edit in all three rounds leaves it at 72, the same as unforced SFT. On the 62 prompts where all four independent Base draws fail, RL reflection recovers 60% (SFT: 12%).
• The multi-improvement reward term matters. Removing it (β = 0) preserves the repair rate but degrades the initial image from 73 to 63, also yielding 79.
• Test-time scaling. SFT's three rounds add roughly +2 points on GenEval and flatten after round 1; after RL, round 1 alone adds +9 points and the model continues to gain through round 3. The largest per-family gains are position (+34) and color_attr (+17).
• What changes inside the model. In a representation analysis, a held-out linear probe on the understanding stream reaches AUC 0.80–0.82 under Base, SFT, and RL checkpoints — the correctness readout is nearly unchanged. What changes is where images go: among initially failing images, the share inside the dense pass region rises from 34% to 62% over three rounds under RL, against 36% to 42% under SFT. The authors conclude RL selects existing repair paths rather than creating a new capability.
• Training dynamics. Protocol compliance converges within the first fifty updates, with invalid trajectories dropping below 1% and staying there. Successful repairs per edit rise from 11% to 38%; damage on initially correct images falls from 20% to 10%; terminal exactness under the training verifier reaches 81%. No reward collapse or protocol regression occurs within 1,000 updates.

Methodology in Plain English

The system is built on BAGEL, a unified model whose understanding and generation experts share attention. The pipeline has three parts.

Defining the reflection protocol. Given a request, the model first produces an image, then enters a loop of at most three repair rounds. Each round it observes the request, the text–image history, and the current image, and emits a tagged reflection whose main fields are [THINKING], [ACTION], and [EDIT] (six fields in total, including [CURRENT_ROUND], [SOURCE_IMAGE], and [SCORE]). An "edit" action carries a natural-language correction that the same model uses to render the next image; a "done" action returns the current image.

Supervised cold start. The target model cannot yet produce these trajectories, so they are generated externally and distilled in. GPT-5.5 acts as critic, turning each request into verifiable constraints and writing structured reflections; Qwen-Image renders the initial image and Qwen-Image-Edit executes each edit. BAGEL takes no part in data generation, so its own failure modes are not distilled back in. The 29,529 accepted trajectories fall into three types: one-shot (9,000, where the initial image already satisfies all constraints), natural-repair (8,645, where a genuinely failing image is fixed in one or two rounds), and planned-progression (11,884, where a complex request is fulfilled over two or three ordered milestones). Prompts come from Puffin-4M, Poster100K, OmniEdit, AnyEdit, and GEdit-Bench, with no overlap with any evaluation benchmark. SFT uses a combined autoregressive loss on reflection text and a flow-matching loss on edited images, for one epoch from the base BAGEL checkpoint (167,363 effective records).

Whole-trajectory RL. For each request, one initial image is sampled and detached from the computation graph; RL optimizes only the reflection-and-editing rounds that follow. From those identical root pixels, K = 16 complete trajectories are sampled. A frozen verifier assigns each image a graded score q in [0, 1] under official thresholds, so a partial repair yields a nonzero change. The trajectory reward combines a terminal score q_T, a bonus for each round that improves the image, a multi-improvement term that credits trajectories improving across more than one edit, a damage penalty for regressions, and a penalty for a premature "done". The weights used are α = β = 0.3 and λ = p = 0.5.

Advantages are computed group-relatively within each shared-root group: A_i = clip((R(τ_i) − μ_R) / max(σ_R, 0.1), −1, 1). The key design choice is that both channels — the text policy and the flow renderer — share the same trajectory-level advantage rather than being normalized separately, so a reflection that leads to a better image raises the advantage for both the diagnostic tokens and the rendering transitions that followed. Each channel uses a clipped surrogate with its own clip range (0.2 text, 0.1 flow) and a KL penalty of 10⁻⁴ against the frozen SFT reference. Flow transitions use the Flow-GRPO SDE sampler with two contiguous stochastic transitions per active repair. The verifier and reference policy are used only during training; at inference only the unified model runs, with no verifier.

Training uses a 3,000-prompt pool over six GenEval families (position 1,305; counting 702; color attribution

Authors’ abstract

Unified multimodal models can both look at and render images, so in principle they can repair their own generations: diagnose what an image gets wrong, revise it, observe the result, and diagnose again. Whether a revision helps is known only after it is rendered, so the reflection text and the image generation must be learned jointly, over the whole loop. Supervised fine-tuning (SFT) on reflection trajectories gives a cold start but does not find the high-success repair paths, and naive RL that optimizes only the renderer or only one head leaves most of the gain untapped. We introduce UMM-Reflection, which applies reinforcement learning (RL) to complete reflection trajectories inside one unified model: sibling trajectories share one initial image, so the group-relative advantage compares reflection strategies, and one trajectory-level advantage updates both the reflection tokens and the flow-based revisions, avoiding the combinatorial blow-up of per-round credit assignment. Unlike single-round editing or pipelines with an external critic, credit flows across rounds and to both roles of the same model, and no verifier is needed at inference. On BAGEL, UMM-Reflection improves GenEval by 12.05 points over SFT, and the gains transfer to WISE (+10.97), OneIG-Bench (+3.48), and T2I-CompBench++ (+4.63), none of which is used in training.