Improved Distribution Matching Distillation for Fast Image Synthesis

Recent approaches have shown promises distilling diffusion models into efficient one-step generators. Among them, Distribution Matching Distillation (DMD) produces one-step generators that match their teacher in distribution, without enforcing a one-to-one correspondence with the sampling trajectories of their teachers. However, to ensure stable training, DMD requires an additional regression loss computed using a large set of noise-image pairs generated by the teacher with many steps of a deterministic sampler. This is costly for large-scale text-to-image synthesis and limits the student's quality, tying it too closely to the teacher's original sampling paths. We introduce DMD2, a set of techniques that lift this limitation and improve DMD training. First, we eliminate the regression loss and the need for expensive dataset construction. We show that the resulting instability is due to the fake critic not estimating the distribution of generated samples accurately and propose a two time-scale update rule as a remedy. Second, we integrate a GAN loss into the distillation procedure, discriminating between generated samples and real images. This lets us train the student model on real data, mitigating the imperfect real score estimation from the teacher model, and enhancing quality. Lastly, we modify the training procedure to enable multi-step sampling. We identify and address the training-inference input mismatch problem in this setting, by simulating inference-time generator samples during training time. Taken together, our improvements set new benchmarks in one-step image generation, with FID scores of 1.28 on ImageNet-64x64 and 8.35 on zero-shot COCO 2014, surpassing the original teacher despite a 500X reduction in inference cost. Further, we show our approach can generate megapixel images by distilling SDXL, demonstrating exceptional visual quality among few-step methods.

한 줄 요약

DMD2는 분포 일치 기반 지도 학습을 통해 단계 수를 500배 줄이며 FID 1.28의 높은 품질 이미지를 생성한다.

핵심 기여도

**Regression loss 제거**: DMD2는 기존 DMD에서 필요했던 수백만 개의 노이즈-이미지 쌍 생성을 요구하는 regression loss를 제거함.
**GAN loss 통합**: 생성된 샘플과 실제 이미지를 구별하는 GAN loss를 도입하여 품질 향상.
**Backward simulation 기법**: 학습-추론 입력 불일치 문제를 해결하기 위해 추론 시 샘플을 학습 시뮬레이션.
**FID 향상**: ImageNet-64x64에서 1.28, COCO 2014에서 8.35의 FID 달성, 기존 teacher 모델을 초과.

핵심 아이디어

기존 DMD는 teacher 모델의 분포를 학습하는 데 성공했으나, 안정적인 학습을 위해 수많은 노이즈-이미지 쌍을 생성하는 regression loss가 필요했다. 이는 계산 비용이 높고, 학습 품질을 teacher의 샘플링 경로에 의존하게 만들었다. DMD2는 이 문제를 해결하기 위해 regression loss를 제거하고, 생성된 샘플의 분포를 정확히 추정하지 못하는 "fake critic"의 불안정성을 해결하기 위해 **Two Time-scale Update Rule (TTUR)** 을 도입했다. 또한, GAN loss를 통합하여 teacher의 불완전한 real score 추정을 완화하고, **distribution-level supervision** 을 추가적으로 제공함으로써 품질을 향상시켰다. 마지막으로, **backward simulation** 을 통해 multi-step 샘플링을 가능하게 하며, 학습-추론 불일치를 해결했다.

기술적 접근법

**DMD2 학습 프레임워크**:
Regression loss 제거
GAN loss 통합 (생성 샘플 vs. 실제 이미지 구분)
Two Time-scale Update Rule (TTUR) 적용
Backward simulation 기법을 통한 multi-step 샘플링 학습
**데이터셋**: ImageNet-64x64, COCO 2014, SDXL 기반 megapixel 이미지
**평가 지표**: FID, CLIP Score, Patch FID
**하이퍼파라미터**: TTUR은 critic과 generator의 업데이트 주기를 다르게 설정하여 안정성 향상

주요 결과

**ImageNet-64x64**: FID 1.28 (기존 teacher 대비 +0.32 개선)
**COCO 2014 (zero-shot)**: FID 8.35
**SDXL 기반 megapixel 이미지 생성**: GAN 기반 방법 대비 뛰어난 시각 품질
**추론 비용**: teacher 대비 500배 감소

의의 및 한계

DMD2는 기존 DMD의 주요 한계인 regression loss와 데이터셋 생성 비용을 제거함으로써, 대규모 텍스트-이미지 생성을 가능하게 한다. GAN loss와 distribution matching의 통합은 teacher 모델의 불완전한 추정을 완화하고, 품질을 향상시켰다. 또한, backward simulation을 통해 multi-step 샘플링을 지원하며, 학습-추론 불일치를 해결함으로써 전체 성능을 높였다. 그러나, GAN loss가 텍스트-이미지 일치도(CLIP Score)에 미치는 영향은 제한적일 수 있으며, 이는 추가 연구가 필요한 부분이다.

실용적 활용

DMD2는 대규모 텍스트-이미지 생성, 고해상도 이미지 생성, 실시간 이미지 생성 시스템 등에 적용 가능하다. 특히, 추론 비용이 500배 감소한 점은 클라우드 기반 이미지 생성 서비스나 모바일 환경에서 유용하게 활용될 수 있다.