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

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 leaving the student's blockwise representation evolution underconstrained.

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 (output extrapolation). In addition, we align the direction and magnitude of representation changes between the student and teacher (representation change alignment), 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.

0.961 GenEval Overall vs. 0.927 DiffusionOPD

0.946 OCR Accuracy surpasses OCR teacher

6.321 Aesthetic Score best among compared

3-in-1 Unified Student beats all task teachers

Motivation & Overview

Standard OPD is capped by exact teacher matching and underconstrains internal dynamics. STEP-OPD pushes the output target beyond the teacher and supervises blockwise representation changes.

STEP-OPD teaser: motivation and overview. — **Upper left:** standard OPD reaches its optimum at exact teacher matching, whereas output extrapolation shifts the learning target beyond the teacher. **Lower left:** representation change alignment improves GenEval and OCR without changing the output supervision target. **Right:** STEP-OPD achieves the strongest overall profile across compositional alignment, text rendering, and human preference.

Beyond-Teacher Targets

Extrapolate along the base→teacher velocity direction so the student is no longer upper-bounded by teacher matching.

Internal Dynamics

Align direction and magnitude of inter-block hidden-state changes, transferring how the teacher progressively transforms representations.

Multi-Expert Consolidation

One unified student consolidates compositional, OCR, and preference teachers—and surpasses each specialist on its own task.

Method

Built on the student's on-policy denoising trajectory, STEP-OPD adds two complementary objectives: task-specific output extrapolation (OE) and representation change alignment (RCA).

Overview of the STEP-OPD framework. — **Overview of STEP-OPD.** The student generates an on-policy trajectory; a latent state is selected by the underlying OPD query strategy. **Left — RCA:** align direction and magnitude of hidden-state changes across adjacent Transformer blocks between the routed teacher and the student. **Right — OE:** the velocity difference between the teacher and shared base model defines a task-specific improvement direction, scaled and added to construct a beyond-teacher target.

Component A · OE

Task-Specific Output Extrapolation

For each task teacher T_k and shared base B, construct ṽ = v_T + α_k(v_T − v_B). When α=0 this recovers standard OPD; when α>0 the target moves beyond the teacher along the base-to-teacher direction. Coefficients are task-specific (GenEval 0.01 / OCR 0.06 / Aesthetic 0.20) with a short warm-up.

Component B · RCA

Representation Change Alignment

Rather than matching absolute hidden states, align inter-block changes Δh = h^l₁ − h^l₀ in both direction (token-wise RMS-normalized MSE) and magnitude (log-RMS). Applied to 12 adjacent pairs covering all 24 Transformer blocks of SD3.5-M, RCA constrains internal evolution without extrapolating internal targets.

Quantitative Results

STEP-OPD consistently improves both DanceOPD and DiffusionOPD across compositional alignment (GenEval), text rendering (OCR), and human preference, and further surpasses single-task teachers.

Model	Sing	Two	Cnt	Col	Pos	C-Attr	Overall	OCR	Pick	Aes	HPS	ImgRwd
(a) Base Models
FLUX.1-Dev	0.778	0.427	0.628	0.649	0.068	0.253	0.467	0.496	22.648	5.737	0.292	0.796
SD3.5-L	0.703	0.389	0.525	0.628	0.140	0.265	0.442	0.635	21.973	5.299	0.264	0.688
SD3.5-M+CFG	0.869	0.495	0.569	0.715	0.148	0.355	0.525	0.559	22.398	5.389	0.280	0.853
(b) Single-Task Teachers
GenEval Teacher	0.994	0.952	0.944	0.923	0.962	0.833	0.935	0.401	22.069	5.247	0.249	0.614
OCR Teacher	0.872	0.472	0.594	0.723	0.145	0.355	0.527	0.938	22.282	5.256	0.272	0.918
Aes Teacher	0.744	0.422	0.588	0.676	0.163	0.323	0.486	0.576	24.031	6.220	0.346	1.504
(c) Diffusion/Flow RL & Multi-Objective OPD
Flow-GRPO	1.000	0.972	0.950	0.924	0.965	0.845	0.943	0.924	23.500	5.922	0.321	1.346
DiffusionNFT	1.000	0.968	0.954	0.937	0.963	0.818	0.940	0.910	23.800	6.010	0.331	1.490
DanceOPD	0.975	0.939	0.934	0.904	0.940	0.803	0.916	0.943	23.836	6.159	0.328	1.423
+ STEP-OPD (Ours)	1.000	0.960	0.950	0.941	0.970	0.845	0.944	0.957	23.994	6.297	0.344	1.518
DiffusionOPD	1.000	0.962	0.938	0.926	0.948	0.790	0.927	0.941	23.941	6.208	0.340	1.503
+ STEP-OPD (Ours)	1.000	0.990	0.988	0.957	0.980	0.853	0.961	0.946	24.053	6.321	0.349	1.533

Best in bold blue, second-best underlined. Evaluation at 512×512.

Qualitative Comparison

Across compositional alignment, text rendering, and human preference, STEP-OPD consolidates the strengths of specialized teachers with more accurate prompt adherence and text rendering while keeping strong visual quality.

Ablation Study

RCA and OE emphasize complementary aspects of generation. Combining them yields the best balance across all three capability groups.

Component Ablation (on DiffusionOPD)

RCA lifts GenEval/OCR; OE drives preference gains; both together win overall.

DiffusionOPD0.927 / 0.941 / 6.208
+ RCA only0.959 / 0.944 / 6.217
+ OE only0.957 / 0.943 / 6.312
+ STEP-OPD (OE+RCA)0.961 / 0.946 / 6.321

Metrics: GenEval / OCR / Aesthetic

RCA Gap Reduction

Median paired relative reduction of representation-change gap under RCA.

GenEval trajectory−96.4%
OCR trajectory−74.2%
Aesthetic trajectory−63.5%

Output matching alone leaves blockwise dynamics underconstrained.

**Qualitative ablation of RCA and OE.** Without RCA, the bowl shows spurious internal structures; RCA yields cleaner geometry. Without OE, fabric texture is flatter; OE produces richer fine-grained detail. Each pair uses the same prompt and random seed.

Representation change gap distributions with and without RCA. — **Representation change gap distributions.** Adding RCA to DiffusionOPD shifts discrepancy distributions toward lower values across GenEval, OCR, and Aesthetic, confirming that the student explicitly learns the teacher's inter-block transformation dynamics.

BibTeX

@misc{wei2026stepopd,
  title         = {STEP-OPD: Rethinking Output Targets and Internal Dynamics
                   in On-Policy Distillation for Diffusion Models},
  author        = {Qingyan Wei and Guangzhao Li and Xiaobing Tu and Yinggui Wang
                   and Xiantao Zhang and Jinkui Ren and Xiaohong Liu and Linfeng Zhang},
  year          = {2026},
  eprint        = {2608.04887},
  archivePrefix = {arXiv},
  primaryClass  = {cs.CV},
  url           = {https://arxiv.org/abs/2608.04887}
}