Learning Few-Step Diffusion Models by Trajectory Distribution Matching

Figure 1

• PixArt-α images use 50 NFE, while TDM images use 4 NFE.The TDM generator was distilled from PixArt-α with 500 training iterations and 2 A800 hours.Source: Figure 1 · p. 1
• All images in the comparison are generated from the same initial noise.Source: Figure 1 · p. 1

Author notes

• The TDM image locations are listed in a footnote, and the work was partly completed during an internship at Huawei Noah’s Ark Lab.The note also identifies a corresponding author.Source: §1 · p. 1

Motivation

• Distribution matching methods are mainly designed for one-step generation, while trajectory distillation struggles to match instance-level trajectories in few steps and performs suboptimally.Few-step generation is presented as a better speed–quality balance for complex text-to-image tasks.Source: §1 · p. 2

TDM

• TDM aligns the student trajectory with the teacher at the distribution level through a data-free score distillation objective.The approach combines trajectory and distribution matching.Source: §1 · p. 2
• A sampling-steps-aware objective decouples learning targets across sampling steps, enabling flexible deterministic sampling.The introduction also describes a surrogate objective using the Pseudo-Huber metric.Source: §1 · p. 2

Motivation

• Existing deterministic methods offer less-adjustable step counts, whereas adjustable stochastic methods sacrifice image quality for flexibility.Source: §1 · p. 2

Results

• TDM reports state-of-the-art results across text-to-image backbones and improves CogVideoX-2B’s VBench total score from 80.91 to 81.65 with 4 NFE.For PixArt-α, the 4-step generator outperforms its teacher in HPS, AeS, and user preference after 500 iterations and 2 A800 hours.Source: §1 · p. 2

Diffusion Models

• A diffusion process adds Gaussian noise to data over T steps, with the noise level specified by a schedule.Noisy samples are formed as x_t = x + σ_tϵ, where ϵ is sampled from a standard normal distribution.Source: §2 · pp. 2–3
• Diffusion networks are trained to denoise noisy samples, and sampling typically solves a probability-flow ODE in at least 25 steps.Source: §2 · p. 3

Diffusion Distillation

• Trajectory distillation simulates teacher ODE trajectories at the instance level, while distribution matching transfers knowledge through a distribution divergence.Distribution matching mostly ignores intermediate trajectory steps, making it harder to extract extra information when increasing sampling steps.Source: §2 · pp. 3–4

3.1. Trajectory Distribution Matching

• Point-to-point trajectory matching is difficult with limited model capacity and requires costly multi-step teacher ODE solving that can introduce numerical error.TDM instead aligns corresponding trajectory distributions, with the stated aim of inheriting teacher knowledge while reducing learning difficulty.Source: §3.1 · p. 4
• TDM aligns student-trajectory distributions with the teacher at corresponding timesteps, avoiding teacher ODE-solving errors while training from student samples without data.The method uses a marginal reverse-KL objective and treats the student as a discrete K-step generator.Source: §3.1 · p. 4
• The data-free objective has a performance ceiling set by the pretrained diffusion teacher, whose samples define the method’s upper bound.The authors suggest fine-tuning SD-v1.5 on a compact high-quality dataset before distillation to address its mismatch with human preferences.Source: §3.1 · p. 5

3.2. Enabling Flexible Sampling Steps

• Sampling-steps-aware training targets multiple desired step counts, while conditioning student and fake scores on the count addresses bias from sharing fake scores.The method is intended to support flexible deterministic sampling; shared fake scores can introduce training bias even when optimal.Source: §3.2 · p. 5

3.3. Surrogate Training Objective

• A Pseudo-Huber surrogate compares generated samples with revised samples, and its gradient normalization may stabilize training.The surrogate connects TDM with consistency-model training, while TDM obtains revised samples differently; the method also distinguishes its score-distillation approach from DMD2 and MMD.Source: §3.3 · p. 6

4. Experiment

• The image-distillation experiments use JourneyDB prompts without image data and evaluate image quality, image-text alignment, and human preference.The setup also compares incorporating high-quality data through pre-distillation fine-tuning or GANs during distillation.Source: §4 · p. 6

