PaperlayerView exact revision on arXiv

Source-linked AI summary

One-step Diffusion Models with $f$-Divergence Distribution Matching

Yilun Xu, Weili Nie, Arash Vahdat

arXiv:2502.15681v2cs.LGcs.AIcs.CV

TL;DR

Diffusion sampling remains slow because it iteratively calls denoising networks, while reverse-KL distribution matching can miss teacher modes. f-distill generalizes distribution matching to f-divergences with a density-ratio-dependent weighting of teacher–student score differences. Less mode-seeking divergences outperform previous variational score distillation methods across image tasks, and Jensen-Shannon reaches state-of-the-art one-step results on ImageNet-64 and zero-shot MS-COCO.

  • Problem

    Diffusion models require iterative denoising, and existing reverse-KL distribution matching is mode-seeking and can ignore diverse modes learned by the teacher.

    Source: §1 · p. 1

  • Method

    f-distill minimizes an f-divergence between teacher and student distributions using a gradient that weights their score difference by a function of the density ratio and chosen divergence.

    Source: §1 · p. 2, §3 · p. 4

  • Results

    Less mode-seeking f-divergences consistently outperform previous variational score distillation approaches, while Jensen-Shannon achieves state-of-the-art one-step generation on ImageNet-64 and zero-shot MS-COCO.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The authors conclude that f-distill demonstrates strong one-step generation capabilities across various image-generation tasks.

    Source: §7 · p. 8

Abstract

from arXiv · show

Sampling from diffusion models involves a slow iterative process that hinders their practical deployment, especially for interactive applications. To accelerate generation speed, recent approaches distill a multi-step diffusion model into a single-step student generator via variational score distillation, which matches the distribution of samples generated by the student to the teacher's distribution. However, these approaches use the reverse Kullback-Leibler (KL) divergence for distribution matching which is known to be mode seeking. In this paper, we generalize the distribution matching approach using a novel $f$-divergence minimization framework, termed $f$-distill, that covers different divergences with different trade-offs in terms of mode coverage and training variance. We derive the gradient of the $f$-divergence between the teacher and student distributions and show that it is expressed as the product of their score differences and a weighting function determined by their density ratio. This weighting function naturally emphasizes samples with higher density in the teacher distribution, when using a less mode-seeking divergence. We observe that the popular variational score distillation approach using the reverse-KL divergence is a special case within our framework. Empirically, we demonstrate that alternative $f$-divergences, such as forward-KL and Jensen-Shannon divergences, outperform the current best variational score distillation methods across image generation tasks. In particular, when using Jensen-Shannon divergence, $f$-distill achieves current state-of-the-art one-step generation performance on ImageNet64 and zero-shot text-to-image generation on MS-COCO. Project page: https://research.nvidia.com/labs/genair/f-distill

1. Introduction

Diffusion sampling is slow, and reverse-KL distribution matching can ignore diverse teacher modes. The paper introduces f-distill, generalizing distribution matching to f-divergences, and reports state-of-the-art one-step results with Jensen-Shannon divergence on ImageNet-64 and zero-shot MS-COCO.

1. Introduction

  • Distribution matching distills a one-step student by matching its sample distribution to that imposed by a pretrained teacher, rather than learning the deterministic noise-to-data mapping.The paper motivates this approach by describing that mapping as complex and hard to learn.Source: §1 · pp. 1–2
  • Reverse-KL is commonly used in distribution matching but is mode-seeking and can potentially ignore diverse modes learned by the teacher.Applying arbitrary f-divergences to diffusion distillation and practically estimating the student gradient remain open challenges.Source: §1 · p. 2
  • f-distill expresses the distribution-matching gradient as a teacher–student score difference multiplied by a density-ratio- and divergence-dependent weight.The density ratio can be obtained from the auxiliary GAN objective; the framework includes previous DMD as a constant-weight special case.Source: §1 · p. 2
  • Forward-KL offers better mode coverage but has large gradient variance, while Jensen-Shannon has moderate mode-seeking and lower variance.Jensen-Shannon also exhibits gradient saturation, particularly early in training.Source: §1 · p. 2
  • Jensen-Shannon f-distill achieves state-of-the-art one-step generation performance on ImageNet-64 and zero-shot MS-COCO.The paper also reports state-of-the-art FID scores on both benchmarks.Source: §1 · p. 2

