PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2608.04887v1cs.CV

TL;DR

Existing OPD matches teacher output velocities, limiting its target to the teacher and leaving internal representation evolution underconstrained. STEP-OPD extrapolates teacher outputs using the teacher–base velocity difference and aligns blockwise representation changes; it improves standard OPD across capability groups and surpasses corresponding single-task teachers.

  • Problem

    Existing OPD optimizes students to match teacher velocities, while output-only supervision leaves blockwise representation evolution underconstrained.

    Source: Introduction · p. 1

  • Method

    STEP-OPD adds a scaled teacher–base velocity difference to the teacher target and aligns the direction and magnitude of student and teacher representation changes across blocks.

    Source: Introduction · p. 2, Task-Specific Output Extrapolation · p. 4

  • Results

    STEP-OPD consistently improves standard OPD across compositional alignment, text rendering, and human preference, and its DiffusionOPD variant raises GenEval overall from 0.927 to 0.961 while surpassing corresponding single-task teachers across capability groups.

    Source: Introduction · p. 2, Conclusion · p. 7

  • Takeaways & Limitations

    The unified student surpasses corresponding single-task teachers across compositional alignment, text rendering, and human preference.

    Source: Conclusion · p. 7

Abstract

from arXiv · show

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.

Introduction

OPD consolidates task-specialized image-generation models into one student, but output matching limits learning to the teacher and leaves internal representation evolution underconstrained. STEP-OPD adds beyond-teacher output targets and blockwise representation supervision, improving benchmark performance across capability groups.

Introduction

  • Standard OPD’s velocity-matching objective makes exact teacher matching its optimum, so it provides no target for improvement beyond the teacher.Even perfect optimization cannot let the student surpass the teacher under this objective.Source: Introduction · p. 1
  • Output matching leaves blockwise representation updates underconstrained, allowing layerwise deviations to offset and weakening transfer of compositional and text-rendering capabilities.Adding representation change alignment to DiffusionOPD improves GenEval and OCR without changing the output target.Source: Introduction · pp. 1–2
  • STEP-OPD extrapolates each task teacher’s velocity using its difference from the shared base model and aligns the direction and magnitude of student and teacher hidden-state changes.These objectives extend the output target beyond teacher matching while directly supervising representation evolution.Source: Introduction · p. 2
  • 0.927 to 0.961: STEP-OPD raises DiffusionOPD’s GenEval overall score and surpasses corresponding single-task teachers across all three capability categories.The method also improves all human-preference metrics and has the strongest overall performance across compositional alignment, text rendering, and human preference.Source: Figure 1 · p. 1, Introduction · p. 2

Related Work

Prior work studies on-policy distillation, guidance through prediction-field differences, and representation-level distillation. STEP-OPD differs by extrapolating beyond teacher outputs during training and aligning changes in representations across network layers.

On-Policy Distillation for Image Generation

Guidance in Diffusion and Flow Models

Representation-Level Distillation

Method

STEP-OPD extends on-policy distillation with output extrapolation and teacher-guided alignment of blockwise representation changes. It reports improvements over OPD baselines and corresponding single-task teachers across compositional alignment, text rendering, and human preference.

Preliminaries

  • The student generates on-policy denoising trajectories, and routed teachers are queried at visited states; ρ selects query states while w_i weights queried timesteps.The trajectory-level reverse KL objective is decomposed into transition-level KL objectives along the student trajectory.Source: Preliminaries · pp. 3–4
  • DiffusionOPD queries all trajectory states uniformly with transition-induced weights, whereas DanceOPD samples one semantic query, favoring low-noise states.DanceOPD uses unit timestep weights.Source: Preliminaries · p. 4

Task-Specific Output Extrapolation

  • Output extrapolation adds a scaled base-to-teacher velocity difference to the teacher target, moving beyond the teacher when αk(s) > 0.At αk(s) = 0, the target recovers standard OPD; the base model and teachers remain frozen, and the coefficient is task-specific.Source: Task-Specific Output Extrapolation · p. 4

