[논문 리뷰] Knowledge Distillation in Iterative Generative Models for Improved Sampling Speed

Eric Luhman, Troy Luhman|arXiv (Cornell University)|2021. 01. 07.

Generative Adversarial Networks and Image Synthesis참고 문헌 40인용 수 62

한 줄 요약

이 논문은 다단계의 결정론적 DDIM 샘플링 프로세스를 단일 단계의 Denoising Student로 압축하는 지식 증류 접근법을 소개하며, 적대적 학습 없이 CIFAR-10, CelebA, LSUN에서 GAN과 유사한 샘플링 속도와 높은 품질의 샘플을 달성한다.

ABSTRACT

Iterative generative models, such as noise conditional score networks and denoising diffusion probabilistic models, produce high quality samples by gradually denoising an initial noise vector. However, their denoising process has many steps, making them 2-3 orders of magnitude slower than other generative models such as GANs and VAEs. In this paper, we establish a novel connection between knowledge distillation and image generation with a technique that distills a multi-step denoising process into a single step, resulting in a sampling speed similar to other single-step generative models. Our Denoising Student generates high quality samples comparable to GANs on the CIFAR-10 and CelebA datasets, without adversarial training. We demonstrate that our method scales to higher resolutions through experiments on 256 x 256 LSUN. Code and checkpoints are available at https://github.com/tcl9876/Denoising_Student

연구 동기 및 목표

• Iterative 생성 모델(예: DDPMs, NCSNs)의 denoising 단계를 줄여 샘플링 속도를 높이는 것을 목표로 한다.
• 빠른 학생이 교사의 DDIM 출력을 일치하도록 학습하는 지식 증류 프레임워크를 제안한다.
• 아키텍처나 학습 역학을 바꾸지 않는 단순하고 비적대적 objective를 제공하여 증류를 가능하게 한다.
• 의미 있는 잠재 표현을 유지하면서 256x256 LSUN과 같은 더 높은 해상도까지 확장 가능함을 입증한다.

제안 방법

• 교사를 결정론적 다단계 생성 프로세스를 갖는 DDIM으로 모델링한다.
• 학습 가능한 평균 F_student(xT)와 단위 분산을 갖는 p(x0|xT)의 가우시안 근사를 출력하는 학생을 정의한다.
• KL(p_teacher(x0|xT) || p_student(x0|xT))를 최소화하여 학습시키는 대신 F_student와 교사의 출력 F_teacher 사이의 회귀 손실로 단순화된다.
• 교사의 노이즈 예측기와 동일한 아키텍처/가중치로 학생을 초기화하여 지식을 전달한다.
• 훈련을 위해 교사 출력에 가우시안 노이즈를 추가하여 0이 아닌 출력 분포를 보장한다.
• 교사와 학생 모두 timestep T(가장 높은 노이즈 수준)에 조건화한다.
• DDIM의 결정론적 특성을 활용하여 xT에서 x0로의 샘플링을 학생의 단일 평가로 축소한다.

실험 결과

연구 질문

• RQ1 지식 증류가 다단계 DDIM 디노이징 프로세스를 비적대적 학습 없이 단일 단계 모델로 전달할 수 있는가?
• RQ2 표준 벤치마크(CIFAR-10, CelebA) 및 고해상도 데이터셋(LSUN 256x256)에서 샘플 품질(FID/IS)과 속도 면에서 단일 단계 디노이징 학생의 성능은 어떠한가?
• RQ3 증류된 모델이 잠재 공간 구조를 보존하고 의미 있는 보간을 가능하게 하는가?

주요 결과

• Denoising Student는 CIFAR-10에서 1단계로 경쟁력 있는 FID와 IS를 달성한다 (FID 9.36, IS 8.36).
• CelebA에서 경쟁력 있는 FID 10.68을 달성하며(텍스트에 IS는 기재되지 않음).
• 더 높은 해상도 LSUN(256x256)에서도 구조와 색상이 일관되지만 픽셀 단위 재현 목표로 인해 일부 질감 흐림이 남는다.
• 샘플링 속도는 크게 가속되어 교사보다 약 100배, CIFAR-10의 경우 DDPM보다 약 1000배 빠르며 50k 이미지를 생성하는 데 걸리는 시간은 51.5초이다.
• 256x256 LSUN 이미지로의 확장이 가능하며 구를 이용한 보간 결과로 의미 있는 잠재 표현을 학습한다.
• 이 방법은 적대적 학습에 의존하지 않으며 암시적 모델과 유사한 잠재 공간 조작 능력을 유지한다.