2. Background

Diffusion models generate samples by iteratively denoising noisy data, while variational score distillation trains a one-step generator to match a teacher distribution. The background introduces the score-based update and its f-divergence context.

2.1. Diffusion models

  • Diffusion models iteratively denoise Gaussian noise through a probability-flow ODE, optionally with a Langevin SDE, and samplers typically still require more than 20 steps for decent samples.The ODE guides samples from high to low noise; the SDE acts as an equilibrium sampler across noise levels.Source: §2.1 · pp. 2–3
  • The f-distill update multiplies the teacher–fake score difference by a weight determined by the divergence and density ratio, downweighting low-density regions with larger score differences.The figure caption notes that the density ratio is available in the auxiliary GAN objective.Source: Figure 2 · p. 3, Figure 3 · p. 3

2.2. Variational score distillation

  • Variational score distillation trains a one-step generator to map prior noise directly to a clean sample, bypassing iterative teacher sampling.Its generator update matches teacher and student distributions by minimizing reverse-KL, using the teacher score minus the student-distribution score.Source: §2.2 · p. 3
  • The student score is estimated by a fake-score network trained on generator samples, with training alternating between generator and score updates.Prior work also adds a GAN loss using a lightweight classifier on fake-score-network middle features.Source: §2.2 · p. 3
  • The f-divergence family includes forward-KL, reverse-KL, Hellinger, and Jensen-Shannon divergences, represented by different choices of convex function.The function is constrained to be convex on (0, +∞) and satisfy f(1) = 0.Source: §2.3 · pp. 3–4

3. Method: general f-divergence minimization

f-distill generalizes distribution matching to arbitrary f-divergences by deriving a density-ratio-weighted score-difference gradient. The framework includes reverse-KL variational score distillation as a special case and supports practical estimation of its terms.

3. Method: general f-divergence minimization

  • The gradient expression is analytical for teacher and student distributions perturbed through the diffusion forward process, with x formed from the generator output plus Gaussian noise.The theorem assumes a twice continuously differentiable f and a differentiable generator mapping.Source: §3 · p. 4
  • The f-divergence gradient is a teacher–student score difference weighted by h(r), where the weight depends on the density ratio and chosen divergence.The framework defines h(r) as f''(r)r^2 and recovers reverse-KL variational score distillation when h(r) is constant at 1.Source: §3 · p. 4
  • The framework can use any continuous, non-negative weighting function h, although the paper’s study is limited to canonical f-divergences.Proposition 1 connects such an h multiplied by the score difference to the gradient of an f-divergence.Source: §3 · p. 5
  • f-distill minimizes the divergence across the diffusion time range, using time-dependent weights to equalize gradient magnitudes across times.The student score is approximated by an online diffusion model, and an auxiliary GAN discriminator estimates the density ratio used by the weighting function.Source: §3 · p. 5

4. Comparing properties of f-divergence

The paper compares f-divergences by mode-seeking behavior, gradient saturation, and training variance. These properties involve trade-offs, and a two-stage normalization scheme is proposed to reduce variance for less mode-seeking divergences.

4. Comparing properties of f-divergence

  • Forward-KL is mode-covering, while reverse-KL and JS are mode-seeking; JS is less mode-seeking than reverse-KL.Less mode-seeking divergences tend to downweight samples in low-density teacher regions through their weighting function h(r).Source: §4 · p. 5
  • Squared Hellinger and JS have smaller gradients at extreme density ratios, a saturation concern when teacher and student distributions are poorly aligned early in training.Diffusion distillation can mitigate saturation by initializing the student from a pretrained diffusion model.Source: §4 · pp. 5–6
  • As the mean gap between unit-variance Gaussians grows, forward-KL and Jeffreys variance increases significantly, while JS and squared Hellinger remain relatively stable.The comparison uses normalized variance of h(p/q) under q.Source: §4 · p. 6, Figure 5 · p. 6
  • Two-stage normalization reduces forward-KL training-loss variance on ImageNet-64 by normalizing density ratios within time bins and h within each minibatch.The scheme maintains the f-distill objective’s relative importance against the GAN objective.Source: §4 · p. 6

5. Experiment

