StepX-OPD: Rethinking Multi-Teacher On-Policy Distillation in Diffusion Models

StepX-OPD. Rethinking Multi-Teacher On-Policy Distillation in Diffusion Models. By Chenru Wang and seven others, from StepFun and Westlake University. Published on the third of November.

A text-to-image diffusion model that is ready to deploy has to do several things at once. It must follow instructions, render text, and match human preference, and it must also sample in only a few steps. Today those two demands are met in separate stages. First, a full-step model consolidates the capabilities through multi-teacher on-policy distillation. Then that model becomes a fixed teacher for step distillation. The authors point to two costs. Every full-step model they evaluate scores lower once compressed. And the second stage cannot start until the first has finished.

StepX-OPD merges the stages. A single student acquires every capability while it is learning to sample in a few steps. The consequence is that frozen specialist teachers supervise the student on the states its own rollout actually visits, rather than on a longer trajectory that the student never runs. Put another way, the teachers grade the student on the path it will take when deployed, not on a path belonging to some other model. And because there is only one stage, nothing has to wait for an earlier model to finish before the few-step student can begin learning its capabilities.

One refinement follows from an analysis of the objective. When the rollout's KL is decomposed, it weights velocity mismatch most heavily near the clean end. In response, the authors also regress the student's velocity on the low-noise part of the rollout. A separate observation concerns a side effect. Integrating capabilities makes models write prompt words into images, even for prompts that never asked for any text. This is a finding in its own right, about capability integration in general rather than about this method in particular. The paper gives a variant of the method that lowers the effect.

Evaluation covers two architectures, SD3.5-Medium and FLUX.1-dev, measured on instruction following, text rendering, image quality and human preference. On both, StepX-OPD reaches the highest normalised average score among the four-step methods compared. It also exceeds the forty-step Flow-OPD model, a model allowed far more sampling steps than the four-step student.

Cost is reported too. On SD3.5-Medium, StepX-OPD trains with about thirty-seven percent fewer GPU-hours than the sequential route of DiffusionOPD followed by step distillation.

Here the order of operations is the thing being questioned. Teaching capabilities first and speed second means the specialists grade a trajectory the deployed model will never take, and the compression afterwards gives back some of what was learned. Folding both into one student keeps the teachers looking at the few-step path that actually ships, and on the reported numbers it costs less compute to get there.