Intermediate Representation Change Alignment

  • Intermediate representation change alignment matches student and teacher updates across non-overlapping adjacent block pairs covering all Transformer blocks.It aligns both update direction and magnitude, matching hidden-state changes directly to the teacher rather than extrapolating them from the base model.Source: Intermediate Representation Change… · pp. 4–5
  • STEP-OPD improves DanceOPD and DiffusionOPD across GenEval, OCR, and preference metrics, and surpasses corresponding single-task teachers in all three capability groups.Table 1 compares methods at 512×512 resolution.Source: Table 1 · p. 5

Experiments

Across compositional alignment, text rendering, and human preference, STEP-OPD improves OPD baselines and surpasses the corresponding task-specific teachers. Its component and sensitivity analyses show complementary benefits from representation alignment and task-specific extrapolation strengths.

Experimental Setup

  • Evaluation covers GenEval compositional alignment, OCR text rendering, and four DrawBench preference metrics, using matched prompts, seeds, resolution, sampler, and guidance.Baselines include pretrained generators, the three specialized teachers, diffusion/flow RL methods, DanceOPD, and DiffusionOPD.Source: Experimental Setup · pp. 5–6

Main Results

  • 0.961 vs 0.927 GenEval overall and 0.946 vs 0.941 OCR: STEP-OPD beats DiffusionOPD, surpasses task teachers across groups, and improves all four preference metrics.Both STEP-OPD variants also outperform the representative diffusion/flow RL methods.Source: Main Results · p. 6
  • STEP-OPD more faithfully follows compositional prompts, renders specified text more accurately, and produces detailed, coherent preference-oriented images.Figure 3 compares models on the same prompt in each row across the three capability groups.Source: Main Results · p. 6, Figure 3 · p. 6

Analysis

  • Representation change alignment alone raises GenEval from 0.927 to 0.959 and OCR from 0.941 to 0.944; output extrapolation gives larger preference gains, and combining them retains those gains while further improving alignment and text rendering.Table 2 ablates output extrapolation (OE) and representation change alignment (RCA).Source: Analysis · pp. 6–7, Table 2 · p. 6
  • With RCA, the bowl has cleaner structure and more coherent geometry; with OE, the coat shows richer fabric detail and stronger visual quality.Figure 4 pairs outputs generated from the same prompt and random seed.Source: Analysis · p. 7, Figure 4 · p. 7
  • GenEval performs best at αG = 0.01, OCR at αO = 0.06, and preference at αA = 0.20; performance generally declines beyond each task’s optimum.The results support task-specific coefficients rather than one shared value.Source: Analysis · p. 7, Table 3 · p. 7
  • RCA reduces median paired relative Egap by 96.4% on GenEval, 74.2% on OCR, and 63.5% on Aesthetic, shifting discrepancies toward closer alignment with teacher representation changes.Lower Egap indicates closer alignment with the teacher’s internal representation changes.Source: Analysis · p. 7, Figure 5 · p. 7

Conclusion

STEP-OPD extends on-policy distillation beyond teacher matching by extrapolating targets and aligning inter-layer representation changes. Experiments show consistent improvements over standard OPD and a unified student that surpasses the corresponding single-task teachers across all three capability groups.

Conclusion

  • STEP-OPD addresses teacher-matching and output-only supervision limits with targets extrapolated from the shared base model toward each task-specific teacher, plus alignment of inter-layer representation changes.The alignment matches both the direction and magnitude of those changes.Source: Conclusion · p. 7
  • Experiments across compositional alignment, text rendering, and human preference show consistent improvements over standard OPD methods.Source: Conclusion · p. 7
  • The unified student surpasses the corresponding single-task teachers across all three capability groups.Source: Conclusion · p. 7

Related papers

Loading 2608.04887v1…