PaperlayerView exact revision on arXiv

Source-linked AI summary

Learning Native Reflection in Unified Models with Interleaved Reinforcement Learning

Yijia Fan, Ziqi Huang, Zhongang Cai, Yan Li, Zimo Wen, Wanqi Yin, Haiwen Diao, Ziwei Liu

arXiv:2609.35767v1cs.CVcs.AI

TL;DR

Unified models can inspect and revise images, but learning effective reflection requires crediting revisions for their rendered outcomes across the loop. UMM-Reflection uses trajectory-level RL to update reflection and rendering together, improving GenEval over SFT and transferring gains to other benchmarks.

  • Problem

    Unified models can inspect and generate images, but prior RL addresses single renders and imitation of multi-round reflection does not ensure effective correction.

    Source: §1 · p. 1

  • Method

    UMM-Reflection applies whole-trajectory RL to sibling trajectories sharing an initial image, using one outcome-driven advantage to update reflection tokens and flow-based revisions together.

    Source: §1 · p. 3

  • Results

    GenEval improves by 12.05 points over SFT, with gains also transferring to WISE, OneIG-Bench, and T2I-CompBench++.

    Source: §1 · p. 2

  • Takeaways & Limitations

    RL selects repair paths the model can already produce, making its revisions more reliable without changing its correctness readout substantially.

    Source: §7 · p. 10

Abstract

from arXiv · show

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.

1 Introduction

Unified multimodal models can inspect and revise their own images, but effective multi-round reflection requires learning which revisions work from their rendered outcomes. UMM-Reflection applies trajectory-level reinforcement learning to jointly train reflection and image generation, improving repair and transfer across benchmarks.

1 Introduction

  • Unified models can look at and generate images in one network, enabling a loop of inspection, diagnosis, and revision beyond one-shot generation.A single render may contain flaws that a one-shot text-to-image model cannot notice or fix.Source: §1 · p. 1
  • The missing capability is reinforcement learning over a unified model’s multi-round reflection, crediting each reflection for the visual improvement it produces rather than for matching a demonstration.Optimizing only the renderer or one head leaves most of the gain untapped.Source: §1 · p. 1
  • Because a reflection’s value is known only after its image is rendered, and sometimes after further rounds, outcome credit must reach both diagnosis and rendering across the trajectory.The text head writes the diagnosis and the flow head renders the next image from the same model parameters.Source: §1 · p. 2
  • Sibling trajectories share one initial image, so group-relative advantage compares reflection strategies; one trajectory-level advantage jointly updates reflection tokens and flow-based revisions.This avoids per-round branching and its rollout cost.Source: §1 · pp. 2–3
  • A single SFT trajectory repairs 20.59% of initially incorrect images, while RL raises the conditional repair rate to 64.94%.SFT trajectories already contain correct repairs, but they do not reliably select them.Source: §1 · p. 2
  • GenEval improves by 12.05 points over SFT and WISE by 10.97; gains also transfer to OneIG-Bench and T2I-CompBench++, neither used in training.UMM-Reflection improves over reflection SFT on four benchmarks while training on one.Source: §1 · pp. 2–3

2 Related Work

Prior work studies self-correction, RL for visual generation, and reflection in unified models, but generally separates the roles or stops at one render. UMM-Reflection applies outcome-driven RL to complete inspect-and-revise trajectories within one unified policy.

Self-correction loops, in text and in pixels

  • Language-model self-correction can lower accuracy without external feedback, while image-generation loops often rely on a separate online critic and renderer [14] [17] [34] [37] [46].UMM-Reflection trains both roles as one policy under one outcome reward and drops the verifier at inference.Source: §2 · p. 3

Reinforcement learning for visual generators

  • Visual-generator RL methods optimize a single prompt-to-image pass that never observes its own render [1] [9] [30] [39] [28] [20] [40].UMM-Reflection places Flow-GRPO rendering transitions inside a multi-round trajectory whose advantage credits reflections and their resulting renders.Source: §2 · p. 3

Reasoning and reflection in unified generators

  • Unified-generator work either improves a single render [18] [44] [22] or imitates trajectories that inspect intermediate images [5] [7] [12] [16] [25].UMM-Reflection applies outcome-driven RL to complete inspect-and-revise trajectories in one unified policy.Source: §2 · p. 3

3 Methodology

The method trains a unified model to inspect, diagnose, and revise its own images, first teaching the interleaved protocol with external trajectories and then optimizing complete reflection loops with reinforcement learning.

3.1 Reflection protocol

  • At each of up to three repair rounds, the model observes the request and image history, then emits a reflection and either an edit or a done action.Reflections use [THINKING], [ACTION], and [EDIT] fields; verifier scores do not enter the policy observation.Source: §3.1 · pp. 3–4

