일단계 확산 모델을 위한 EM 증류 기법
EM Distillation for One-step Diffusion Models
May 27, 2024
저자: Sirui Xie, Zhisheng Xiao, Diederik P Kingma, Tingbo Hou, Ying Nian Wu, Kevin Patrick Murphy, Tim Salimans, Ben Poole, Ruiqi Gao
cs.AI
초록
확산 모델은 복잡한 분포를 학습할 수 있지만, 샘플링에는 계산 비용이 많이 드는 반복적 과정이 필요합니다. 기존의 증류 방법들은 효율적인 샘플링을 가능하게 하지만, 매우 적은 샘플링 단계에서 성능 저하, 학습 데이터 접근에 대한 의존성, 또는 전체 분포를 포착하지 못할 수 있는 모드 추구 최적화와 같은 뚜렷한 한계가 있습니다. 우리는 최대 우도 기반 접근법인 EM 증류(EMD)를 제안하며, 이는 확산 모델을 지각적 품질 손실을 최소화하면서 단일 단계 생성기 모델로 증류합니다. 우리의 접근법은 기대값 최대화(EM) 관점에서 유도되었으며, 여기서 생성기 매개변수는 확산 교사 사전 분포와 추론된 생성기 잠재 변수의 결합 분포에서 추출된 샘플을 사용하여 업데이트됩니다. 우리는 재매개변수화된 샘플링 기법과 노이즈 제거 기술을 개발하여 증류 과정을 안정화합니다. 또한, 우리의 방법이 기존의 모드 추구 KL 최소화 방법과 흥미로운 연관성을 가짐을 밝혔습니다. EMD는 ImageNet-64 및 ImageNet-128에서 FID 점수 측면에서 기존의 단일 단계 생성 방법들을 능가하며, 텍스트-이미지 확산 모델 증류에 관한 선행 연구와 비교하여 유리한 성능을 보입니다.
English
While diffusion models can learn complex distributions, sampling requires a
computationally expensive iterative process. Existing distillation methods
enable efficient sampling, but have notable limitations, such as performance
degradation with very few sampling steps, reliance on training data access, or
mode-seeking optimization that may fail to capture the full distribution. We
propose EM Distillation (EMD), a maximum likelihood-based approach that
distills a diffusion model to a one-step generator model with minimal loss of
perceptual quality. Our approach is derived through the lens of
Expectation-Maximization (EM), where the generator parameters are updated using
samples from the joint distribution of the diffusion teacher prior and inferred
generator latents. We develop a reparametrized sampling scheme and a noise
cancellation technique that together stabilizes the distillation process. We
further reveal an interesting connection of our method with existing methods
that minimize mode-seeking KL. EMD outperforms existing one-step generative
methods in terms of FID scores on ImageNet-64 and ImageNet-128, and compares
favorably with prior work on distilling text-to-image diffusion models.Summary
AI-Generated Summary