Search papers, labs, and topics across Lattice.
This paper introduces STEP-OPD, an innovative framework for on-policy distillation in image generation that enhances the learning process by expanding the student's targets beyond merely matching teacher outputs. By incorporating velocity differences between task-specific teachers and a shared base model, STEP-OPD constrains the student's internal representation evolution, leading to improved capability transfer across layers. Experimental results demonstrate that STEP-OPD significantly boosts performance metrics, including a GenEval score increase from 0.927 to 0.961, outperforming single-task teachers in multiple domains.
Output extrapolation in STEP-OPD enables a unified student model to surpass the performance of its specialized teachers across all evaluated tasks.
On-policy distillation (OPD) has become an effective approach for consolidating multiple task-specialized image generation models into a single student. However, existing OPD methods optimize the student mainly to match the teacher's output velocity, making the teacher the upper limit of the optimization objective. While output-level supervision alone leaves the student's blockwise representation evolution underconstrained, which weakens the transfer of capabilities that must be progressively developed across layers. We propose STEP-OPD, an on-policy distillation framework for image generation that extends the student's learning target beyond the teacher and introduces explicit constraints on its internal representation evolution. Instead of treating the teacher as the final target, we use the velocity difference between each task-specific teacher and the shared base model as a direction for further learning and add a scaled version of this difference to the teacher velocity. In addition, we align the direction and magnitude of representation changes between the student and teacher, enabling the student to learn how representations are progressively transformed across network blocks. Experiments on compositional alignment, text rendering, and human preference show that our method consistently improves Standard OPD methods. In particular, it increases the GenEval score of DiffusionOPD from 0.927 to 0.961, while also improving OCR and all preference-based metrics. The resulting unified student surpasses the corresponding single-task teachers across all three capability groups, showing that output extrapolation enables beyond-teacher learning. And representation change alignment provides complementary guidance for the student's internal transformations.
DAMO