PaperlayerView exact revision on arXiv

Source-linked AI summary

From Structure to Detail: Hierarchical Distillation for Efficient Diffusion Model

Hanbo Cheng, Peng Wang, Kaixiang Lei, Qi Li, Zhen Zou, Pengfei Hu, Jun Du

arXiv:2511.08930v1cs.CVcs.AI

TL;DR

Diffusion models’ iterative denoising creates high inference latency, while single-step trajectory and distribution distillation face complementary limitations. The paper combines them in a hierarchical pipeline, using trajectory distillation for structural initialization followed by distribution matching and adversarial refinement. It reports that the single-step model rivals its 250-step teacher on ImageNet.

  • Problem

    Conventional diffusion inference often requires 50 to 100 denoising steps, while single-step trajectory distillation loses fine details and distribution-based methods can collapse to a few modes.

    Source: §1 · pp. 1–2

  • Method

    Hierarchical Distillation first uses trajectory distillation to provide a structural initialization, then refines the student through distribution matching and adversarial training.

    Source: §1 · p. 2

  • Results

    FID 2.26 on ImageNet 256×256 lets the single-step method rival the 250-step teacher.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The authors present HD as a state-of-the-art approach to single-step generation, with competitive text-to-image benchmark results.

    Source: §1 · p. 2

    The current design is a two-stage process; the authors identify unified end-to-end training as future work.

    Source: Limitation and Future work · p. 13

Abstract

from arXiv · show

The inference latency of diffusion models remains a critical barrier to their real-time application. While trajectory-based and distribution-based step distillation methods offer solutions, they present a fundamental trade-off. Trajectory-based methods preserve global structure but act as a "lossy compressor", sacrificing high-frequency details. Conversely, distribution-based methods can achieve higher fidelity but often suffer from mode collapse and unstable training. This paper recasts them from independent paradigms into synergistic components within our novel Hierarchical Distillation (HD) framework. We leverage trajectory distillation not as a final generator, but to establish a structural ``sketch", providing a near-optimal initialization for the subsequent distribution-based refinement stage. This strategy yields an ideal initial distribution that enhances the ceiling of overall performance. To further improve quality, we introduce and refine the adversarial training process. We find standard discriminator structures are ineffective at refining an already high-quality generator. To overcome this, we introduce the Adaptive Weighted Discriminator (AWD), tailored for the HD pipeline. By dynamically allocating token weights, AWD focuses on local imperfections, enabling efficient detail refinement. Our approach demonstrates state-of-the-art performance across diverse tasks. On ImageNet $256\times256$, our single-step model achieves an FID of 2.26, rivaling its 250-step teacher. It also achieves promising results on the high-resolution text-to-image MJHQ benchmark, proving its generalizability. Our method establishes a robust new paradigm for high-fidelity, single-step diffusion models.

1. Introduction

Diffusion inference latency motivates single-step distillation, but trajectory and distribution methods face complementary fidelity and stability limits. The section frames HD as a staged combination that uses trajectory distillation for initialization before distribution matching and adversarial refinement.

1. Introduction

  • Trajectory distillation preserves global structure but loses fine detail, whereas distribution-based methods can collapse to a few high-density modes without a strong initial distribution.The paper describes trajectory distillation as lossy information compression and attributes distribution-based collapse to reverse-KL zero-forcing.Source: §1 · p. 2
  • HD uses trajectory distillation to create a structurally sound sketch, then refines that initialization through distribution matching and adversarial training.The paper says the initialization provides a well-posed starting distribution for refinement.Source: §1 · p. 2
  • FID 2.26 on ImageNet 256×256 lets the single-step method rival the 250-step teacher, while the paper reports competitive text-to-image results.The contribution summary also characterizes the method as state of the art in single-step generation.Source: §1 · p. 2

2. Related Works

  • Prior trajectory methods pursue teacher-trajectory fidelity, while distribution matching aligns generated and real distributions but faces mode collapse and initialization sensitivity.Earlier trajectory approaches sought improved trajectory approximation; the paper contrasts them with direct distribution alignment.Source: §2 · pp. 2–3

