PaperlayerView exact revision on arXiv

Source-linked AI summary

Simplifying and Improving On-Policy Distillation for Few-Step Diffusion

Jiawei Zhang, Ziyu Wan, Nanye Ma, Long Zhao, Dong Chen, Hongyu Liu, Qifeng Chen, Dongdong Chen

arXiv:2610.18472v1cs.CVcs.LGeess.IV

TL;DR

Diffusion and flow sampling is costly, motivating methods that compress generation into a few steps. SMD uses on-policy velocity regression with training-only corrections and repulsion; at 4-NFE it outperforms strong few-step baselines and surpasses multi-step teachers on several benchmarks.

  • Problem

    Diffusion and flow sampling commonly requires tens to hundreds of sequential evaluations, while few-step approaches may rely on online score estimation or added components for within-interval velocity.

    Source: §1 · p. 9

  • Method

    SMD trains few-step students by on-policy teacher-velocity regression, adding a training-only Pivot proxy for within-interval variation and repulsion from weaker teacher predictions.

    Source: §1 · p. 9

  • Results

    Across image, 3D, video, and audio-video tasks and models up to 33B, SMD at 4-NFE beats strong few-step baselines and exceeds multi-step teachers on several benchmarks.

    Source: §1 · p. 10, §8 · p. 38

  • Takeaways & Limitations

    SMD extends on-policy regression beyond self-distillation to cross-model, multi-teacher, and causal-video distillation.

    Source: §8 · p. 38

    For high-curvature trajectories, teacher velocity can change substantially within one transition, creating conflicting regression targets that may make optimization less stable.

    Source: §3.2 · p. 14

Abstract

from arXiv · show

Few-step distillation reduces the sampling cost of diffusion and flow models, but existing approaches often require auxiliary score models or specialized architectures. We show that simple on-policy velocity regression already provides a strong foundation for 4-NFE generation. To represent teacher velocity variation within each coarse transition, we introduce Pivot, a training-only residual proxy that separates local corrections from the student's interval-mean prediction. We further repel the student from weaker teacher predictions, providing a lightweight mechanism for avoiding undesirable modes. The resulting framework, Simplified Matching Distillation (SMD), remains entirely regression-based, requires neither online score estimation nor adversarial training, and leaves the deployed generator unchanged. We evaluate SMD on image generation and editing, 3D generation, video synthesis, and joint audio--video generation, scaling to 33 billion parameters. At 4-NFE, SMD outperforms strong few-step baselines and surpasses multi-step teachers on several benchmarks. It also supports cross-model, multi-teacher, and causal-video distillation.

1 Introduction

The introduction motivates few-step distillation as a way to reduce costly sampling, then presents on-policy velocity matching and SMD’s Pivot and repulsive additions. It reports strong 4-NFE results across modalities and scales.

1 Introduction

  • Diffusion and flow models require tens to hundreds of sequential evaluations, motivating few-step distillation to reduce sampling cost for deployment and interactive applications.The introduction frames few-step distillation as compressing a pretrained teacher into a process with a few steps.Source: §1 · p. 9
  • DMD uses an online fake model to track student scores, while trajectory-based alternatives remove online score estimation but add components to represent dense within-interval velocity.The introduction notes that DMD and variants can suffer from low diversity, unstable training, and over-saturated output.Source: §1 · p. 9
  • Simple on-policy velocity matching is a strong few-step baseline: it matches student and teacher predictions at the same states while keeping student velocity constant within each interval.The baseline requires no auxiliary fake score model, architectural modification, or elaborate policy design, and produces photorealistic, varied 4-NFE outputs.Source: §1 · p. 9
  • SMD adds a training-only Pivot proxy for within-interval velocity variation and repulsion from weaker teacher predictions, while remaining regression-based and leaving inference unchanged.It requires neither score estimation nor adversarial objectives and does not modify the generator architecture or inference pipeline.Source: §1 · p. 9
  • At 4-NFE, SMD outperforms compared few-step baselines on GenEval and DPG-Bench, while its FLUX.1-dev student surpasses the 50-NFE teacher on both.SMD also supports image generation and editing, 3D generation, video synthesis, and joint audio-video generation, scaling to models with up to 33B parameters.Source: §1 · p. 10