4.1. Main Results

• TDM-unify reports superior performance across architectural variants and evaluation metrics, with deterministic sampling providing state-of-the-art results at multiple step counts.The comparison covers text-to-image generation methods, including TDM-unify-SFT initialized from fine-tuned SD-v1.5 and TDM-unify-GAN initialized from original SD-v1.5.Source: §4.1 · p. 6, Table 1 · p. 7
• TDM-LoRA significantly outperforms existing approaches on unseen customized models while maintaining high-fidelity generation and style preservation.Style preservation is assessed by FID between teacher and student samples.Source: Table 2 · p. 7, §4.1 · p. 7
• 1.25% and 0.57% of DMD2 and YOSO training costs, respectively, are required by TDM on SDXL and PixArt-α; LCM deterministic sampling is recovered in 100 iterations.TDM’s cost comparison includes fine-tuning for TDM-unify-SFT; the LCM result is reported for 4-step generation.Source: Table 3 · p. 8, §4.1 · p. 9, Figure 8 · p. 9

4.2. Additional Results on Text-to-Video Acceration

• With 4 NFE, TDM surpasses CogVideoX-2B by a notable margin on VBench in text-to-video evaluation.The paper reports that the extension benefits from TDM’s quick convergence and data-free property.Source: Table 4 · p. 9, §4.2 · p. 9

4.3. Ablation Study

• 20.81 vs. 28.90 HPS at 1 step and 29.08 vs. 30.52 HPS at 2 steps show the reported benefit of unify training; importance sampling raises 4-step HPS from 29.27 to 31.35.Without importance sampling, the paper also reports degraded mode coverage and training stability.Source: §4.3 · p. 9
• 26.11 vs. 28.90 at 1 step and 29.39 vs. 31.31 at 4 steps show the reported benefit of sampling-steps-aware training; the surrogate objective yields 31.31 vs. 30.85 HPS.The authors attribute the step-conditioning result as likely due to distribution interference across steps when sharing the fake score.Source: §4.3 · p. 9

5. Conclusion

• TDM unifies trajectory distillation and distribution matching in a new distillation paradigm.The conclusion presents this as a data-free approach.Source: §5 · p. 10
• With two A800 hours, TDM distills a four-step generator that outperforms PixArt-α in quantitative and human evaluations.The conclusion describes this as extremely low training cost.Source: §5 · p. 10
• The authors report that TDM achieves state-of-the-art performance with fewer steps and extremely low training costs.They say this efficiency may make future diffusion-distillation research more inclusive.Source: §5 · p. 10

A. Training Algorithm

• Algorithm 1 samples noise, generates a K-step ODE trajectory, samples timesteps and noisy inputs, then updates the fake score and generator.The loop repeats for N iterations.Source: A. Training Algorithm · p. 12

B. Derivations

• The derivation considers fake-score learning by score matching without noise and identifies the optimal fake score through the zero-gradient condition.The stated derivation proceeds from the score-matching objective to its gradient and optimum.Source: B. Derivations · p. 12

C. Additional Related Works

• Trajectory-distillation methods simulate diffusion trajectories, while distribution matching targets distributions; TDM matches trajectories at the distribution level.The passage notes trajectory methods' numerical errors solving pretrained PF-ODEs and that score-distillation methods typically ignore intermediate trajectory steps.Source: C. Additional Related Works · p. 12
• Unlike Hyper-SD, TDM unifies trajectory distillation and distribution matching in one objective rather than using separate objectives and training stages.The passage also distinguishes TDM from DMD2, which predicts clean images at different timesteps without aligning to the teacher trajectory.Source: C. Additional Related Works · p. 12
• Unlike MMD, TDM uses reverse KL for the generator and score matching for the fake score, with deterministic sampling from noise rather than ancestral sampling from noisy real data.The passage characterizes MMD's stochastic intermediate samples as suboptimal for few-step sampling.Source: C. Additional Related Works · p. 13
• TDM trains at only K timesteps with distribution-level matching, supporting flexible deterministic sampling in contrast to a concurrent method's point-to-point self-consistency.The passage says that method required arbitrary-timestep training and that its text-to-image extension remained unclear.Source: C. Additional Related Works · p. 13

