Diversity-Preserved Distribution Matching Distillation for Fast Visual Synthesis

TL;DR

This work tackles the diversity loss observed in distribution matching distillation (DMD) for fast diffusion-based synthesis by introducing Diversity-Preserved DMD (DP-DMD). DP-DMD splits distillation roles between early and late denoising steps: the first step is guided by a target-prediction objective to preserve diversity, while subsequent steps apply the standard DMD loss with gradients from the latter steps blocked at the first step. The approach is lightweight (no perceptual/backbone modules or discriminators) and operates in latent space, delivering competitive visual quality under few-step inference while maintaining higher sample diversity than prior DMD variants. Across multiple text-to-image backbones and benchmarks, DP-DMD demonstrates improved diversity–quality trade-offs and strong GenEval performance, highlighting its practical value for fast, diverse visual synthesis.

Abstract

Distribution matching distillation (DMD) aligns a multi-step generator with its few-step counterpart to enable high-quality generation under low inference cost. However, DMD tends to suffer from mode collapse, as its reverse-KL formulation inherently encourages mode-seeking behavior, for which existing remedies typically rely on perceptual or adversarial regularization, thereby incurring substantial computational overhead and training instability. In this work, we propose a role-separated distillation framework that explicitly disentangles the roles of distilled steps: the first step is dedicated to preserving sample diversity via a target-prediction (e.g., v-prediction) objective, while subsequent steps focus on quality refinement under the standard DMD loss, with gradients from the DMD objective blocked at the first step. We term this approach Diversity-Preserved DMD (DP-DMD), which, despite its simplicity -- no perceptual backbone, no discriminator, no auxiliary networks, and no additional ground-truth images -- preserves sample diversity while maintaining visual quality on par with state-of-the-art methods in extensive text-to-image experiments.

Figures (8)

• Figure 1: Training pipeline of DP-DMD. The first denoising step of the student is guided by a Flow-Matching diversity loss using a teacher-derived intermediate state, with gradients stopped thereafter. The remaining steps are optimized via the DMD objective, leveraging teacher and fake-model scores to refine sample quality. This role separation preserves diversity while maintaining high-fidelity generation in few-step distillation.
• Figure 2: Gradient stopping in DP-DMD. Training dynamics of diversity and preference for DMD, DP-DMD, and a variant without gradient stopping after the first step. All curves start from the 100-th training iteration and are smoothed using exponential moving average.
• Figure 3: Qualitative comparison of diversity supervision methods. Visual results of four distillation variants on the same prompts: (a) vanilla DMD, (b) DMD-LPIPS, (c) DMD-GAN, and (d) the proposed DP-DMD. While perceptual and GAN-based approaches provide limited or unstable diversity gains and often suffer quality degradation, DP-DMD preserves rich sample diversity while maintaining high visual fidelity, demonstrating a more favorable diversity–quality trade-off.
• Figure 4: Qualitative comparison with open-source few-step distillation methods.
• Figure A: Progressive denoising dynamics. Visualization of SD3.5-M esser2024scaling inference exhibits a stage-wise denoising pattern. The left panel shows a trajectory from step 1 to step 17, while the right panel highlights early steps under different noise initializations. Early steps recover the global structural layout, already showing variation across samples and suggesting a strong link to sample diversity, whereas later steps refine fine-grained appearance details and textures.