2 Preliminaries

The preliminaries formulate diffusion and flow generation as transport along a learned velocity field, then describe sampling, guidance, and the finite time grid. Few-step sampling uses only a small number of coarse transitions, making each transition’s accuracy important.

2 Preliminaries

  • A linear path interpolates data and Gaussian noise, with conditional velocity x1 − x0; flow matching trains a model to predict it, while the marginal field is its conditional expectation.The teacher prediction is denoted vT, and generation integrates the learned velocity field from noise toward data.Source: §2 · p. 10
  • Generation starts from Gaussian noise and integrates the ODE from t=1 to t=0; each first-order update uses the current state, teacher velocity, and interval width.Each network call counts as one NFE, and few-step sampling retains only a small number of coarse transitions.Source: §2 · pp. 10–11
  • Classifier-free guidance combines conditional and unconditional teacher predictions; w=1 recovers the conditional prediction, while larger values generally strengthen conditioning at possible cost to diversity.Guidance modifies the local velocity used by the sampler without changing the underlying integration procedure.Source: §2 · p. 10
  • Inference evaluates the velocity field on a finite decreasing time grid, whose placement allocates evaluations across transport regions; few-step sampling makes each coarse transition more consequential.Some models warp nominal time before network evaluation, and the preliminaries absorb that transformation into the sampling grid.Source: §2 · p. 11

3 Method

SMD trains a few-step student by matching its on-policy velocity predictions to a frozen teacher, then refines that regression with a proxy path and reference repulsion. Both refinements are training-only, so inference retains the student’s coarse sampler.

3.1 A Simple On-Policy Baseline

  • On-policy velocity matching substantially improves 4-NFE generation, rapidly improving Qwen-Image samples and improving Wan2.1-T2V-1.3B despite checkpoint fluctuations.The baseline uses teacher velocities queried at intermediate states along the student’s own coarse transition, without explicit endpoint supervision.Source: §3.1 · pp. 11–13

3.2 Imitation from a Proxy Path

  • When teacher trajectories have high curvature, velocity direction can change substantially within a transition, making regression targets potentially conflicting and optimization less stable.With low curvature and moderately varying velocity, local teacher predictions provide more compatible supervision for refining the coarse endpoint.Source: §3.2 · p. 14

3.3 Repulsion from a Reference Prediction

  • Reference repulsion complements imitation of the target teacher prediction by separating the student from a weaker prediction from the same teacher.The reference can be unconditional or conditional under classifier-free guidance, a lower-guidance prediction for embedded-guidance models, or a prediction from a layer-skipped forward pass.Source: §3.3 · p. 15

3.2 Imitation from a Proxy Path

  • Pivot adds a time-varying residual to the student’s coarse velocity, matching local proxy velocities to teacher predictions while preserving the transition endpoints.The zero-integral residual changes intermediate query states and velocities without changing the coarse transition’s endpoint.Source: §3.2 · p. 13, §3.3 · p. 15

4 Few-step Distillation

SMD is evaluated under few-step budgets across image generation and editing, 3D generation, video synthesis, and joint audio–video generation. Across these settings, results include benchmark comparisons and qualitative assessments, with some 3D and reference-conditioned evaluations qualitative only.

4.1 Image Generation

  • At 4-NFE, SMD leads the compared few-step methods on GenEval and DPG-Bench for both image backbones; FLUX.1-dev scores rise from 0.675 to 0.715 on GenEval and from 84.3 to 84.8 on DPG-Bench versus π-Flow.Both FLUX.1-dev scores exceed the corresponding 50-NFE teacher scores; Qwen-Image reaches 0.873 on GenEval and 88.7 on DPG-Bench.Source: §4.1 · p. 16, Table 1 · p. 17, Table 2 · p. 18