D. Experiment details

• In the described setup, both the generator and fake score use the AdamW optimizer; by default, β1 is 0 and β2 is 0.999.Source: D. Experiment details · p. 13
• SD-v1.5 uses batch size 256, CFG 3.5, and gradient clipping at 1.0, with up to 20,000 unified-training or 3,000 specific four-step iterations.The generator and fake-score learning rates are 2e-6 and 2e-5, respectively.Source: D. Experiment details · p. 13
• SDXL uses batch size 64 and CFG 8, with 1,000 fine-tuning iterations at 512 resolution followed by 1,000 at 1024 resolution.The fake score is initialized from pretrained SDXL in both stages; gradient clipping is 1.0.Source: D. Experiment details · p. 13
• PixArt-α uses batch size 32, CFG 3.5, gradient clipping at 1.0, and training within 500 iterations.The generator and fake-score learning rates are 2e-6 and 2e-5, respectively.Source: D. Experiment details · p. 13

E. Additional Experiments

• TDM-unify shows systematic behavior changes across conditioning-step and sampling-step combinations, positioning itself along the ODE trajectory according to the specified condition steps.Figure 10 visualizes varying condition and sampling steps using the prompt “A corgi with sunglasses, traveling in the sea.”Source: E. Additional Experiments · p. 13, Figure 10 · p. 13

F.1. Ablation Details

• The ablation variants use the same hyperparameters and training iterations, differing only in the components being ablated.Source: §F.1 · p. 13
• The GAN ablation uses latent discriminators based on the SD-v1.5 UNet encoder, and injects the desired sampling steps K into the discriminator.ADVLoss is identified as the adversarial loss.Source: §F.1 · p. 13
• Tables 6 and 7 compare SD-v1.5 variants on HPS for four-step generation, while Table 8 addresses the effect of Fisher Divergence.The supplied captions do not report table values or comparison outcomes.Source: Table 6 · p. 14, Table 7 · p. 14, Table 8 · p. 14
• The original loss uses the normalization proposed in DMD [41], whereas the proposed surrogate loss does not use that normalization.Source: §F.1 · p. 14
• The DMD2 implementation updates the fake score 10 times per iteration, with other hyperparameters kept consistent with the authors' configuration.Source: §F.1 · p. 14

F.2. Additional Ablation

• Importance sampling improves mode coverage and image quality; without it, the fake score may not track the student distribution accurately.DMD2 also shows mode collapse, though less severely; its repeated fake-score training is noted as a possible reason and increases training time.Source: §F.2 · p. 14, Figure 11 · p. 14
• DDIM and DPM training yield similar performance, with DPM sampling available at test time for either training solver.The authors attribute the similarity partly to backpropagating through only one ODE step, which prevents higher-order solvers from using higher-order training information.Source: §F.2 · p. 14
• Matching noisy rather than clean samples supports deterministic sampling; clean-sample matching produces poor results with notable artifacts.The comparison is shown for four-step generation on SD-v1.5.Source: Figure 12 · p. 14, §F.2 · p. 15
• TDM can use divergences other than reverse KL, and the more expensive Fisher divergence further improves performance.Source: §F.2 · p. 15

G. User Study Details

• The user study used 20 randomly selected prompts and approximately 40 responses to compare anonymous images for quality and prompt alignment.Images were manually checked for inappropriate or dangerous content.Source: G. User Study Details · p. 15

• Figure 16 visualizes ODE trajectories using clean samples at different timesteps.Source: H. Additional Qualitative Results · p. 15
• On unseen customized models, the method shows better visual quality and style preservation than the competing baseline.The additional samples cover dreamshaper and realisticvision.Source: H. Additional Qualitative Results · p. 15, Figure 14 · p. 16, Figure 15 · p. 17
• The ODE-trajectory comparison reports less CFG artifact and better visual quality for the method.Source: Figure 16 · p. 17