Consistent Distribution Matching for Data-Free Diffusion Distillation

AI summary

Overview

Research area: Generative modeling — specifically distillation of diffusion and flow-matching models for fast (one-step / few-step) sampling.

Technical level: Advanced. The paper combines flow-map parameterization, Fisher/reverse-KL divergence objectives, Jacobian-vector products, and Wasserstein-2 convergence analysis.

Scope: A data-free, simulation-free distillation method (Consistent Distribution Matching Distillation, CDMD) that uses only a frozen teacher and a single trainable student, with theory, a 2D synthetic study, ImageNet 256×256 benchmarks, and text-to-image results.

What This Paper Is About

Diffusion and flow models produce high-quality images but require hundreds of neural function evaluations at inference because they numerically integrate a differential equation. Distillation trains a student to imitate the teacher in one or a few steps, but most existing methods need curated real data, expensive teacher rollouts, auxiliary "fake score" networks, or adversarial objectives. This paper asks whether a teacher can be distilled into a few-step student using only a frozen teacher and a trainable student — no external data, no simulation of the teacher, no proxy networks — and answers yes with a single objective and two models total.

Key Contributions

1. CDMD framework: A data-free and simulation-free distillation method for diffusion and flow models that uses only a frozen teacher and a trainable student, avoiding external distillation data, teacher rollouts, and auxiliary proxy networks.
2. Unified student: One student network parameterizes an average velocity field over arbitrary time intervals, so the same model serves both as the one-step/few-step generator and as its own instantaneous-velocity surrogate (replacing the auxiliary fake-score model used by DMD-style methods).
3. Theory: A distribution-matching analysis (Theorem 3.1) showing that, under regularity assumptions, driving the CDMD objective to zero implies the squared Wasserstein-2 distance between student flow-map pushforwards and teacher marginals — averaged over time — goes to zero, and that every global minimizer recovers the teacher's average velocity. Notably, this does not require supervision states to be sampled from the teacher's path.
4. JVP variance analysis: An analysis of the variance/second moment of the Jacobian-vector-product correction in the data-free setting, identifying when the correction stabilizes training and when student–teacher path mismatch can hurt, validated on synthetic data.

Main Findings

• State-of-the-art data-free results on ImageNet 256×256: CDMD reaches a 1-NFE FID of 2.04 and a 4-NFE FID of 1.37 with a LightningDiT-XL/2 teacher (teacher FID 1.35), within 40 epochs of training.
• Best 1-NFE FID across all three teachers: 3.47 with SiT-XL/2 (teacher FID 2.06), 2.93 with SiT-XL/2+VA-REPA (teacher FID 1.31), and 2.04 with LightningDiT-XL/2. At 2 NFEs these improve to 2.87 and 2.60 respectively without retraining.
• Scaled baselines under matched budgets: Against DMD2 (3.43), FreeFlow (4.28), MFD (6.67), MeanFlow (4.19, adapted), sCM (4.21, adapted), and FACM (2.15 at 1 NFE / 1.79 at 2 NFE, adapted) with the LightningDiT-XL/2 teacher, CDMD's 1-NFE 2.04 is best. With SiT-XL/2: DMD2 3.93, FreeFlow 4.13, MeanFlow 4.07 vs. CDMD 3.47. With SiT-XL/2+VA-REPA: DMD2 3.86, FreeFlow 3.72, MeanFlow 4.47 vs. CDMD 2.93.
• Better teacher scaling: Swapping SiT-XL/2 for SiT-XL/2+VA-REPA improves CDMD by 0.54 FID, versus 0.41 for FreeFlow and 0.07 for DMD2, while MeanFlow degrades.
• Endpoint matching is not transport matching: On a 2D spiral with the exact velocity field as teacher, CDMD and FreeFlow reduce endpoint Wasserstein-2 distance at the same rate, but only CDMD keeps reducing teacher average-velocity MSE while FreeFlow plateaus much higher. Marginal matching alone cannot distinguish the methods.
• Teacher tangent is the best JVP direction: Among the three tangent choices examined, using the teacher velocity v(x̂_t, t) yields the lowest and most stable JVP second moment on synthetic data, even on student-bootstrapped states.
• Mixed training states work best: Interpolating between the two state constructions with ratio m gives the lowest FID at m = 0.5 (3.41, IS 254.8) and the highest IS at m = 0.875 (261.4, FID 3.57); the extremes m = 0 (FID 3.62) and m = 1.0 (FID 3.84) are worse across 50K training iterations.
• Multi-step behavior depends on guidance: At guidance w = 1.5, FID keeps improving with more sampling steps; at w = 2.0 FID is best at very few steps and worsens with more steps, suggesting strong guidance amplifies student–teacher mismatch over multi-step sampling. A 10% fraction of t = r pairs improves 1-NFE FID but degrades multi-step sampling, while 50% gives more stable few-step performance.
• Text-to-image: CDMD leads most preference and alignment metrics and FID at 4 NFEs when training data is inaccessible. The available content reports baseline numbers for the SANA-1.6B teacher and competing methods (e.g., teacher 20 NFE: FID 30.39, Aesthetic 6.335; teacher 4 NFE: FID 35.65; DMD2 4 NFE: 6.489 / 0.2006 / 0.2439 / 0.8042 / 29.103 / 32.34; Consistency Distillation 4 NFE: 54.38 FID), but the table is truncated before CDMD's own text-to-image row, so its exact values are not reported in the provided text.
• Qualitative results: A 4-step CDMD generator distilled from SANA-1.6B LoRA produces 1024×1024 samples, and uncurated one-NFE class-conditional ImageNet samples are shown.

