Learning Native Reflection in Unified Models with Interleaved Reinforcement Learning

What it claims

Whole-trajectory RL on BAGEL raises the rate at which one reflection trajectory fixes an initially wrong image from 20.59% (SFT) to 64.94% (Table 3). The authors find that SFT already teaches useful revisions: 16 SFT rollouts contain a correct repair for 78% of failing training images, and 95% of SFT trajectories follow the protocol. What SFT lacks is a reliable choice among them, and individual trajectories often go in circles (tie, dog, tie). Representation analysis says RL leaves perception and the model's internal correctness readout nearly unchanged. RL selects repair paths the backbone already has; it does not create a new capability.

How it works

1. Protocol. The model renders x0. Each round it emits a tagged reflection ([THINKING], [ACTION], [EDIT]) with action edit or done. An edit carries a natural-language instruction that the same model's flow head renders. The cap is three repair rounds. BAGEL cannot interleave natively, so an external controller feeds each round back as a new turn.
2. SFT data. The 29,529 trajectories are distilled from external models: - GPT-5.5 acted as critic. - Qwen-Image rendered the initial images and Qwen-Image-Edit executed the edits. BAGEL took no part. - The mix is 9,000 one-shot, 8,645 natural-repair and 11,884 planned-progression trajectories. - Prompts come from Puffin-4M, Poster100K, OmniEdit, AnyEdit and GEdit-Bench, with no overlap with evaluation benchmarks. - Training is one epoch from base BAGEL with an AR loss plus a flow-matching loss.
3. RL. - Sampling. Each root image is detached, so the first render gets no gradient. K=16 sibling trajectories are sampled from it. - Reward. A frozen verifier gives graded scores q ∈ [0,1]. The trajectory reward is R = q_T + α·Σ[Δ]+ + β·S_multi − λ·Σ[−Δ]+ − p·1[premature done], with α=β=0.3 and λ=p=0.5. S_multi deliberately favours improvement spread over several edits. - Advantage. It is clipped to [−1,1] with a σ floor of 0.1. The same A_i is applied to both the text tokens and the flow transitions, which use the Flow-GRPO SDE sampler on two contiguous stochastic steps per repair. - Regularisation. Each channel has its own KL penalty against the frozen SFT reference. - Why per-trajectory. Per-round credit would need 16³ = 4,096 rollouts per root, or a learned value model.

flowchart LR
  R[Shared root x0, detached] --> S[16 sibling trajectories]
  S --> F[Reflect]
  F -->|edit| G[Flow head renders]
  G --> F
  F -->|done or 3 rounds| V[Frozen verifier: training only]
  V --> A[One advantage per trajectory]
  A --> T[Text head]
  A --> H[Flow head]

Evidence

• GenEval: +12.05 over SFT. In the authors' protocol, UMM-Reflection reaches about 84.
• Single-round accuracy: Base, SFT and RL all sit at 70–73 on single-round output, so the gain comes from the reflection rounds.
• Direct T2I-RL ablation: with the same number of updates, generator-only RL lifts single-shot GenEval from 71 to 76. It does not transfer: its WISE score is 54 against 55 for Base, while UMM-Reflection reaches 74.
• Best-of-4: the paper says Best-of-4 selection also does not match the method.
• Unseen benchmarks: +10.97 on WISE, +3.48 on OneIG-Bench and +4.63 on T2I-CompBench++, none used in training.
• RL setup: a 3,000-prompt pool over six GenEval families, two roots × 16 siblings per update, and 20 training denoising steps.

How to use it

The page links these artifacts:

• Code: https://github.com/waltstephen/UMM-Reflection
• Models and data: https://huggingface.co/collections/YijiaFan/umm-reflection
• Project page: https://waltstephen.github.io/UMM-Reflection

Inference needs only the unified model, with no verifier, but it does require the multi-round external controller loop. The page gives no install commands, licence or compute figures.

Limits

• Non-standard GenEval protocol. Scores use the authors' one-image, instruction-voice protocol. BAGEL reports 0.82 under its native protocol, so the local numbers are not directly comparable to published leaderboard figures.
• Trained on one benchmark's families. RL prompts come from GenEval families and the reward uses a GenEval-style detector verifier. Gains on the unseen benchmarks are smaller (+3.5 to +4.6), except WISE.
• Proprietary dependency. SFT data depends on GPT-5.5 and Qwen-Image or Qwen-Image-Edit.
• Cost in the reward design. The reward intentionally favours multi-edit trajectories, which may add rounds at inference.
• Single backbone. Only BAGEL is tested.
• Truncated fetch. The page text stops at Table 1, so the full results tables, Appendix details and any stated limitations section were not available.