3.1. Flow Matching with Mean Velocity

  • MeanFlow models interval-average velocity so an accurate full-trajectory prediction reduces inference to one function evaluation.The interval is [r, t], and setting it to [0, 1] yields the single-step path.Source: §3.1 · p. 3

3.2. Distribution Matching Distillation

  • DMD minimizes KL divergence between generated and real distributions using their score fields, approximated by a frozen teacher and a separately trained fake network.The DMD loss uses a stop-gradient implementation for equivalent gradient optimization.Source: §3.2 · pp. 3–4

4. Method

The method organizes trajectory distillation and distribution matching hierarchically: MeanFlow-based TD initializes structure, then distribution refinement restores detail. Its analysis treats mainstream TD objectives as mean-velocity estimation, motivating this division of labor.

4.1. A Unified View of Trajectory Distillation

  • Mainstream TD methods reduce to modeling dynamic mean velocity, while finite-capacity single-step students face a bottleneck that loses fine-grained detail.The authors select MeanFlow as a representative because it explicitly models mean velocity.Source: §4.1 · pp. 4–5

4.2. Stage 1: Structured Initialization via TD

  • MeanFlow-based TD uses pretrained-teacher velocity to give the student a structural prior and a well-posed initialization for refinement.The teacher's fixed, deterministic mapping provides a lower-variance learning signal than training MeanFlow from scratch.Source: §4.2 · pp. 5–6

4.3. Stage 2: Distribution Refinement

  • The final HD student generates minimal-step images matching the multi-step teacher's quality while preserving diversity.DMD refines the initialized distribution toward real data; feature-space adversarial training and AWD's token weighting target localized artifacts.Source: §4.3 · p. 6

5. Experiment

The experiments examine Trajectory Distillation’s capacity bottleneck, compare Hierarchical Distillation with prior methods, and ablate its stages. Results report single-step benchmark performance and show benefits from both TD initialization and the proposed discriminator.

5.1. The Analysis of Bottleneck of Trajectory Distillation

  • A 2D mapping experiment distills a flow-matching teacher from a rectangular prior to an upper semicircle, evaluating generated-point distance to the unit circle.Lower Euclidean distance indicates better performance.Source: §5.1 · p. 6
  • Even with over 50× model capacity, MeanFlow TD does not perfectly replicate the teacher trajectory; student performance improves with capacity while the multi-step teacher changes marginally.Figure 3 states that TD’s upper bound increases with model parameters.Source: §5.1 · p. 7, Figure 3 · p. 7

5.2. Comparison Study

  • FID 2.26 in one-step ImageNet 256×256 generation is comparable to the 250-step teacher; on MJHQ, the single-step model approaches the 50-step teacher and outperforms step-distillation methods.The authors also report that the method preserves global structure and fine-grained details across these tasks.Source: §5.2 · p. 7, Table 1 · p. 7, Table 2 · p. 8

5.3. Ablation Study

  • TD initialization provides a significant advantage, and the proposed Attention-Weighted Discriminator boosts performance; the standard adversarial setup is suboptimal within HD.The ablations compare models with and without TD initialization while keeping subsequent DMD and adversarial distillation steps identical.Source: §5.3 · p. 8, Table 3 · p. 8

6. Conclusion

The paper unifies Trajectory Distillation (TD) under a mean-velocity objective, then uses Hierarchical Distillation (HD) to combine TD’s structural priors with distribution-matching refinement and an Adaptive Weighted Discriminator (AWD). Its experiments report single-step fidelity comparable to a multi-step teacher.

6. Conclusion

  • HD uses TD to initialize structural priors, then refines the model through distribution matching with an Adaptive Weighted Discriminator (AWD) focused on local artifacts.The paper presents this pipeline as a way to refine details while retaining the structure supplied by TD.Source: §6 · p. 8
  • The unified analysis identifies a shared mean-velocity objective for TD and an information-theoretic bottleneck linked to difficulty preserving fine-grained details.The paper reports that the resulting single-step student achieves fidelity comparable to its multi-step teacher.Source: §6 · p. 8

Equivalence of Mean Velocity Estimation and Trajectory Distillation

