Command Palette
Search for a command to run...
高速画像生成のための改良型分布マッチング蒸留
高速画像生成のための改良型分布マッチング蒸留
Tianwei Yin Michaël Gharbi Taesung Park Richard Zhang Eli Shechtman Frédo Durand William T. Freeman
概要
近年、高コストな拡散モデルを効率的なワンステップ生成器へと蒸留する手法が有望視されている。中でもDistribution Matching Distillation(DMD)は、分布において教師モデルと一致するワンステップ生成器を生成する。すなわち、その蒸留過程は教師モデルのサンプリング軌跡との一対一対応を強制しない。しかし、実際に安定した訓練を確保するため、DMDは、決定論的サンプラーを用いて教師モデルが多数のステップで生成した大規模なノイズ・画像ペア群を用いて計算される追加の回帰損失を必要とする。これは大規模なテキストからの画像生成において計算コストが高いだけでなく、生徒モデルの品質を制限し、教師モデルの元のサンプリング経路に過度に拘束する。我々はDMD2を導入する。これはこの制限を取り除き、DMD訓練を改善する一連の技術である。第一に、回帰損失と高コストなデータセット構築の必要性を排除する。結果として生じる不安定性は、「偽の」批評家が生成サンプルの分布を十分な精度で推定していないことに起因することを示し、その対策として二重時間スケール更新則を提案する。第二に、GAN損失を蒸留手順に統合し、生成サンプルと実画像を識別する。これにより、生徒モデルを実データで訓練できるため、教師モデルによる不完全な「実」スコア推定を緩和し、品質を向上させる。第三に、生徒モデルにおける複数ステップサンプリングを可能にし、訓練時に推論時の生成器サンプルを模擬することで、先行研究における訓練と推論の入力不一致に対処する新しい訓練手順を導入する。これらの改善を総合することにより、ワンステップ画像生成において新たなベンチマークを確立し、ImageNet-64×64でFIDスコア1.28、ゼロショットCOCO 2014で8.35を達成し、推論コストを500分の1に削減しながら元の教師モデルを上回った。さらに、SDXLを蒸留することでメガピクセル画像を生成できることを示し、少数ステップ手法の中でも卓越した視覚品質を示し、教師モデルを凌駕した。コードと事前学習済みモデルを公開する。
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.