4.2 Image Editing

  • On instruction-based editing, SMD raises ImgEdit-Bench overall from Lightning’s 4.48 to 4.51, matching the multi-step teacher, and improves all three GEdit-Bench metrics.GEdit-Bench overall quality rises from 8.14 to 8.19; removal and adjustment show the largest gains over Lightning, which remains stronger on extraction.Source: Table 3 · p. 19, Figure 15 · p. 19, §4.2 · p. 20

4.3 3D Shape and Texture

  • For Hunyuan3D-2.1, SMD preserves recognizable structure where direct 4-step teacher sampling often fragments geometry; TRELLIS.2 distills only its high-resolution shape and texture stages.TRELLIS.2’s preceding sparse-structure and coarse-shape stages remain unchanged, so 4-NFE applies separately to each distilled stage, not the full pipeline.Source: §4.3 · pp. 20–21, Figure 16 · p. 20, Figure 17 · p. 21

4.4 Video Generation

  • At both Wan2.1-T2V scales, SMD improves the Quality aggregate over the multi-step teacher while remaining close in Semantic and competitive with strong 4-NFE baselines.Qualitative 14B comparisons show less pronounced saturation and contrast, with softer lighting transitions than the baselines in the illustrated scenes.Source: §4.4 · pp. 22–23, Table 4 · p. 22, Figure 18 · p. 22

4.5 Joint Audio–Video Generation

  • On H3’s common evaluation set, SMD improves temporal coherence, visual naturalness, audio clarity, and audio quality over LightX2V-DMD at 4- and 8-NFE, while LightX2V-DMD achieves higher prompt alignment.Structural coherence is comparable at 4-NFE and improves substantially with SMD at 8-NFE; the audio difference is particularly pronounced at 4-NFE.Source: §4.5 · p. 23, Table 5 · p. 24
  • Qualitative reference-conditioned H3 results show that a 4-NFE student retains appearance editing, composition, and reference-guided motion and appearance transfer without ground-truth target outputs.Training uses instructions and reference audio–video clips.Source: §4.5 · p. 24

5 Beyond Self-Distillation

The section tests whether on-policy velocity supervision extends beyond self-distillation to cross-model, multi-teacher, and causal-video settings. Results show benefits that depend on the teacher-student pairing, while routed multi-teacher supervision improves on the multi-step base teacher across all six reported metrics.

5 Beyond Self-Distillation

  • The experiments test cross-model transfer, four reward-specialized LoRA teachers routed by prompt category, and on-policy compression of an existing causal video generator.For causal video, the frozen multistep causal teacher supplies the imitation target while a frozen bidirectional teacher supplies global correction; neither is used at inference.Source: §5.1 · p. 25, §5.2 · p. 26, §5.3 · pp. 26, 27, 29

5.1 Cross-Model Distillation

  • A 9B FLUX.2-klein teacher improves a FLUX.2-klein-4B student on six of seven metrics, while Lens-RL-3.8B gains diversity and style but loses on other reported scores.Cross-model velocity supervision is feasible, but gains depend on the teacher-student pairing.Source: §5.1 · p. 25, Table 6 · p. 25, Figure 22 · p. 26

5.2 Multiple Experts as Teachers

  • Routed multi-teacher distillation improves all six reported metrics over self-distillation and exceeds the original 28-NFE base teacher on every reported metric.The routed student also gives the best OCR score, while individual specialists remain strongest on several corresponding criteria.Source: §5.2 · p. 26, Table 7 · p. 27

5.3 Few-Step Causal Video Generation

  • On Wan2.1-T2V-1.3B, the 4-NFE student improves Quality and Total over DMD∗ and comes within 0.09 of Causal Forcing in Total without training a fake-score model.The student uses a correction weight of 0.2 without reference repulsion; neither teacher is used at inference.Source: §5.3 · pp. 26, 29, Table 8 · p. 29
  • With causal-teacher guidance fixed and repulsion disabled, correction weight 0.2 gives the best Semantic score, but the highest Total occurs without correction.Larger correction weights do not improve monotonically, and visual effects vary by prompt.Source: Figure 25 · p. 28, §5.3 · p. 29

6 Ablation Studies