Under ideal fitting, the analysis shows that continuous consistency models and progressive distillation can both be understood as estimating mean velocity over a time interval. The derivations specify the interval each method averages over.

Equivalence of Mean Velocity Estimation and Trajectory Distillation

  • The analysis frames mainstream trajectory-distillation methods as mean-velocity estimation, including continuous consistency models and progressive distillation.The proof discusses these methods under the assumption that models can perfectly fit their training objectives.Source: Equivalence of Mean Velocity… · p. 11
  • Continuous consistency models implicitly estimate the mean instantaneous velocity over [0, t].The paper derives this from the model’s consistency constraint.Source: Equivalence of Mean Velocity… · p. 11
  • As progressive distillation rounds increase, the student converges to modeling mean velocity over the full interval [0, T].Each round trains the student to match the average of two consecutive teacher steps; the derivation assumes ideal fitting.Source: Equivalence of Mean Velocity… · pp. 11–12

Details of the Experiment in Section 5.1

The Section 5.1 experiments use residual-MLP models whose capacity is scaled by width or depth. Student performance improves with capacity, but even much larger students do not perfectly reproduce the teacher’s trajectory.

Details of the Experiment in Section 5.1

  • The Section 5.1 model stacks residual-connected MLPs and scales capacity by increasing hidden-state dimensionality or adding layers.The model architecture and prior-versus-target distributions are shown in Figure 5; Table 4 contains model-size hyperparameter configurations.Source: Details of the Experiment in Section 5.1 · p. 12, Figure 5 · p. 13, Table 4 · p. 13
  • Scaling student capacity improves performance, but even models tens of times larger than the smallest configuration do not perfectly replicate the teacher’s trajectory.On the relatively simple task described, the multi-step teacher’s performance has already saturated as its parameter count increases.Source: Details of the Experiment in Section 5.1 · p. 12

Implementation Details

Implementation uses separate two-stage training setups for class-conditional and text-to-image generation, with task-specific data, schedules, and guidance settings.

Class Conditional Generation

  • Class Conditional Generation trains for 20 epochs in Stage 1 and fine-tunes for 10 epochs in Stage 2, using CFG 1.75 across timesteps.Stage 1 uses Adam with batch size 256 and learning rate 1 × 10^-4; Stage 2 uses a generator learning rate of 1 × 10^-5 and fake-branch and discriminator rates of 5 × 10^-5.Source: Class Conditional Generation · p. 13

Text-to-Image Generation

  • Text-to-Image Generation trains on CC3M for 25,000 Stage 1 iterations and fine-tunes for 8,000 Stage 2 iterations, with CFG 4.5.Stage 1 uses CAME, learning rate 1 × 10^-5, and batch size 256; the one-step setting uses maximum timestep T = 1.0, except SANA-sprint at T = 1.5708 (π/2).Source: Text-to-Image Generation · p. 13

Limitation and Future work

The paper reports that its single-step generation quality rivals the 250-step teacher, while identifying future work on unifying its two-stage design, exploring non-adversarial alignment, and testing broader scale and domains.

Limitation and Future work

Evaluation with FLOPs

On ImageNet at 256×256, the method reduces inference cost from 29.98 GFLOPs to 0.43 GFLOPs while retaining performance comparable to the 250-step teacher.

Evaluation with FLOPs

  • 29.98 GFLOPs to 0.43 GFLOPs, a 69.72× reduction: the single-step student retains performance comparable to the 250-step teacher on ImageNet 256×256.Source: Evaluation with FLOPs · p. 14, Table 5 · p. 14
  • The teacher's 250-step inference procedure is compressed into one student step while preserving performance comparable to the multi-step baseline.Source: Evaluation with FLOPs · p. 14
  • The evaluation reports FLOPs and parameter counts for SiT-XL/2 at 256×256 input resolution with batch size 1.Source: Table 5 · p. 14

Qualitative Results on ImageNet 256 × 256

The section presents qualitative results comparing the one-step HD student with the 250-step SiT teacher on ImageNet at 256×256.

Qualitative Results on ImageNet 256 × 256

Related papers

Loading 2511.08930v1…