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

Qingyan Wei · Guangzhao Li · Xiaobing Tu · Yinggui Wang
Xiantao Zhang · Jinkui Ren · Xiaohong Liu · Linfeng Zhang

---

📢 Updates

2026-08-05  ·  Paper Release

• 🎉 Paper released on arXiv: STEP-OPD
• 📄 PDF: arxiv.org/pdf/2608.04887
• ⏳ Code & checkpoints coming soon

Full History

Date	Milestone
2026-08-05	Paper released on arXiv. Repo online. Code coming soon.

---

Figure 1. Motivation and overview of STEP-OPD. Output extrapolation enables beyond-teacher learning, representation change alignment improves internal capability transfer, and the unified student achieves the strongest overall multi-task profile.

---

💡 Overview

STEP-OPD improves on-policy distillation (OPD) for multi-task diffusion / flow alignment. Standard OPD only matches the teacher velocity, so the teacher becomes the learning upper bound, and the student's block-wise representation evolution is underconstrained.

We introduce two complementary components:

	Component	Idea
🎯	Output Extrapolation (OE)	Extrapolate along the base → teacher direction to build a beyond-teacher target
🔗	Representation Change Alignment (RCA)	Align direction & magnitude of inter-block hidden-state changes

What you get

• Beyond-teacher learning on the student's own trajectory
• Better transfer of compositional / text-rendering capabilities
• Plug-and-play on top of DiffusionOPD and DanceOPD
• GenEval 0.927 → 0.961, with gains on OCR and preference metrics

---

🧠 Method

Figure 2. Overview of STEP-OPD. Left: Representation Change Alignment aligns the direction and magnitude of hidden-state changes across selected adjacent blocks. Right: Task-Specific Output Extrapolation constructs a beyond-teacher velocity target along the teacher–base improvement direction.

---

🎯 Output Extrapolation

Treat the teacher–base velocity gap as an improvement direction and move the target beyond the teacher:

$$ \tilde{\mathbf{v}}_{k,i}(s)=\mathbf{v}_{T_k,i}+\alpha_k(s),(\mathbf{v}_{T_k,i}-\mathbf{v}_{B,i})=\mathbf{v}_{B,i}+(1+\alpha_k(s)),(\mathbf{v}_{T_k,i}-\mathbf{v}_{B,i}) $$

• $\alpha = 0$ recovers standard OPD
• $\alpha > 0$ enables beyond-teacher learning
• Task-specific coefficients with a short warm-up (main setting: $\alpha_G=0.01$, $\alpha_O=0.06$, $\alpha_A=0.20$)

---

🔗 Representation Change Alignment

Align how representations evolve across adjacent Transformer blocks. For an adjacent block pair $p=(\ell_0,\ell_1)$:

$$ \Delta\mathbf{h}^{p}_{M,i}=\mathbf{h}^{\ell_1}_{M,i}-\mathbf{h}^{\ell_0}_{M,i},\quad M\in{S_\theta,T_k} $$

We match both the direction and magnitude of $\Delta\mathbf{h}$ between student and teacher, so the student learns progressive internal transformations rather than only the final velocity.

---

📊 Results

All numbers at 512×512. STEP-OPD improves both DanceOPD and DiffusionOPD, and can surpass single-task teachers.

Main Comparison

Model	GenEval ↑	OCR ↑	PickScore ↑	Aesthetic ↑	HPSv2.1 ↑	ImgRwd ↑
GenEval Teacher	0.935	0.401	22.069	5.247	0.249	0.614
OCR Teacher	0.527	0.938	22.282	5.256	0.272	0.918
Aesthetic Teacher	0.486	0.576	24.031	6.220	0.346	1.504
DanceOPD	0.916	0.943	23.836	6.159	0.328	1.423
+ STEP-OPD	0.944	0.957	23.994	6.297	0.344	1.518
DiffusionOPD	0.927	0.941	23.941	6.208	0.340	1.503
+ STEP-OPD (Ours)	0.961	0.946	24.053	6.321	0.349	1.533

Qualitative

Compared with the base model, task teachers, and DiffusionOPD, STEP-OPD better consolidates prompt adherence, text rendering, and visual quality.

---

🛠️ Code & Models

⏳ Coming Soon — training / evaluation code and checkpoints will be released here.

Planned release:

• STEP-OPD training (OE + RCA)
• Evaluation scripts (GenEval / OCR / preference)
• Configs & student checkpoints

---

📚 Citation

@article{wei2026step,
  title={STEP-OPD: Rethinking Output Targets and Internal Dynamics in On-Policy Distillation for Diffusion Models},
  author={Wei, Qingyan and Li, Guangzhao and Tu, Xiaobing and Wang, Yinggui and Zhang, Xiantao and Ren, Jinkui and Liu, Xiaohong and Zhang, Linfeng},
  journal={arXiv preprint arXiv:2608.04887},
  year={2026}
}

---

🙏 Acknowledgement

Built upon and inspired by: DiffusionOPD · Flow-GRPO · DanceOPD · DiffusionNFT

---

📬 Contact

Feel free to reach out for discussion or collaboration:

📧 qingyanwei@sjtu.edu.cn

---

⭐ If you find STEP-OPD useful, please star this repo and cite our paper.
Future updates will appear in Updates.