The ablations examine prompt coverage, Pivot design, and teacher repulsion on Qwen-Image and Wan2.1-T2V-1.3B. Outcomes vary by category, backbone, and setting: even single-prompt students respond to unseen prompts, while some larger Pivot or repulsion settings reduce performance.

6.1 Prompt Coverage

  • At 30K prompts, Text-focused training gives the highest Overall among the four categories; raising other categories’ training sets from one to 300 prompts substantially improves their Text scores.The Text-trained student has Overall 0.538 with one prompt versus 0.541 at 30K.Source: §6.1 · pp. 29–30, Table 9 · p. 30, Figure 26 · p. 30
  • Students distilled on one prompt still respond to unseen object and portrait prompts, and single-prompt video students respond to unseen prompts but can retain training-specific visual traits.The evidence is qualitative; the desert-bloom video student retains dot-like patterns.Source: §6.1 · p. 30, Figure 27 · p. 31, Figure 28 · p. 31

6.2 Pivot Design

  • Adding Pivot improves Qwen-Image Overall and Wan2.1-T2V-1.3B Total with either pretrained or random initialization at K = 4, M = 3.The gain does not depend on copying pretrained blocks.Source: §6.2 · p. 31, Table 10 · p. 32
  • For Qwen-Image at M = 7, K = 8 improves all five OneIG dimensions over K = 1; for Wan2.1-T2V-1.3B at M = 3, K = 4 gives the highest Total.For Qwen-Image at K = 1, increasing Pivot blocks from three to seven lowers Overall from 0.540 to 0.534; some larger-K examples also show more similar compositions across seeds.Source: §6.2 · pp. 32–33, Figure 29 · p. 32, Figure 30 · p. 33
  • Piecewise with P = 32 gives the best Qwen-Image Overall, whereas Polynomial with P = 4 gives the best Wan2.1-T2V-1.3B Total.Increasing Piecewise from 4 to 32 elements improves Qwen-Image but not Wan2.1-T2V-1.3B; different bases have relatively similar aggregate scores.Source: §6.2 · p. 34, Table 11 · p. 34

6.3 Teacher Repulsion

  • An unconditional reference with λ = 0.2 gives the best Qwen-Image Overall and Wan2.1-T2V-1.3B Total in the sweeps, while a conditional reference does not give the same benefit.Increasing the unconditional weight to 0.4 reduces Wan2.1-T2V-1.3B Total; moderate repulsion can help, but higher weight can weaken the benefit.Source: §6.3 · pp. 34–35, Table 12 · p. 34, Figure 32 · p. 35

7 Related Work

Related work spans trajectory- and distribution-level distillation, with additional methods that bridge these paradigms or reduce inference cost. The section also situates SMD against regression-based approaches and describes connections between few-step distillation and reinforcement learning.

Trajectory-Level Distillation

  • Trajectory-level methods compress sampling paths: progressive distillation halves step counts over merged intervals [12], while consistency models enforce trajectory-wide agreement and have been extended across schedules, latent spaces, and multi-step settings [46] [47] [48] [49] [50].Flow-map approaches learn transport between time pairs, with MeanFlow enabling one-step generation from a time-averaged velocity identity [60] [61] [62] [63].Source: §7 · p. 36
  • π-Flow and PDD are close regression-based formulations [14] [1], but SMD puts within-interval residual capacity in a training-only branch whose basis integrates to zero.Unlike π-Flow’s substep integration or PDD’s parallel denoising heads, SMD deploys the original student head without Pivot or substep integration.Source: §7 · p. 36

Distribution-Level Distillation

  • Distribution-level distillation compares student-generated distributions with teacher or data distributions using discriminator feedback, score estimates, or sample statistics; SwiftBrush showed pure variational score distillation enables one-step synthesis surpassing the teacher, without adversarial training [96] [97].Alternatives include Fisher divergence, implicit matching, f-divergence weighting, proper scoring rules, EM distillation, and moment matching [107] [108] [109] [110] [112] [113] [114].Source: §7 · pp. 36–37