Methodology in Plain English

• Model the average, not the instant. Instead of learning the instantaneous velocity of the generative ODE, the student learns the average velocity between any two times t and r. The difference between those two times lets you read the instantaneous velocity back out mathematically, while the average itself directly gives you a jump from time t to time r — in one network evaluation.
• One network, two jobs. Because the average velocity can be turned into the instantaneous velocity (via a derivative, computed with one Jacobian-vector product), the student doubles as the "fake score" estimator that DMD-style methods normally train separately. That removes the auxiliary network and the second optimizer.
• Build training states without data or teacher rollouts. Starting from pure noise, the student maps to an earlier time s; the result is then analytically renoised back up to time t using the closed-form flow-matching noising formula. Setting s = t puts the state exactly on the student's own path; setting s = 0 reduces to a noised one-step student sample, as in score distillation. In between, an interpolation ratio m mixes the two.
• Supervise with the teacher's velocity. At those bootstrapped states, the student's velocity surrogate (average velocity plus the detached derivative correction) is regressed toward the frozen teacher's instantaneous velocity. The stop-gradient on the correction term avoids higher-order gradients and stabilizes training. The training distribution evolves with the student, so as the student improves, the supervision states become more informative.
• Theory ties it together. The paper proves that minimizing this local consistency residual shrinks the time-averaged Wasserstein-2 distance between student-induced and teacher marginals, and that the teacher's average velocity is the unique global minimizer — without needing supervision states drawn from the teacher's own path.

Why This Matters

• Research impact: It removes the two heaviest scaffolding pieces in distillation pipelines — real/curated data and the auxiliary fake-score network — collapsing training to one student, one frozen teacher, and one optimizer. It also shows theoretically why supervising on student-generated states is principled, and empirically that matching endpoints is not the same as matching transport.
• Real-world applications:
  • Real-time or on-device image generation and editing, where one or four network evaluations instead of hundreds is the difference between usable and unusable.
  • Distilling proprietary or closed models whose training data cannot be released, since no external data is required.
  • Large text-to-image and video pipelines where memory and inference cost dominate deployment budgets.
  • Interactive creative tools and content pipelines that need low-latency sampling at high resolution (the paper shows 1024×1024 samples).
• Industry relevance: Lower training memory (the paper measures this on a single A100 80GB), a simpler single-objective training loop, and one fewer model to maintain make this markedly cheaper to scale than adversarial or two-player distillation. The ability to add sampling steps at inference time from the same trained student also gives a deployment-time quality/latency knob.

Future Directions

• Closing the multi-step gap under strong guidance: FID degrades with more steps at w = 2.0, so understanding and mitigating the amplified student–teacher mismatch at high guidance is an open problem.
• Characterizing path mismatch: The JVP variance analysis identifies when student–teacher path divergence degrades training; turning that diagnostic into an explicit correction or adaptive scheduling is a natural next step.
• Beyond images: The framework is stated for diffusion and flow models generally, but evaluation covers 2D synthetic data, ImageNet 256×256, and SANA-1.6B text-to-image — video, audio, and 3D generation are untested.
• Fully reported text-to-image benchmarking: The truncated table leaves CDMD's own preference/alignment/FID numbers for SANA-1.6B incomplete in the available content, and how the method scales with larger teachers and longer training schedules remains to be established.

Target Audience

Researchers and engineers working on diffusion/flow-model acceleration, distillation, and efficient generative model deployment. The theory sections (Wasserstein-2 convergence, Fisher divergence, JVP variance) suit readers with graduate-level familiarity with generative modeling and optimal transport, while the method section and ImageNet table are accessible to practitioners who mainly want the recipe and the benchmark comparisons.

Authors’ abstract

Flow and diffusion models suffer from slow inference due to computationally expensive numerical integration. Distillation provides a promising way for a student model to learn from a teacher's dynamics, enabling one-step or few-step generation. However, existing methods often depend on curated distillation datasets, costly teacher rollouts, or auxiliary proxy networks, which complicate model training and scaling. In this work, we propose Consistent Distribution Matching, a simulation-free and data-free distillation method for accelerating diffusion and flow models while preserving strong generative capacity. Our key insight is to unify sample generation and score estimation with one student network. Thus, our framework uses only two models, a frozen teacher and a trainable student, and optimizes one objective. We prove that minimizing our objective indicates Wasserstein convergence of the student flow-map pushforwards to the teacher marginals. On ImageNet 256$\times$256, our method attains an FID of 2.04 with a single function evaluation (1-NFE) and a 4-NFE FID of 1.37 within 40 epochs of training, surpassing the state-of-the-art distillation baselines without data. Our code code and model are available at https://consistentdmd.github.io/.