3.2 Trajectory data construction

  • External models produce 29,529 accepted trajectories across one-shot, natural-repair, and planned-progression types, with prompts drawn from five sources and no evaluation-benchmark overlap.The target model is BAGEL; GPT-5.5 critiques and writes edits while Qwen-Image and Qwen-Image-Edit generate and revise images.Source: §3.2 · p. 4

3.3 Supervised initialization

  • SFT initializes BAGEL on the constructed trajectories using autoregressive loss for reflection text and flow-matching loss for each edited image.This stage teaches the interleaved protocol and provides the checkpoint from which RL starts.Source: §3.3 · p. 4

3.4 Whole-trajectory reinforcement learning

  • RL samples K = 16 complete reflection trajectories from identical detached initial images, so sibling rollouts compare subsequent repair trajectories rather than different initial draws.The initial text-to-image generation receives no policy-gradient signal, and siblings are not pruned or selected best-of-K at deployment.Source: §3.4 · pp. 4–5
  • The trajectory reward combines final-image quality and per-round improvements, penalizes regressions and premature stopping, and rewards improvement across multiple edits.The specified coefficients are α=β=0.3 and λ=p=0.5.Source: §3.4 · p. 5
  • One group-relative trajectory advantage updates both reflection tokens and flow transitions, avoiding the KN rollouts required for per-round group comparisons; the verifier and reference policy are training-only.The same advantage credits diagnostic text and the image-rendering transitions that followed it.Source: §3.4 · p. 5

4 Experiments

Experiments compare UMM-Reflection with BAGEL Base and reflection SFT, evaluate transfer to three benchmarks unseen in RL training, and test how repair performance changes across rounds.

4.1 Setup

  • RL starts from the reflection-SFT checkpoint and uses one 1,000-update checkpoint across GenEval, WISE, OneIG, and CompBench, with an external controller enforcing interleaving.Base and SFT use the same BAGEL backbone; evaluation uses one image per prompt and at most three repairs.Source: §4.1 · pp. 5–6

4.2 Main results

  • GenEval reaches 0.84 versus 0.72 for reflection SFT (+12 points), and UMM-Reflection beats SFT on unseen WISE (+10.97), CompBench (+4.63), and OneIG (+3.48).Table 1 reports GenEval scores across six compositional requirements; Table 2 reports transfer on benchmarks unseen in RL training.Source: §4.2 · p. 6, Table 1 · p. 6, Table 2 · p. 6

4.3 Multi-round test-time scaling

  • After RL, round 1 adds +9 GenEval points and gains continue through round 3; conditional repair rises from 20.59% (SFT) to 64.94% (RL 1000).The largest reported family gains are position (+34) and color_attr (+17).Source: §4.3 · p. 7

5 Ablations

The ablations show that reflection-round RL improves repair beyond SFT, with the strongest results requiring joint training of the text and flow heads. The gains are not explained by extra edits or selecting among multiple generated images.

5 Ablations

  • Trajectory RL on top of SFT adds +11 GenEval points and raises the repair rate from 21% to 65%.Direct generator RL before both stages reaches 81 after 500 updates, while fine-tuning on final images from the same trajectories leaves single-pass accuracy at 75, the same as Base.Source: §5 · p. 8, Table 3 · p. 8
  • GenEval rises from 72 (SFT) to 75, 79, and 82 after 100, 200, and 500 RL updates, reaching 84 at 1,000 updates.Repair increases from 21% to 61% at 500 updates and 65% at 1,000, while damage stays between 8% and 10%.Source: §5 · p. 8, Table 3 · p. 8
  • Jointly training both heads reaches 84 GenEval accuracy and a 64.9% repair rate, six points above the best single-head setting.Training only the flow head yields 73 and 22.8% repair; training only the text head yields 78 and 49.4%.Source: §5 · p. 8, Table 3 · p. 8
  • At the same four-image budget, Best-of-4 reaches 80 while UMM-Reflection reaches 84; on prompts where all four Base draws fail, RL reflection recovers 60% versus 12% for SFT.The Best-of-4 comparison has paired McNemar p ≤0.04.Source: §5 · pp. 8–9, Table 3 · p. 8
  • Removing the multi-improvement reward term preserves the repair rate but lowers initial-image accuracy from 73 to 63, yielding 79 overall.Figure 4 plots GenEval macro accuracy against reflection rounds; p20 reports that initial-image accuracy stays at 71–73 as gains accrue through reflection.Source: Figure 4 · p. 7, §5 · pp. 8–9, Table 3 · p. 8

6 Study: What Changes Inside the Model

Training first establishes protocol compliance, then improves repair quality across reflection rounds. The backbone’s correctness readout changes little under RL, while its initially failing images move more often into the readout’s pass region.