Other Distillation and Acceleration Techniques

  • Other extensions include flow-map surrogates, video adaptations, adaptive matching, diversity-preserving objectives, caching, and architectural or sampling accelerations [98] [99] [100] [101] [102] [103].The survey also covers self-training and single-sample acceleration [122] [123] [124] [126].Source: §7 · p. 37

Connections to Reinforcement Learning

  • Recent work connects diffusion reinforcement learning to reverse- and forward-process policies, preference rewards, and frameworks that combine rewards with distribution matching [145] [146] [147] [148] [149] [150].Some frameworks also reinterpret distribution matching itself as a reward signal for distilled generators [159] [160] [161] [162] [163].Source: §7 · pp. 37–38

8 Conclusion

The conclusion presents SMD as regression-based few-step distillation that avoids auxiliary score models, adversarial training, and specialized inference architectures. The accompanying appendix passages explain endpoint-preserving Pivot residuals and report that late model blocks are especially influential in the tested swaps.

8 Conclusion

  • SMD combines on-policy velocity regression, a training-only Pivot proxy, and repulsion from weaker teacher predictions while leaving student inference unchanged.Across tasks and models up to 33B parameters, it delivers strong few-step performance; at 4-NFE it outperforms strong few-step baselines and surpasses multi-step teachers on several benchmarks.Source: §8 · p. 38

A.1 Residual Parameterizations

  • Pivot’s zero-integral velocity residual changes intermediate states while preserving both endpoints of the coarse transition.Its integral is a path correction to the straight-line trajectory, with the fixed coarse student velocity determining the endpoint.Source: §A.1 · p. 47
  • Residual parameterizations trade basis structure and coefficient count: cosine modes have zero integral, while centered nodal bases remove their interpolant’s interval mean.Piecewise interpolation yields a piecewise-quadratic path correction; collocation yields a polynomial correction, and centered nodal constructions have P − 1 independent modes.Source: §A.1 · pp. 47–49
  • In the controlled square-wave fit, increasing P sharpens transitions, but finite-dimensional continuous fits retain overshoot and cannot reproduce jumps exactly.The example measures approximation error with the coarse path fixed, not distillation performance.Source: Figure 34 · p. 48, §A.1 · p. 49

A.2 Block-Swap Study

  • Across SANA, Krea-2, and SD3.5-Large Base/Turbo pairs, inserting Turbo’s late blocks into Base gives the highest tested scores, while the reverse replacement gives the lowest.The study uses mean HPSv2 scores on 100 prompts and reports this late-block sensitivity as support for using late student features for Pivot.Source: §A.2 · pp. 49–50, Table 14 · p. 50

A.3 Training Architecture

The appendix compares training computations and documents per-model setups for SMD across image, editing, 3D, and video distillation. The illustrated SMD architecture includes a training-only Pivot branch.

A.3 Training Architecture

  • Figure 36 contrasts SMD’s training-only Pivot branch with distribution matching’s fake-score model and a flow map using a JVP.L denotes backbone depth; Pivot’s residual projection produces P basis coefficients defining rϕ(τ).Source: §A.3 · p. 50, Figure 36 · p. 51

B Few-step Implementation and Evaluation Details

  • Distillation runs use eight NVIDIA A100 GPUs, and the configuration tables specify trainable parameterizations, target and reference settings, and repulsion weights.NFE counts deployed student evaluations; teacher guidance, reference predictions, and Pivot are training-only.Source: B Few-step Implementation and… · p. 51

B.1 Image Generation

  • Image-generation students use on-policy trajectories from random noise, start from teacher initialization, and are selected using GenEval, DPG-Bench, and OneIG-Bench-EN.Training uses text prompts and image dimensions only; reported Overall is the mean of five OneIG-Bench-EN dimensions.Source: §B.1 · p. 51

B.2 Image Editing

  • Image editing trains on source images and instructions through on-policy rollouts without using target edited images.Evaluation uses the official ImgEdit-Bench and GEdit-Bench-EN procedures, with GPT-4o and GPT-4.1 respectively.Source: §B.2 · pp. 51–52