The experiments evaluate f-distill across class-conditioned and text-to-image generation, comparing divergences and existing methods. They report strong one-step results and examine how the weighting function relates to teacher–student score differences.

5.1. Image generation

  • The evaluation covers CIFAR-10 and ImageNet-64 class-conditioned generation and zero-shot MS COCO text-to-image generation, using EDM and Stable Diffusion teachers.COYO-700M is the text-to-image training set; the reported comparison baselines include multi-step diffusion, distillation methods, and a GAN-only ablation.Source: §5.1 · pp. 6–7
  • 1.16 FID on ImageNet-64 and 7.42 FID on MS COCO are reported for JS; it beats most baselines and approaches eDiff-I at 6.95 FID on MS COCO.The ImageNet-64 comparison excludes GDD-I; the MS COCO result uses CFG=1.75.Source: Table 4 · p. 7, §5.1 · p. 7
  • On CIFAR-10, all f-distill variants outperform adversarial distillation, while forward-KL and Jeffreys generally yield lower FID and higher Recall.The comparisons use FID for sample quality and Recall for diversity.Source: §5.1 · p. 7
  • Forward-KL and JS achieve better FID scores than reverse-KL, while JS is more stable than forward-KL on ImageNet-64.At CFG=5, JS outperforms most baselines except SwiftBrush v2, which uses SD2.1 and an additional CLIP loss.Source: §5.1 · p. 7
  • Generated one-step student images generally show richer detail and closer text-prompt alignment than teacher samples in the visual comparison.Figure 7 also compares reverse-KL and JS outputs for a COYO train prompt.Source: §5.1 · p. 7, Figure 7 · p. 8

5.2. Behavior of weighting function h

  • With less mode-seeking divergences, weighting h generally decreases as teacher–fake score differences rise, downweighting samples in low-density data regions.The comparison uses 6.4K generated samples; the authors note teacher score estimates are prone to inaccuracy in low-density regions.Source: §5.2 · pp. 7–8, Figure 6 · p. 8

6. Related work

Related work distinguishes trajectory distillation from distribution matching for reducing diffusion sampling to one step. Within distribution matching, f-distill unifies f-divergence objectives that prior approaches treated separately.

6. Related work

  • Numerical solvers reduce diffusion sampling steps but usually still require more than 20 steps because of discretization error.Distillation instead targets a single network call.Source: §6 · p. 8
  • Trajectory distillation teaches a one-step student to mimic the teacher’s deterministic sampling process, including noise-to-data mappings and consistency functions.Progressive distillation iteratively halves the number of sampling steps.Source: §6 · p. 8
  • Distribution matching aligns the one-step student’s sample distribution with the teacher’s, using adversarial training or divergence minimization.Variational score distillation has been used for reverse-KL, forward-KL, and Fisher divergence objectives.Source: §6 · p. 8
  • Unlike prior methods that minimize a particular divergence and may require different training strategies, f-distill unifies f-divergences in a single framework.The framework includes reverse-KL, forward-KL, Jensen-Shannon, and other divergences.Source: §6 · p. 8

7. Conclusions

The conclusion presents f-distill as a general f-divergence framework for distribution-matching distillation. Its gradient combines score differences with density-ratio weighting, and experiments demonstrate one-step generation capabilities.

7. Conclusions

  • f-distill frames distribution-matching distillation as minimization of an f-divergence.Source: §7 · p. 8
  • Its gradient update multiplies the teacher–student score difference by a weighting function and naturally downweights regions with larger score-estimation errors.Source: §7 · p. 8
  • Experiments across image-generation tasks demonstrate strong one-step generation capabilities for f-distill.Source: §7 · p. 8

Supplementary Material

The supplementary material states assumptions for the gradient proof, derives the f-divergence gradient, and establishes that continuous non-negative weighting functions correspond to f-divergence gradients.

Supplementary Material

  • Theorem 1 assumes a twice continuously differentiable f and a student distribution induced by a differentiable generator mapping from a prior.The sampled intermediate variable is x = Gθ(z) + σ(t)ϵ, with z drawn from the prior and ϵ standard Gaussian.Source: A. Proofs · p. 12
  • The proof expands the total derivative into terms associated with differentiating f with respect to x and q, then simplifies them using an integration identity.The derivation applies the log-derivative trick in a later step.Source: A. Proofs · p. 13
  • Proposition 1 states that any continuous non-negative weighting function h on (0, +∞) yields an expectation corresponding to an f-divergence gradient.The proof constructs a convex f with f(1)=0 and f''(r)=h(r)/r^2.Source: A. Proofs · p. 13

