[논문 리뷰] Progressive Distillation for Fast Sampling of Diffusion Models

Tim Salimans, Jonathan Ho|arXiv (Cornell University)|2022. 02. 01.

Generative Adversarial Networks and Image Synthesis인용 수 190

한 줄 요약

본 논문은 diffusion-model 샘플링을 압축하기 위한 점진적 증류(progressive distillation)를 도입하여, 단계 수를 반복적으로 절반으로 줄여 최소 4단계로도 고품질 샘플을 달성하면서 성능을 유지한다.

ABSTRACT

Diffusion models have recently shown great promise for generative modeling, outperforming GANs on perceptual quality and autoregressive models at density estimation. A remaining downside is their slow sampling time: generating high quality samples takes many hundreds or thousands of model evaluations. Here we make two contributions to help eliminate this downside: First, we present new parameterizations of diffusion models that provide increased stability when using few sampling steps. Second, we present a method to distill a trained deterministic diffusion sampler, using many steps, into a new diffusion model that takes half as many sampling steps. We then keep progressively applying this distillation procedure to our model, halving the number of required sampling steps each time. On standard image generation benchmarks like CIFAR-10, ImageNet, and LSUN, we start out with state-of-the-art samplers taking as many as 8192 steps, and are able to distill down to models taking as few as 4 steps without losing much perceptual quality; achieving, for example, a FID of 3.0 on CIFAR-10 in 4 steps. Finally, we show that the full progressive distillation procedure does not take more time than it takes to train the original model, thus representing an efficient solution for generative modeling using diffusion at both train and test time.

연구 동기 및 목표

• 무조건적 및 클래스-조건부 이미지 생성을 위한 확산 모델의 느린 샘플링 병목 현상을 동기 부여하고 해결한다.
• 느리고 다단계인 교사로부터 정확도를 빠른 학생 모델로 전달하는 증류 방법을 제안한다.
• 점진적 증류가 단계 수를 줄이면서도 표준 벤치마크(CIFAR-10, ImageNet, LSUN)에서 샘플 품질을 유지함을 입증한다.
• 감소 없이 빠른 증류를 지원하는 안정적인 확산 매개변수화와 손실 가중치 전략을 제공한다.
• 전체 증류 과정이 원래 모델을 훈련시키는 것에 비해 계산적으로 효율적임을 보인다.

제안 방법

• 하나의 학생 단계로 두 교사 DDIM 단계를 맞추도록 학생 모델을 훈련시키는 점진적 증류를 형식화한 뒤, 단계 수를 반복적으로 절반으로 축소한다.
• 주어진 z_t에서 시작하는 교사를 두 DDIM 단계를 실행하고 이를 역전시켜 학생을 위한 선명한 타깃을 얻는 증류 타깃을 사용한다.
• 확산 디노이징 모델의 다양한 매개변수화(x 직접, epsilon, 또는 v)와 낮은 신호-잡음비에서도 유효한 손실 가중치를 실험한다.
• 코사인 노이즈 스케줄과 BigGAN 스타일의 업/다운샘플링을 갖춘 U-Net 아키텍처로 표준 확산 모델 학습 절차를 따른다.
• 증류된 모델을 DDIM 및 확률적 기준선과 비교하고 데이터셋 전반에서 FID(가능한 곳에서 IS) 를 보고한다.
• 오픈 소스 구현과 상세 재현성 메모를 제공한다.

실험 결과

연구 질문

• RQ1 점진적 증류가 확산 모델의 샘플링 단계를 의미 있게 줄이면서 샘플 품질의 큰 손실 없이 가능할까?
• RQ2 무조건적 및 클래스-조건부 생성 전반에서 안정적이고 효과적인 증류를 제공하는 매개변수화 및 손실 가중치 전략은 무엇인가?
• RQ3 표준 벤치마크(CIFAR-10, ImageNet, LSUN)에서 기존의 빠른 샘플링 방법과 비교했을 때 증류 모델의 품질은 어떠한가?
• RQ4 원래 확산 모델을 훈련시키는 것에 비해 점진적 증류의 계산 비용은 어느 정도이며 실제로 효율적인가?

주요 결과

• 증류된 모델은 CIFAR-10에서 단 4개의 샘플링 단계로도 고품질 샘플을 생성하고 경쟁력 있는 FID 점수를 달성한다.
• 점진적 증류는 각 반복에서 단계 수를 절반으로 줄여 상당한 속도 향상을 제공하며 전체 절차가 원래 모델의 학습 시간를 초과하지 않는다.
• 다양한 안정적인 매개변수화(x 직접, epsilon, 또는 v)와 손실 가중치(SNR 기반 변형)가 견고한 성능을 보이고, 차폐 실험에서 불안정한 조합 하나가 확인되었다.
• CIFAR-10, 64×64 ImageNet, 128×128 LSUN 벤치마크에서 증류 모델은 비슷하거나 더 낮은 단계 수에서 더 빠른 기준선보다 우수하거나 견주어 보이며, 특히 4-에서 8단계 구간에서 두드러진다.
• 증류 과정은 여전히 효율적이며: 총 소요 시간은 원래 모델을 훈련하는 데 필요한 시간보다 크지 않아 테스트 시 효율 향상을 실용적으로 만든다.
• 증류된 모델은 확률적 샘플러와 함께 사용할 수도 있어 DDIM과 확률적 기준선 사이의 중간 성능을 보인다.