Self-OPD: On-Policy Distillation for Flow Matching Models without Teacher
Shiyi Zhang, Mushui Liu, Yunze Tong, Wanggui He, Siyu Zou, Jinlong Liu, Yunlong Yu, Jian Song, Hao Jiang, Pipei Huang, Bo Zheng
cs.CV
2026-08-27
Self-OPD turns student SDE branches and an ODE self-reference into dense pull-push supervision. On SD3.5-Medium: GenEval 0.95, OCR 97.5%; mixed-reward wall-clock ~44–48 h vs 97 h.
Flow matching is already a default backbone for high-quality visual generation. Aligning it to downstream goals such as OCR, composition, and aesthetics is still awkward. RL methods like Flow-GRPO treat denoising as a sequential decision problem and score a full trajectory at the end. Credit assignment is high-variance, and mixed rewards fight each other. On-policy distillation (OPD) regresses the student velocity toward a teacher at every step, which is more stable. Each new objective then needs its own specialist teacher, the student is capped by that teacher, and merging several teachers at the field level can point the update in opposing directions.
The question from Tsinghua, Zhejiang University, and Alibaba is whether OPD can keep dense per-step supervision after the external teacher is removed.
Self-OPD explores locally on the student's own trajectory. At timestep $tj$ it computes the deterministic next state $x{t{j+1},\theta}$, then draws $K$ SDE candidates around it ($K=8$, $\eta=0.7$ in the experiments). Each candidate is rolled out with a deterministic ODE to a clean image and scored by the task reward. A pure ODE from the same parent state is the self-reference $r^{\mathrm{ode}}$. Branch advantage is the standardized gap to that baseline.
The update is all-branch pull-push, not regression onto the best branch only. Positive-advantage velocities pull the student; negative ones push it away. A direction-aware gate $dk\in[0,1]$ damps repulsion when a bad branch points the same way as the best branch, so the push does not cancel the pull. The loss is normalized by SDE transition variance and the affine map into velocity space, which matches the precision term of a per-step reverse KL, so the timescale across steps is not a free knob. Timestep weights favor early-to-mid denoising, where global layout is set, while every stage still appears in training.
Multi-objective alignment does not add gradients in parameter space. Each reward is z-scored over the branch set, then mixed with task weights (on OCR, $\lambda{\mathrm{OCR}}:\lambda{\mathrm{PickScore}}:\lambda{\mathrm{HPSv2}}=3:1:1$) into a scalar used only to rank branches. The regression target remains a concrete velocity that already sits in the joint high-reward region, so black-box rewards are admissible.
The backbone is SD3.5-Medium at 512×512, with LoRA on the transformer. Comparators include Flow-GRPO, GRPO-Guard, DiffusionNFT, and teacher-based Flow-OPD and DiffusionOPD.
Under single-reward training Self-OPD reports GenEval 0.95, OCR 97.5%, PickScore 24.79, and HPSv2 0.3665, all above Flow-GRPO and GRPO-Guard with the same base model and rewards. A single mixed-reward model still posts GenEval 0.95 and OCR 96.0%. On the same GenEval/OCR test images it reaches PickScore 23.87 and HPSv2 0.3214, above DiffusionOPD at 22.72 / 0.2676 and Flow-OPD at 23.43 / 0.3042.
The two fusion styles show up in preference distributions. When DiffusionOPD moves from aesthetic prompts to GenEval/OCR prompts, mean PickScore drops by 1.23 and HPSv2 by 0.105. Self-OPD's two distributions almost overlap ($\Delta=0.48$ / $0.020$), because branches are chosen inside the joint high-reward region rather than by compromising teacher gradients in parameter space.
Ablations: Best-of-K regression is unstable and barely beats the base; all-branch pull-push stabilizes training; adding $|\Delta t|$-aligned KL normalization speeds convergence. Unbounded repulsion $dk\in[0,2]$ collapses later. Uniform timestep weights eventually catch up but converge more slowly; a TIS variant that oversamples early steps falls below the base.
Wall-clock: DiffusionOPD's slowest teacher takes 85.8 h in parallel, plus 11.3 h of student distillation, 97 h total, with the student idle during teacher training. Self-OPD from scratch hits OCR 0.946 in about 62 h and GenEval 0.915 in about 90 h. Warm-starting from parallel single-reward specialists (37 h wall-clock) then mixing reaches those marks in about 11 h and 7 h more, 44–48 h total, roughly 2× faster with a higher finish.
For flow-matching alignment, a new reward no longer has to start with a teacher. Dense per-step supervision can come from the student's own neighborhood, and mixed objectives can be composed at the reward level without field-level gradient fights. For a product that needs one image to pass OCR, composition, and aesthetics together, the "same test images" protocol is closer to the actual goal than winning each metric on a different prompt family.
The cost is K branched rollouts and scorers at every trained step, so training compute is not free. What it removes is tens of hours of teacher pretraining, and the teacher ceiling.
The main HTML tables are garbled, so per-method cell values beyond the numbers stated in the prose cannot be trusted. The study is tied to SD3.5-Medium, 512×512, and LoRA; transfer to FLUX or higher resolution is unknown. Rewards are still OCR / GenEval / PickScore / HPSv2. Black-box composition does not prove stability under human feedback. K=8 SDE branches plus ODE rollouts may eat the sample-efficiency gain of dense supervision; there is no FLOP-matched comparison. Sensitivity is real: Best-of-K is unstable, unbounded repulsion collapses. There is no human eval and no report of regression on generic generation from the base model.