Command Palette
Search for a command to run...
Distillation par correspondance de distribution améliorée pour la synthèse rapide d’images
Distillation par correspondance de distribution améliorée pour la synthèse rapide d’images
Tianwei Yin Michaël Gharbi Taesung Park Richard Zhang Eli Shechtman Frédo Durand William T. Freeman
Résumé
Des travaux récents ont montré qu’il est prometteur de distiller des modèles de diffusion coûteux en générateurs efficaces à une étape. Parmi eux, la distillation par correspondance de distribution (DMD) produit des générateurs à une étape qui égalent leur professeur en distribution, c’est-à-dire que le processus de distillation n’impose pas de correspondance biunivoque avec les trajectoires d’échantillonnage du professeur. Cependant, pour garantir un entraînement stable en pratique, DMD nécessite une perte de régression supplémentaire calculée sur un vaste ensemble de paires bruit–image, générées par le professeur au moyen de nombreuses étapes d’un échantillonneur déterministe. Cette approche est non seulement coûteuse en calcul pour la synthèse texte–image à grande échelle, mais elle limite également la qualité de l’élève en le liant trop étroitement aux trajectoires d’échantillonnage originales du professeur. Nous présentons DMD2, un ensemble de techniques qui lèvent cette limitation et améliorent l’entraînement de DMD. Premièrement, nous éliminons la perte de régression et la nécessité de construire un ensemble de données coûteux. Nous montrons que l’instabilité qui en résulte provient du fait que le critique « faux » n’estime pas la distribution des échantillons générés avec une précision suffisante, et nous proposons une règle de mise à jour à deux échelles de temps comme remède. Deuxièmement, nous intégrons une perte GAN dans la procédure de distillation, en discriminant entre les échantillons générés et les images réelles. Cela nous permet d’entraîner le modèle élève sur des données réelles, atténuant ainsi l’estimation imparfaite du score « réel » par le modèle professeur et améliorant par conséquent la qualité. Troisièmement, nous introduisons une nouvelle procédure d’entraînement qui permet l’échantillonnage multi-étapes chez l’élève et résout le problème de décalage des entrées entre l’entraînement et l’inférence présent dans les travaux antérieurs, en simulant pendant l’entraînement les échantillons du générateur tels qu’ils sont produits à l’inférence. Prises dans leur ensemble, nos améliorations établissent de nouvelles références en génération d’images en une étape, avec des scores FID de 1,28 sur ImageNet-64×64 et de 8,35 sur COCO 2014 en zéro-shot, surpassant le professeur original malgré une réduction de 500× du coût d’inférence. De plus, nous montrons que notre approche peut générer des images de l’ordre du mégapixel en distillant SDXL, démontrant une qualité visuelle exceptionnelle parmi les méthodes à peu d’étapes et surpassant le professeur. Nous publions notre code et nos modèles pré-entraînés.
One-sentence Summary
DMD2, introduced by researchers at MIT and Adobe Research, enhances distribution matching distillation by eliminating the regression loss through a two time-scale critic update, adding a GAN loss for training on real data, and enabling multi-step sampling, achieving state-of-the-art one-step image generation with FID scores of 1.28 on ImageNet-64×64 and 8.35 on zero-shot COCO 2014, surpassing the teacher model while reducing inference cost by 500×.
Key Contributions
- The method eliminates the regression loss and expensive dataset construction from DMD, and resolves the resulting instability by identifying the fake critic's inaccurate distribution estimate and proposing a two time-scale update rule as a remedy.
- It integrates a GAN loss that discriminates between generated samples and real images, enabling training on real data, mitigating imperfect real score estimation from the teacher, and enhancing visual quality.
- A new training procedure introduces multi-step sampling in the student and addresses the training–inference input mismatch by simulating inference-time generator samples during training, setting new benchmarks with FID scores of 1.28 on ImageNet-64×64 and 8.35 on zero-shot COCO 2014.
Introduction
Diffusion models deliver high-quality visual generation but require many iterative denoising steps, making high-resolution synthesis slow and expensive. Prior distillation methods accelerate sampling by training a few-step student, yet they often degrade quality. Distribution matching distillation (DMD) achieved state-of-the-art one-step results but still relied on a regression loss that mandates costly precomputation of millions of noise-image pairs and caps student quality at the teacher’s level. The authors introduce DMD2, a refined distribution matching approach that eliminates the regression loss entirely while maintaining training stability through a two time-scale update rule. They further integrate a GAN objective that operates at the distribution level, surpassing the teacher’s quality, and devise backward simulation to enable multi-step generation without the domain mismatch typical of previous methods. The result is a fast, scalable generator that sets new benchmarks in one‑step and few‑step image synthesis.
Method
The authors propose an improved distribution matching distillation framework to distill a costly multi-step diffusion model into an efficient few-step generator.
First, they remove the regression loss used in standard DMD, which required expensive precomputed noise-image pairs. To stabilize the training without this regularization, they employ a Two Time-scale Update Rule. Specifically, the fake score estimator is updated more frequently than the generator to ensure it accurately tracks the non-stationary output distribution of the generator.
To surpass the teacher model's quality, the authors integrate a GAN objective. Refer to the framework diagram:
The architecture features a few-step generator Gθ that produces fake images. These are diffused and compared against diffused real images using a frozen real score function (the teacher) and a dynamically trained fake score function. The gradient of the distribution matching objective is computed from the difference between these scores. Additionally, a discriminator branch is added on top of the bottleneck of the fake diffusion denoiser. This discriminator is trained to distinguish between real and fake diffused samples, while the generator minimizes the standard non-saturating GAN loss.
For larger models, the authors extend the method to support multi-step generation using a fixed timestep schedule. During inference, the model alternates between denoising and noise injection. However, training multi-step generators on noisy real images creates a mismatch with inference, where inputs come from previous generator steps. As shown in the figure below:
The authors address this by simulating the inference-time backward process during training. Instead of using forward diffusion on real images, they generate noisy synthetic images by running the current student generator for several steps. The generator then denoises these simulated images, ensuring alignment between training and testing domains.
Finally, the training process alternates between two main steps: optimizing the generator using the distribution matching gradient and the GAN loss, and optimizing the fake score estimator along with its attached discriminator using a denoising score matching objective on fake data and the GAN classification loss.
Experiment
Experiments evaluate DMD2 on class-conditional ImageNet-64×64 and zero-shot text-to-image synthesis on COCO using FID, CLIP scores, and human preference studies. Distilled SDXL and SDv1.5 models with as few as 1–4 sampling steps match or surpass teacher diffusion models in image quality and prompt alignment, while drastically reducing compute. Ablation results show that removing the original DMD regression loss and employing a two-time-scale update rule with an added GAN term prevents training instability and oversaturated outputs, and the integrated distribution matching and GAN framework achieves higher realism and text coherence than either component alone.
On ImageNet-64×64, single-step StyleGAN-XL achieves the lowest FID (1.52) among prior one-step methods, while multi-step diffusion models like RIN require 1000 steps to reach 1.23. Existing distillation techniques perform poorly, with FID scores above 15. The proposed approach outperforms all distillation baselines and even surpasses the teacher diffusion model’s ODE sampler. StyleGAN-XL attains the best FID (1.52) among all single-step baselines, but still trails multi-step diffusion models like RIN (1.23) and ADM (2.07). Prior distillation methods (Progressive Distillation, BOOT) yield FID above 15, while the proposed one-step model significantly outperforms them and the teacher diffusion model.
DMD2 achieves strong zero-shot text-to-image generation on COCO 2014 with a single forward pass, surpassing the SDXL teacher and other acceleration methods in FID while maintaining competitive CLIP scores. The one-step model yields a 4.91-point improvement in FID over the teacher, and a user study confirms it outperforms the teacher in image quality for 24% of samples with 25× fewer steps. DMD2 with one forward pass achieves the lowest FID (19.01) among all methods, improving on the SDXL teacher by 4.91 points. DMD2's CLIP score (0.336) is nearly identical to SDXL-Turbo (0.337), indicating comparable text-image alignment. In a user study, the DMD2 four-step generator outperforms its teacher in image quality for 24% of samples while using 25× fewer forward passes.
Removing the regression loss from DMD raises ImageNet FID from 2.62 to 3.48, indicating training instability. The two-timescale update rule (TTUR) stabilizes training and restores FID to 2.61 without the regression loss, matching the original DMD. Adding a GAN loss on top of TTUR further improves FID to 1.51, a 1.1-point gain over the TTUR-only variant and the best result among all ablations. Omitting the regression loss increases FID from 2.62 to 3.48, reflecting instability. TTUR without regression loss recovers a FID of 2.61, matching the baseline DMD. Combining TTUR and a GAN loss achieves the best FID of 1.51, outperforming adding GAN to the original DMD (2.56) and adding TTUR alone (2.52).
Removing the GAN loss from DMD2 leads to a sharp FID increase and oversaturated images. A pure GAN variant (without distribution matching) obtains the lowest FID but the worst patch-level and text-alignment scores, while omitting backward simulation degrades patch FID. The full DMD2 model achieves the best patch FID and joint highest CLIP, offering a strong balance of fidelity and text alignment. Without the GAN term, FID worsens to 26.90 and generated images become oversaturated and oversmoothed. Ablating distribution matching (pure GAN) yields the lowest FID (13.77) but the worst Patch FID (27.96) and lowest CLIP (0.307), indicating poor text alignment and patch-level quality. Removing backward simulation increases Patch FID to 24.21, showing its importance for image fidelity. The full DMD2 attains the best Patch FID (20.86) and tied-highest CLIP (0.332) while keeping a competitive FID (19.32).
On ImageNet-64×64, the proposed one-step model significantly outperforms prior distillation methods and even surpasses the teacher diffusion ODE sampler, although multi-step diffusion models still achieve the best quality. For zero-shot text-to-image on COCO, DMD2 with a single forward pass achieves a large FID improvement over the SDXL teacher and is preferred in user studies while using 25× fewer steps. Ablation studies reveal that removing the regression loss leads to training instability recoverable by TTUR, and adding a GAN loss further boosts performance. In DMD2, the GAN loss prevents oversaturation, and the full model combining distribution matching, backward simulation, and GAN provides the best balance between patch-level fidelity and text alignment.