FocusOPD: Structure-Guided On-Policy Distillation for Diffusion Models

Xinrui Lin; Shumin Wang; Suqi Song; Fan Li; Yanyong Zhang

cs.CVOct 15, 2026

FocusOPD: Structure-Guided On-Policy Distillation for Diffusion Models

Authors: Xinrui Lin, Shumin Wang, Suqi Song, Fan Li, Yanyong Zhang

Organizations: University of Science and Technology of China, Hefei, China · The Hong Kong University of Science and Technology (Guangzhou) · Huawei Noah’s Ark Lab

Abstract

Multi-teacher on-policy distillation enables rectified-flow models with diffusion-transformer architectures to combine task-specific capabilities within a single student. However, matching model outputs does not guarantee that their internal representations are aligned, and it remains unclear where and when additional feature supervision is most useful. We analyze the teacher-student feature mismatch across Transformer blocks and denoising timesteps, and find that the mismatch is concentrated in deeper Transformer blocks during early denoising and persists after OPD distillation, indicating ineffective distillation. Motivated by this observation, we introduce FocusOPD, which augments the base OPD objective with token-wise directional feature alignment restricted to deep layers and early denoising steps. Across two rectified-flow backbones and multiple OPD objectives, FocusOPD improves capability transfer in compositional generation, text rendering, and preference-based evaluation. On the primary benchmark, it improves over DiffusionOPD from 0.953 to 0.967 on GenEval, from 23.797 to 23.944 on PickScore, and from 6.052 to 6.199 on LAION-Aesthetic. Ablations show that selective deep and early supervision is more effective and efficient than alignment over the full network and trajectory.

Figures & tables

Figure 1: Teacher–student feature MSE across (a) transformer blocks and (b) denoising steps, ordered from high to low noise. Different curves indicates different groups of (a)denoising steps (b)blocks.

Panel of Figure 2 — Figure 2: Hidden-NMSE heatmaps. Columns show checkpoints 0, 500, 1000, and the checkpoint-1000 minus checkpoint-0 difference.

Figure 3: Deep-block decomposition across three teachers. Terms are normalized by total mean discrepancy energy. Block indices are 18 for SD3.5-M and 22 for Z-Image. Error bars: 95% prompt-bootstrap intervals.

Model	Teacher	$t$ -ctrl.	Scaling
SD3.5-M	PickScore	1.303	0.748
	OCR	0.728	0.547
	GenEval	5.932	0.879
Z-Image	PickScore	1.938	0.947
	OCR	3.380	0.786
	GenEval	8.243	0.862

Table 1: Held-out prediction NMSE. Lower is better.

Figure 4: Overview of FocusOPD. Teacher and student are evaluated on shared states from student trajectories. Output supervision follows the base OPD protocol, while directional feature supervision targets deep layers during early denoising. Only the trained student is used at inference.

SD3.5-M
Method	GenEval							OCR	DrawBench
Method	Sing.	Two	Cnt.	Col.	Pos.	C-Att.	Ovr.	Acc.	Pick	Aes.	HPS	ImgR
Single-task teacher references
Base model	0.963	0.889	0.613	0.809	0.250	0.570	0.682	0.675	22.411	5.388	0.281	0.901
GenEval teacher	1.000	1.000	0.950	0.936	0.960	0.900	0.958	0.507	22.061	5.202	0.248	0.611
OCR teacher	0.988	0.838	0.663	0.787	0.310	0.480	0.678	0.961	22.351	5.273	0.273	0.938

Table 2: Multi-teacher distillation on GenEval, OCR, and DrawBench. The table reports all GenEval categories, OCR accuracy, and four independent DrawBench quality metrics. All metrics are higher-is-better; Scores are reported at checkpoint 1000. Gray rows highlight our method applied to each baseline.

Figure 5: Qualitative comparison of generation results. Our method produces images with better compositional alignment, text rendering quality, and visual aesthetics compared to baseline methods.

Configuration	Depth	Time	GenEval	OCR	Pick	Aes	HPS	ImgR	Time/up (s)	Cache (MiB)
Output KD only	–	–	0.951	0.961	23.654	6.146	0.339	1.502	16.54	0
w/o depth selection	All-24	Front 40%	0.965	0.959	23.565	6.208	0.341	1.505	19.88	864
w/o temporal selection	Last-12	Full	0.951	0.960	23.556	6.183	0.341	1.500	19.88	2160
All blocks, full time	All-24	Full	0.953	0.958	23.574	6.198	0.342	1.501	21.77	2160
FocusOPD	Last-12	Front 40%	0.960	0.959	23.545	6.214	0.340	1.497	18.99	432

Table 3: Ablation of feature supervision components on SD3.5-M with DiffusionOPD. Training time per update is measured over steps 11–1000; teacher feature cache is the CPU-stored tensor payload per rank. All feature variants use cosine distance with

\lambda=0.3

; all methods train the full 24-block transformer.

Table 9

Feature loss type	GenEval	OCR	PickScore	Aes	HPSv2.1	ImgRwd
Output KD only (no feature)	0.951	0.961	23.654	6.146	0.339	1.502
Teacher-energy-normalized MSE	0.955	0.963	23.509	6.210	0.340	1.491
Cosine distance	0.960	0.959	23.545	6.214	0.340	1.497

Table 6: Feature loss formulations on SD3.5-M with Last-12 and front-40% coverage. Both feature variants use

\lambda=0.3

. Normalized MSE divides each sample/layer squared error by teacher mean squared energy, then averages over selected layers; the denominator floor is

10^{-8}

. This compares formulations at a fixed coefficient, not matched gradient magnitudes.

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Figure 6: Deep-block decomposition across three teachers for all tested

\sigma

. Terms are normalized by total mean discrepancy energy. Block indices are 18 for SD3.5-M and 22 for Z-Image. Error bars: 95% prompt-bootstrap intervals.

Diffusion Model Distillation

Diffusion model distillation transfers a large model’s generative behavior to faster or smaller models that need fewer sampling steps.

See all 191 papers

STEP-OPD: Rethinking Output Targets and Internal Dynamics in On-Policy Distillation for Diffusion Models
Aug 5, 2026Qingyan Wei, Guangzhao Li, Xiaobing Tu +5Diffusion Model Distillation Diffusion-Based Image Generation
Shanghai Jiao Tong University · Shanghai Innovation Institute · Alibaba Group
Poly-OPD: Heterogeneous Multi-Teacher On-Policy Distillation for Capability-Selectable Flow Models
Aug 5, 2026Siming Fu, Haojun Xu, Ruizhe He +9Diffusion Model Distillation Multi-Teacher Knowledge Distillation
1Joy Future Academy · 2Zhejiang University · 3Beihang University
GFD-OPD: Guidance-Folded On-Policy Distillation of Diffusion Models Across Scales
Sep 30, 2026Zhenxing Zhang, Jiayan Teng, Wenxu Wu +6Diffusion Model Guidance Classifier-Free Guidance
Hefei University of Technology · Zhipu AI · Tsinghua University

Aug 5, 2026J/K move · Enter open · S save

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

Qingyan Wei, Guangzhao Li, Xiaobing Tu, Yinggui Wang+4

Shanghai Jiao Tong University · Shanghai Innovation Institute · Alibaba Group

Diffusion Model Distillation Diffusion-Based Image Generation On-Policy Distillation

Open paper AskPDF arXiv

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.

Abstract

Diffusion Model Distillation

Explore similar work

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

Poly-OPD: Heterogeneous Multi-Teacher On-Policy Distillation for Capability-Selectable Flow Models

GFD-OPD: Guidance-Folded On-Policy Distillation of Diffusion Models Across Scales