B. Training details

The supplementary sections document training choices, explain how divergence choices affect mode coverage and weighting, and report diversity and image-quality evaluations. Across text-to-image datasets and guidance settings, JS is reported to outperform reverse-KL on in-batch similarity and nearly all HPSv2 prompt categories.

B. Training details

  • ImageNet-64 training uses 2e-6 for 200k iterations, then 5e-7 for 180k; batch sizes vary by dataset.Batch sizes are 2048/512/1024/384 for CIFAR-10/ImageNet-64/COYO SD v1.5/COYO SDXL; COYO-Aesthetic-6.5 filters images below aesthetic score 6.5.Source: B. Training details · p. 13, Table 5 · p. 14

C.1.1. Classification by mode-seeking

  • Mode-seeking describes fitting a generative model that captures only dominant data modes, which can mean less diversity; the authors classify divergences using a limit-based criterion.The criterion treats a smaller limit as greater tolerance for large density ratios.Source: §C.1.1 · pp. 14–15

C.1.2. Relation to the weighting function h

  • For less mode-seeking divergences, h(r) tends to increase with the density ratio, downweighting regions where teacher density is lower.The text gives JS and forward-KL as examples of increasing h, while reverse-KL has constant h.Source: §C.1.2 · p. 15, Figure 9 · p. 16

C.2. f-divergence and Fisher divergence

  • Fisher-divergence methods match teacher and student score functions rather than their probability densities, but direct optimization is challenging.Prior approaches use assumptions such as stop-gradient on the sampling distribution or Monte Carlo estimation of an intractable gradient term.Source: §C.2 · pp. 15–16

D.1. CIFAR-10 result with more baselines

  • The CIFAR-10 comparison uses a simple classification head on teacher features and is intended to analyze f-distill variants without hyperparameter tuning.The supplementary comparison adds baselines to the main-text table; Table 7 reports FID and Recall scores.Source: §D.1 · p. 16, Table 7 · p. 17

D.2. Diversity evaluation based on in-batch similarity

  • JS outperforms reverse-KL in in-batch similarity across datasets and CFGs, with a larger diversity gain at higher CFG.The authors interpret this as evidence that JS preserves more modes.Source: §D.2 · p. 16, Table 8 · p. 17

D.3. Image quality evaluation based on HPSv2

  • JS consistently outperforms reverse-KL on almost all HPSv2 prompt categories and, at CFG=5, competes with the 50-step SD v1.5 teacher at CFG=8 and SwiftBrush v2.HPSv2 is used to assess image quality; higher scores are better.Source: §D.3 · p. 16, Table 9 · p. 17

D.4. Training stability

The supplementary experiments examine training stability, higher classifier-free guidance, and the GAN-loss ablation. Forward-KL fluctuates more and has worse FID than JS on COYO-700M, while both JS and reverse-KL diverge at CFG=8.

D.4. Training stability

  • Forward-KL has greater loss fluctuation and FID 8.70, compared with JS at FID 7.45, on COYO-700M with SD v1.5.The authors hypothesize inaccurate discriminator density-ratio estimates contribute to the instability and leave stabilization techniques for future work.Source: §D.4 · pp. 17–18, Figure 10 · p. 18

D.5. Higher classifier-free guidance

  • At CFG=8, both JS and reverse-KL diverge; preliminary experiments with clipped teacher scores prevented divergence.Removing GAN loss did not resolve the issue; the authors suggest the weak SD1.5 teacher may be involved.Source: §D.5 · p. 18

D.6. Ablation study on GAN loss

  • The relative FID ranking of f-divergences remains the same with or without GAN loss on CIFAR-10 and MS-COCO 2014.Table 10 reports the GAN-loss ablation.Source: §D.6 · p. 18, Table 10 · p. 18

E. Extended Samples

The extended-sample section presents generated examples for CIFAR-10, ImageNet-64, and COYO-700M text prompts. Teacher and student examples use the same random seeds and class labels or prompts.

E. Extended Samples

Related papers

Loading 2502.15681v2…