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
초록
최근 연구들은 계산 비용이 높은 확산 모델을 효율적인 단일 스텝 생성기로 증류하는 가능성을 보여주었다. 그중 분포 매칭 증류(DMD)는 분포 측면에서 교사 모델과 일치하는 단일 스텝 생성기를 생성하는데, 이 증류 과정은 교사 모델의 샘플링 궤적과의 일대일 대응을 강제하지 않는다. 그러나 실제 환경에서 안정적인 학습을 보장하기 위해 DMD는 결정론적 샘플러의 여러 스텝을 거쳐 교사 모델이 생성한 대규모 노이즈-이미지 쌍을 사용하여 계산된 추가적인 회귀 손실을 필요로 한다. 이는 대규모 텍스트-이미지 합성에 계산적으로 많은 비용이 들 뿐만 아니라, 학생 모델의 품질을 교사 모델의 원래 샘플링 경로에 지나치게 밀접하게 묶어 제한한다. 우리는 이러한 한계를 극복하고 DMD 학습을 개선하는 일련의 기법인 DMD2를 소개한다. 첫째, 회귀 손실과 값비싼 데이터셋 구축의 필요성을 제거한다. 이로 인해 발생하는 불안정성이 '가짜' 비평가가 생성된 샘플의 분포를 충분히 정확하게 추정하지 못하기 때문임을 보이고, 해결책으로 두 시간 척도 업데이트 규칙을 제안한다. 둘째, 증류 절차에 GAN 손실을 통합하여 생성된 샘플과 실제 이미지를 판별한다. 이를 통해 실제 데이터로 학생 모델을 학습시켜 교사 모델의 불완전한 '실제' 점수 추정을 완화하고 품질을 향상시킨다. 셋째, 학생 모델에서 다중 스텝 샘플링을 가능하게 하고, 학습 중 추론 시점의 생성기 샘플을 시뮬레이션하여 이전 연구의 학습-추론 입력 불일치 문제를 해결하는 새로운 학습 절차를 도입한다. 이러한 개선 사항들을 통합하여, 우리는 ImageNet-64×64에서 1.28, 제로샷 COCO 2014에서 8.35의 FID 점수를 달성하여 추론 비용을 500배 절감했음에도 원래 교사 모델을 능가하는 단일 스텝 이미지 생성의 새로운 벤치마크를 수립했다. 나아가 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.