STEP-OPD: Rethinking Output Targets and Internal Dynamics in On-Policy Distillation for Diffusion Models

2026-08-05

AI Summary

Existing policy distillation methods focus solely on matching the teacher model’s outputs, neglecting the evolution of internal representations and thereby limiting the student’s performance ceiling. This work proposes the STEP-OPD framework, which leverages the velocity difference between the teacher and a base model to extrapolate an optimal learning direction and explicitly aligns both the direction and magnitude of block-level representational changes between teacher and student. By integrating in-policy distillation, velocity-difference extrapolation, and representation evolution alignment within a diffusion model paradigm, the method achieves superior output fidelity and representational transfer beyond the teacher. Evaluated on GenEval, it improves the score from 0.927 to 0.961, and consistently outperforms single-task teacher models in OCR accuracy and human preference metrics, marking the first unified student model to surpass its teacher across all capability dimensions.

Abstract

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.

Problem

• on-policy distillation
• diffusion models
• representation evolution
• output targets
• model compression

Innovation

• on-policy distillation
• diffusion models
• representation evolution
• output extrapolation
• velocity alignment