UMM-Reflection BAGEL-SFT

The reflection-SFT stage of UMM-Reflection (paper, Learning Native Reflection in Unified Models). It is BAGEL-7B-MoT fine-tuned to inspect its own image, decide whether it is done, and, if not, write an edit and render a corrected image, all in one model and one context.

This checkpoint is the initialization for the RL stage. The final RL model is YijiaFan/UMM-Reflection-BAGEL-RL.

Contents

The layout is the same as the base BAGEL repository. Only ema.safetensors differs from base BAGEL-7B-MoT (revision 265d1d48ec8e850a29d3a1f208c2a2ec3cd7577b); every other file is byte-identical to that revision.

Training

• Data. YijiaFan/UMM-Reflection-SFT-Data: 29,529 multi-round reflection trajectories (text-to-image and editing), plus 1,265 prompt-only anchor images generated by base BAGEL.
• Objective. Cross-entropy on controller turns (<think>, SCORE, ACTION EDIT/DONE, edit payload) and verifier turns, and flow matching on image transitions conditioned on the previous image.
• Schedule. Two stages on 16 GPUs: 1,410 steps with cosine LR 2e-6 to 2e-7, then a constant 2e-7 until one epoch of expanded rows is consumed (step 2,969).

Usage

Download it and point the code at the directory.

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

# GenEval-553 evaluation of the SFT model
ARM=sft MODEL_DIR=pretrained/UMM-Reflection-BAGEL-SFT bash scripts/eval/geneval553.sh

# RL from this checkpoint (see the repository README for the reward service)
RL_INIT_DIR=pretrained/UMM-Reflection-BAGEL-SFT \
    NODE_RANK=0 MASTER_ADDR=<node0> MASTER_PORT=29600 bash scripts/rl/train_rl.sh

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. The training data has its own terms; see the dataset card.