6.1 Training dynamics: the interface locks in first, then repair improves

  • On the full 553-prompt GenEval test set, SFT’s three reflection rounds add +2 macro-accuracy points, while RL adds +11, with gains concentrated in repair rather than initial generation.Protocol compliance converges within 50 updates; successful repairs per edit rise from 11% to 38%, and damage to initially correct images falls from 20% to 10%.Source: §6.1 · p. 9

6.2 What RL changes in the model

  • The understanding-stream correctness probe stays near AUC 0.80–0.82 across Base, SFT, and RL, while the pass-region share among initially failing images rises from 34% to 62% under RL and from 36% to 42% under SFT.RL moves images further at a matched number of edits, and more of its edits succeed.Source: §6.2 · p. 9

7 Conclusion

The paper concludes that RL makes reflection in a unified model an effective repair mechanism. Its reported GenEval gain and transfer to other benchmarks come from reflection rounds, as RL selects repairs that reach the correct region without substantially changing the model’s correctness readout.

7 Conclusion

  • RL turns reflection from an SFT cold start into an effective repair mechanism, with the +12-point GenEval gain and transfer to WISE, OneIG, and CompBench coming from reflection rounds.The backbone’s correctness readout remains nearly unchanged; RL selects among repairs the model can already produce those that land in the correct region.Source: §7 · pp. 9–10

A Implementation and Evaluation Details

The appendix documents the training and evaluation setup, comparison controls, reflection format, and graded reward. Evaluations use a shared RL checkpoint and controlled prompt protocols; reward grading gives partial progress signals while preserving the official pass/fail decisions.

A Implementation and Evaluation Details

  • RL trains on 3,000 prompts from six GenEval families with 16 siblings per root, and the final checkpoint follows 1,000 updates; evaluation uses 512 × 512 images, 50 denoising steps, and at most three repairs.The same checkpoint is evaluated across benchmarks, with one returned image per prompt and a separate 270-prompt UID-disjoint held-out set.Source: A Implementation and Evaluation Details · p. 15
  • The benchmark set comprises GenEval, WISE, OneIG-Bench, and T2I-CompBench++, with distinct metrics and controlled instruction-voice evaluations rather than native leaderboard protocols.WISE curves use a common GPT-4o re-evaluation, while the main table reports the original WISE evaluation.Source: A Implementation and Evaluation Details · pp. 15–16

A.1 Direct-generation RL control

  • Direct T2I-RL is a 1,000-update Flow-GRPO control initialized from Base, without reflection SFT or a controller objective, and it performs no repair or candidate selection at inference.It shares the six-family training pool, K = 16, 20 training denoising steps, and two selected flow transitions with UMM-Reflection.Source: §A.1 · p. 17, Figure 7 · p. 17

A.2 Understanding-based selection protocol

  • Best-of-4 compares choices among four direct T2I-RL images using either Base or UMM-Reflection’s understanding path, with no task-family, detector, candidate-score, or verdict information given before selection.Each selector has one invalid-format response among 553 prompts; the preregistered fallback selects the first presented image.Source: §A.2 · pp. 17–18

B Reflection Format and SFT Details

  • The reflection format has six tagged fields, including THINKING, ACTION, and EDIT; the model-generated SCORE is a self-assessment, and verifier scores or labels are not policy observations.SFT decomposes trajectories into controller, transition, and verifier rows, supervising reflection text, individual edits, and the next reflection, respectively.Source: B Reflection Format and SFT Details · p. 18

C Graded Reward

  • Under standard binary GenEval scoring, five of six families yield ∆t = 0 when edits improve some but not all constraints; the method instead grades failure regions using frozen-verifier outputs at official thresholds.The official thresholds are 0.3 for object detection and 0.9 for counting; no sub-threshold confidence is read. Across 200 image–family pairs, no fail-closed violations occurred.Source: C Graded Reward · pp. 18–19

D Training Dynamics

The training-dynamics figure tracks reward, repair success, and exactness across RL updates, while the external-benchmark figure compares scores across reflection rounds and model variants.

D Training Dynamics

  • Figure 8 tracks mean trajectory reward, successful repairs per edit from an incorrect state, and trajectories reaching terminal exactness across the 1,000-update RL run.Faint lines show individual updates; solid lines are trailing 20-update trends.Source: Figure 8 · p. 19
  • Figure 8 uses pooled numerators and denominators over each 20-update window for its rate trends.Source: Figure 8 · p. 19

E Test-Time Scaling on External Benchmarks

  • Figure 9 plots scores from 0–100 against reflection rounds for Base, reflection SFT, UMM-Reflection-500, and UMM-Reflection-1000 on four benchmarks.Source: Figure 9 · p. 19

F Distributional Analysis

The distributional analysis finds that RL makes larger, more successful image changes than SFT at matched edit counts, while leaving the model’s visual pathway and cross-round attention largely unchanged.