B.3 3D Shape and Texture

  • Hunyuan3D-2.1 uses one rendered view per object to condition shape-latent rollouts, while TRELLIS.2 distills its high-resolution shape and texture stages.TRELLIS.2 retains the preceding sparse-structure and coarse-shape stages; ground-truth shape latents are not used for Hunyuan3D rollouts, nor are ground-truth sparse features used to construct TRELLIS.2 rollout states.Source: §B.3 · p. 52, Table 16 · p. 52

B.4 Video Generation

  • Wan2.1 video students are trained on-policy, with the 1.3B model using an unconditional reference and the 14B model using the conditional teacher prediction.Evaluation follows the Self-Forcing protocol; the reproduced DMD2 baselines use synthetic latents, LoRA students and fake-score models, and combined VSD and GAN losses.Source: §B.4 · p. 53

B.5 Detailed VBench Results

The appendix details VBench’s per-dimension video results and evaluation measures, then describes joint audio–video training and its evaluation setup. The reported audio and visual scores use distinct scales and definitions.

B.5 Detailed VBench Results

  • Table 17 reports all 16 VBench dimensions for Wan2.1-T2V at 4-NFE, with Quality and Semantic aggregates using standard VBench normalization and weights.Dynamic Degree and Motion Smoothness distinguish motion activity from temporal smoothness; the dimensions do not directly measure repeated-generation diversity.Source: §B.5 · p. 53, Table 17 · p. 54

B.6 Joint Audio–Video Generation

  • MiniMax H3 joint audio–video training uses text prompts without real target videos or first/last-frame inputs, while reference-conditioned training uses references as conditions rather than targets.The model has guidance distilled into it, so student inference requires no additional CFG branch.Source: §B.6 · p. 54

B.7 H3 Evaluation

  • H3 evaluation matches prompts, seeds, resolution, and duration between inference budgets, and reports four visual dimensions on a 1–5 scale and two audio dimensions on a 0–5 scale.Figure 19 displays full-clip native 32 kHz stereo spectrograms over 0–16 kHz on a shared −90 to 0 dB scale.Source: §B.7 · p. 55

C Beyond Self-Distillation Implementation and Evaluation Details

The appendix records implementation and evaluation details for cross-model, multi-teacher, and causal-video distillation. The setups vary their prompt pools, teacher assignments, and causal correction signals.

C Beyond Self-Distillation Implementation and Evaluation Details

C.1 Cross-Model Distillation

  • Cross-model distillation uses the π-Flow prompt set, while image benchmarks follow the procedures used for image generation.Lens self-distillation retains its native CFG norm rescaling.Source: §C.1 · p. 56

C.2 Multiple Experts as Teachers

  • Multi-teacher training draws equally from four prompt categories and uses the corresponding specialist teacher for each batch.The 120K-entry pool combines four 30K-prompt categories; each specialist uses one fixed reward teacher.Source: §C.2 · pp. 56–57

C.3 Few-Step Causal Video Generation

  • Causal video distillation uses a causal teacher for imitation and a frozen bidirectional Wan2.1-T2V-14B teacher for correction.The correction sweep varies λcorr across {0, 0.2, 0.4, 0.8}, with reference repulsion disabled.Source: §C.3 · p. 57, Table 20 · p. 57

D Ablation Implementation and Evaluation Details

The appendix specifies settings for Pivot and teacher-repulsion ablations and the evaluation prompts used for video comparisons. The visual assessment prioritizes persistent, material quality differences over isolated defects or stylistic differences.

D Ablation Implementation and Evaluation Details

  • Wan2.1-T2V-1.3B ablations use four-step Euler sampling, two queries per transition, and one evaluation run per setting.The default settings are listed in Table 21; qualitative cells show one sample for single-prompt video and four otherwise.Source: D Ablation Implementation and… · p. 57, Table 21 · p. 57

E Visual Evaluation System Prompt

  • The visual evaluation compares matched chronological frames and breaks equivalence only for persistent quality advantages or major viewing-impairing defects.It assesses prompt alignment, structural coherence, temporal coherence, and visual naturalness independently, and directs evaluators not to force a winner.Source: E Visual Evaluation System Prompt · pp. 58–59

Related papers

Loading 2610.18472v1…