UMM-Reflection BAGEL-RL

The final model of UMM-Reflection (paper, Learning Native Reflection in Unified Models). Starting from UMM-Reflection-BAGEL-SFT, it is trained with multi-round Flow-GRPO. For each prompt, it samples 16 complete reflection trajectories from the same first image and optimizes the text (controller) and flow (image) heads jointly against a graded six-family GenEval reward.

Results

Local evaluation with the same prompts, resolution, and scorer for every row. GenEval uses the 553 official prompts with up to three repair rounds.

Model	GenEval	WISE	OneIG	T2I-CompBench++
BAGEL-Base	0.71	0.55	0.80	0.49
BAGEL-SFT	0.72	0.63	0.79	0.50
UMM-Reflection (this model)	0.84	0.74	0.83	0.55

On GenEval-553, 64.9% of images that are wrong after the first render end up correct after reflection, against 20.6% for the SFT model.

The layout is the same as the base BAGEL repository. ema.safetensors merges the 784 trained language-model tensors and the nine trained auxiliary tensors (time_embedder, vae2llm, llm2vae, latent_pos_embed) of RL checkpoint-1000 into the SFT weights. Every other file is byte-identical to BAGEL-7B-MoT revision 265d1d48ec8e850a29d3a1f208c2a2ec3cd7577b.

File	sha256
`ema.safetensors`	`34b1b232d971c24a279ab4863162ddb9155adc757a2f39b5f934a4396d824345`

Training

Initialization. UMM-Reflection-BAGEL-SFT.
Prompts. A 3,000-prompt pool over the six GenEval families, built from Flow-GRPO's GenEval training metadata and disjoint from the GenEval-553 test prompts.
Schedule. 1,000 steps on 16 GPUs. Each rank takes 2 prompts with 16 sibling trajectories each, 20 denoising steps, and up to three repair rounds. Text and flow learning rates are both a constant 5e-6.

Usage

git clone https://github.com/waltstephen/UMM-Reflection && cd UMM-Reflection
huggingface-cli download YijiaFan/UMM-Reflection-BAGEL-RL \
    --local-dir pretrained/UMM-Reflection-BAGEL-RL

ARM=rl MODEL_DIR=pretrained/UMM-Reflection-BAGEL-RL bash scripts/eval/geneval553.sh

Evaluation uses 512 x 512 images, 50 denoising steps, CFG 3.0, controller temperature 0.5, and at most three repair rounds. REPAIR_RNG_STEP defaults to 500 for ARM=rl, the convention of the reported numbers. See the repository README for the GenEval verifier assets.

Citation

@article{ummreflection2026,
  title   = {Learning Native Reflection in Unified Models
             with Interleaved Reinforcement Learning},
  author  = {Fan, Yijia and Huang, Ziqi and Cai, Zhongang and Li, Yan and
             Wen, Zimo and Yin, Wanqi and Diao, Haiwen and Liu, Ziwei},
  journal = {arXiv preprint arXiv:2609.35767},
  year    = {2026}
}

License

Apache 2.0, the same as BAGEL. ae.safetensors is the FLUX.1 autoencoder shipped with BAGEL, under its own Apache 2.0 license.

YijiaFan/UMM-Reflection-BAGEL-RL