F Distributional Analysis

  • At three edits, RL’s DINO distance is 0.28 versus SFT’s 0.11, and its exact-match gain from R0 to R3 is +16.3 points versus −3.3.The paired differences are 0.17 in DINO distance (95% CI [0.13, 0.21]) and +19.6 points in exact-match gain (95% CI [11.1, 28.1]).Source: F Distributional Analysis · p. 20, Table 5 · p. 20
  • RL raises the share of initially failing images inside the dense pass region from 34% at R0 to 62% at R3, while SFT rises from 36% to 42%.The understanding-stream verifier probe’s held-out AUC is 0.815 under RL, versus 0.807 under SFT and 0.804 under Base.Source: F Distributional Analysis · p. 20

G Visual Pathway Stability

  • Across ten replayed trajectories, SFT and RL attention maps correlate at a median of 0.98 over VAE keys and above 0.99 over ViT keys.Figure 10 shows paired SFT/RL attention to 32 × 32 VAE keys, averaged over 28 layers.Source: G Visual Pathway Stability · p. 20, Figure 10 · p. 21

H Cross-Round Attention to Earlier Images

  • At the third reflection round, each earlier image receives at least 19% of image attention, and SFT and RL divide attention across the visual history similarly.VAE-key attention is split almost equally across the three images, at 29–37%.Source: H Cross-Round Attention to Earlier… · p. 22, Table 6 · p. 22

I Correct Repairs Already in the SFT Rollout Distribution

RL concentrates its policy on successful repairs already present in the SFT rollout distribution, and fixed-image tests identify the RL-written instruction as an effective repair channel.

I Correct Repairs Already in the SFT Rollout Distribution

J Swapping the Editing Instruction on Fixed Images

  • With image, renderer, and sampling noise fixed, the RL instruction repairs 48.4% of failing states, versus 20.5% for the original request and 27.6% for a rule-written instruction.The paired RL-minus-request difference is +27.9 points (95% state bootstrap interval [20.8, 34.7]).Source: J Swapping the Editing Instruction on… · p. 22

K Paired Final-Model Results

The section reports final-model benchmark comparisons and supporting category, qualitative, stopping, and family-level analyses, with the inference prompt specifying the reflection protocol.

K Paired Final-Model Results

  • The final full-weight RL model and retained SFT are compared on the same 553 prompts, with paired GenEval verdicts and stage-wise results reported in separate tables.Table 8 distinguishes initial and final accuracy, conditional repair and damage, and protocol validity.Source: K Paired Final-Model Results · p. 23, Table 8 · p. 23, Table 9 · p. 23

L Category-Level Results

M Detailed External Benchmark Comparisons

  • The detailed tables cover WISE, T2I-CompBench++, and OneIG with benchmark-specific local and reference protocols.The local CompBench evaluation uses one image for each of 2,400 prompts, unlike the published evaluation’s ten images per prompt; the local WISE table uses original-WISE legacy scores, not WISE_Verified.Source: M Detailed External Benchmark… · p. 24, Table 11 · p. 25, Table 12 · p. 25

N External Critic Baseline

  • On GenEval, UMM-Reflection scores 84 and repairs 65% versus the external pipeline’s 79 and 28%; it leads on WISE and CompBench and matches the pipeline on OneIG.The external pipeline uses GPT-5.5 to inspect and edit Base images, with no verifier verdict at inference.Source: N External Critic Baseline · p. 27, Table 15 · p. 27

O Additional Qualitative Examples

  • Figures 13–18 show selected successful repairs, while Figure 19 depicts reflections that miscount giraffes or mistake the second-largest economy, leading to repeated incorrect edits.The examples are unedited model outputs, and frames mark benchmark verdicts.Source: O Additional Qualitative Examples · p. 27, Figure 19 · p. 34

P Stopping Behavior under Terminal Rewards

  • Adding a correct-DONE bonus makes stopping cheap and leads to premature stops on 112 incorrect images, which account for 112 of 116 final failures; accuracy falls from 82 to 79.The primary reward instead penalizes DONE on a verifier-incorrect image by 0.5.Source: P Stopping Behavior under Terminal… · p. 35

Q Per-Family Learning on GenEval

  • GenEval families learn at different rates: position accuracy rises from 55 to 89, while counting reward stays flat and counting accuracy remains 67.5 after 1,000 updates.The passage identifies count-aware edits as the next training target.Source: Figure 20 · p. 35, Q Per-Family Learning on GenEval · p. 36

R Inference System Prompt

  • The inference prompt requires tagged THINKING, ACTION, and EDIT fields, with DONE reserved for images that visibly satisfy the complete original request.It directs the model to preserve unrelated content during edits and judge only visible evidence.Source: R Inference System Prompt · p. 36
Loading 2609.35767v1…