빠른 이미지 생성을 위한 개선된 분포 매칭 증류 기법

Improved Distribution Matching Distillation for Fast Image Synthesis

May 23, 2024

저자: Tianwei Yin, Michaël Gharbi, Taesung Park, Richard Zhang, Eli Shechtman, Fredo Durand, William T. Freeman

cs.AI

초록

최근 연구들은 확산 모델을 효율적인 단일 단계 생성기로 증류하는 데 유망한 결과를 보여주었습니다. 이 중 분포 매칭 증류(Distribution Matching Distillation, DMD)는 교사 모델의 샘플링 궤적과 일대일 대응을 강제하지 않으면서도 분포 수준에서 교사 모델과 일치하는 단일 단계 생성기를 생성합니다. 그러나 안정적인 학습을 보장하기 위해 DMD는 결정론적 샘플러를 사용하여 교사 모델이 생성한 많은 수의 노이즈-이미지 쌍을 기반으로 추가적인 회귀 손실을 계산해야 합니다. 이는 대규모 텍스트-이미지 합성에 비용이 많이 들며, 학생 모델의 품질을 교사 모델의 원래 샘플링 경로에 지나치게 의존하게 만드는 한계가 있습니다. 우리는 이러한 한계를 해결하고 DMD 학습을 개선하는 DMD2 기술 세트를 소개합니다. 먼저, 회귀 손실과 비용이 많이 드는 데이터셋 구축의 필요성을 제거합니다. 우리는 이로 인한 불안정성이 가짜 비평가가 생성된 샘플의 분포를 정확하게 추정하지 못하기 때문임을 보이고, 이를 해결하기 위해 두 가지 시간 척도 업데이트 규칙을 제안합니다. 두 번째로, 증류 과정에 GAN 손실을 통합하여 생성된 샘플과 실제 이미지를 구별합니다. 이를 통해 학생 모델을 실제 데이터로 학습시켜 교사 모델의 불완전한 실제 점수 추정을 완화하고 품질을 향상시킵니다. 마지막으로, 다단계 샘플링을 가능하게 하기 위해 학습 절차를 수정합니다. 이 설정에서 학습-추론 입력 불일치 문제를 식별하고, 학습 시간 동안 추론 시 생성기 샘플을 시뮬레이션함으로써 이를 해결합니다. 종합적으로, 우리의 개선 사항들은 단일 단계 이미지 생성에서 새로운 벤치마크를 설정하며, ImageNet-64x64에서 1.28, 제로샷 COCO 2014에서 8.35의 FID 점수를 달성하여 추론 비용을 500배 줄이면서도 원래 교사 모델을 능가합니다. 또한, 우리의 접근법이 SDXL을 증류하여 메가픽셀 이미지를 생성할 수 있음을 보여주며, 몇 단계 방법 중에서도 탁월한 시각적 품질을 입증합니다.

English

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.