Diffusion Models Beat GANs on Image Synthesis

Abstract

• Diffusion 모델은 생성 모델로, 이미지 생성 태스크를 해결.

• 기존 Diffusion 모델(2015)을 개선하여 이미지 생성 샘플 퀄리티가 GAN류의 SOTA 결과보다 뛰어남.

• Classifier Guidance 방법을 사용하여 생성 이미지에 조건을 줌.

• ImageNet 128128에서 FID 2.97, ImageNet 256256에서 FID 4.59, ImageNet 512512에서 FID 7.72를 달성했으며, Sample당 25번의 Forward Passes로도 BigGAN-Deep(2018 모델)과 비견됨.

Frechet Inception Distance(FID)

• 2개의 데이터셋 분포 간 거리를 측정하는 메트릭.

• Train 데이터셋의 샘플 분포와 생성 데이터셋의 샘플 분포를 비교

• 비교는 Inception V3 모델을 사용함.

• 이미지셋을 Inception V3 모델에 입력, Feature extractor의 출력에 Global Average Pool을 적용하여 1D Tensor로 Flatten, 이미지셋에 대한 임베딩 텐서((2048,))의 평균과 분산을 계산, 두 이미지셋 분포의 Wasserstein Distance를 계산()

Introduction

• GANs가 Likelihood 기반의 SOTA 모델들보다 다양성을 적게 포착한다는 것 또한 드러남. 게다가, GANs는 조심스럽게 선택한 하이퍼 파라미터와 규제기(Regularizers)가 없으면 붕괴(Collapsing)되어 학습하기 어려움.

• GANs는 이미지 생성 태스크에서 SOTA이지만, GANs의 단점들이 새로운 도메인으로 확장하고 적용하기가 어렵게 만듦. 그 결과, 대부분의 연구들은 Likelihood 기반(e.g. VAE)의 모델들로 GAN 수준의 Sample Quality를 달성하려 진행됨.

• Likelihood 기반 모델은 더 많은 다양성을 포착하고 일반적으로 GANs보다 쉽게 확장되고 학습되지만, 시각 샘플 품질 관점에서 부족함이 있음. 게다가 VAEs를 제외하고, 이러한 모델들로 생성(Sampling)하는 것은 실제 소요 시간(Wall-Clock Time) 관점에서 GANs보다 느림.

• Diffusion 모델은 Likelihood 기반의 모델.

• 정적인 학습 목표(Stationary Training Objective), 그리고 쉬운 확장성(Scalability)을 가지면서 고품질 이미지를 생산함.

• 기존 Diffusion 모델은 CIFAR-10에서 SOTA 결과를 냈지만, LSUN과 ImageNet 같은 어려운 생성 데이터셋에 대해서는 아직 GANs에 뒤쳐져 있음.

• 기존 Diffusion 모델은 학습과정에서 해상도를 올리는 등의 연산량을 늘리면 성능이 향상되지만, FID 점수는 당시 SOTA 결과를 내는 BigGAN-deep에 비견되진 않음.

• 저자는 Diffusion과 GAN의 간극은 2가지 요소로부터 발현된다고 함.
1. 최근의 GAN 구조는 깊게 탐구되었고 정제되었음.
2. GAN은 다양성(Diversity)과 품질(Fidelity) 간 Trade Off 할 수 있음. 즉, 높은 퀄리티 샘플을 생성하기 위해 모든 분포를 커버하지 않을 수 있고, 그 반대도 가능.

• 본 논문은 Diffusion 모델의 구조를 개선하고 다양성과 품질 사이에서 Trade Off 할 수 있도록 구조를 만드는 것이 목표.

Background

• Diffusion은 점진적으로 Noising Process를 역행하는 방식(Denoising).

• 샘플링 방식은 노이즈로 시작하여 노이즈가 덜한 샘플들인 들을 생성해내며 마지막 샘플인 를 생성.

• 각 Timestep 는 특정 노이즈 레벨에 대응되며 는 신호 에 Timestep 에 의해 결정된 노이즈 이 섞인 것으로 생각하면 됨.

• 노이즈 은 Diagonal Gaussian Distribution으로부터 뽑힌 값.

• Diffusion은 로부터 더 Denoised 된 을 생성하도록 학습. 이를 함수 로 변수화하며, 이 함수는 노이즈가 포함된 샘플 의 노이즈 성분을 예측.

• 이 모델을 학습시키기 위해, 랜덤하게 [데이터 샘플 , 타임 스텝 ]를 뽑고, 이를 하나로 모아 노이즈가 포함된 샘플 를 만들어 Minibatch를 구성.
• 여기서 는 Markovian Noising Process이며, 점진적으로 데이터에 노이즈를 추가하여 노이즈가 추가된 부터 까지를 생성함. 는 임의의 타임 스텝에 대한 노이즈의 분산이며, 이고, 는 분산 스케쥴.
• 베이즈 이론을 사용하여 사후 확률(Posterior) 또한 평균 와 분산 정규 분포를 따른다는 것을 알 수 있음. 즉, Denosing 해도 정규분포를 따름.

• 학습 목표는 으로, 다시 말하면 실제 노이즈와 예측된 노이즈 간의 Mean-Squared Error Loss.

• 레퍼런스 논문에 따르면, 합리적인 가정 하에, 분포 를 대각 다변량 정규분포 로 모델링 할 수 있음.

• 여기서 평균 는 함수 로 계산. 분산 는 알려진 상수(레퍼런스 숫자)로 고정.

• 레퍼런스 논문은 Simple Mean-Squared Error Objective 이 Actual Variational Lower Bound 보다 실제 샘플에서 더 잘 작동한다는 것을 보였음.

Sample Quality Metrics

• 샘플 퀄리티를 정량적으로 측정하기 위한 몇가지 방법 존재.

• Inception Score(IS)는 모델이 ImageNet 클래스 분포를 얼마나 포착하는지를 평가함. 그러면서 각 생성된 샘플이 얼마나 하나의 클래스를 설명하는지를 평가. IS의 단점은 전체 분포를 커버하는 것 혹은 하나의 클래스의 다양성을 포착하는 것에 대한 보상이 없고, 작은 Subset을 아예 암기하는 모델이 높은 IS를 갖는 것.

• Frechet Inception Distance(FID)는 IS보다 다양성을 더 잘 포착. FID는 Inception-V3의 Latent 공간에서의 두 개의 이미지 분포간 거리에 대해 대칭적.

• sFID는 FID의 일종으로, Standard Pooled Features가 아닌 공간 Features를 사용.

• 본 논문은 FID를 사용하여 Diversity와 Fidelity를 얼마나 포착하는지를 통해 전반적인 샘플 퀄리티를 평가.

Architecture Improvements

• 본 논문은 기존 Diffusion 모델들의 구조를 개선. 레퍼런스 모델로 Diffusion 모델에 UNet 구조를 사용.

• UNet 모델은 Residual Layers와 Downsampling Conv. 스택을 사용하고, 그 뒤에 Residual Layers와 Upsampling Conv. 스택을 붙임. 그리고 Skip Connection은 같은 공간 크기를 가지는 층들을 연결.

• 또한 16x16 해상도에 단일 헤드를 가진 Global Attention Layer를 사용. 그리고 각 Residual Block에 Timestep Embedding의 Projection을 붙임.

• 본 논문은 모델 구조가 샘플 퀄리티를 크게 개선할 수 있기에 여러 실험을 진행합니다.
• 깊이 대 너비(Depth versus width) 증가, 모델 크기 상대적으로 일정하게 유지
• 어텐션 헤드 수 증가
• 어텐션을 16x16 해상도에만 적용하는 것이 아닌 32x32, 16x16, 8x8에 적용
• Upsampling과 Downsampling에 BigGAN의 Residual Block을 사용
• Residual Connection을 로 Rescaling

• 본 논문의 모든 모델은 ImageNet 128x128에 대해 Batch 크기 256, 250 샘플링 스텝으로 학습.

• Residual Connection을 Rescaling 하는 것을 제외하고 모든 변경사항들이 성능 개선을 가져왔고, 긍정적인 복합 효과를 가져옴.

• 채널 수를 고정시키고 헤드 갯수를 변화를 주는 실험과 헤드 갯수를 고정시키고 채널 수를 증가 시키는 실험 결과를 보면 헤드가 많을수록, 채널 수가 적을수록 FID가 개선되었습니다.

• Figure 2를 보면 64개 채널이 최적의 Wall-Clock 시간을 기록했으며, 이에 따라 실험에서 헤드 당 64개의 채널을 사용.

Adaptive Group Normalization

• 저자는 Adaptive Group Normalization(AdaGN)을 사용하여 실험. AdaGN은 그룹 정규화 연산 후 각 Residual Block에 Timestep 및 클래스 임베딩을 통합함.
• 여기서 는 첫번째 Convolution 뒤의 Residual Block의 중간 활성값(Intermediate Activations)이며, 는 Timesetp과 Class Embedding의 Linear Projection으로 얻음.

• 위 결과를 보면 AdaGN이 Skip Connection과 GN의 조합보다 더 작은 FID 점수를 냄.

Classifier Guidance

• 저자는 Diffusion 모델에 조건을 주는 방법을 탐색.

• 분류기 를 사용하여 Diffusion Generator를 개선하며, 이 분류기 를 노이즈가 많은 이미지 로 학습하고 분류기의 Gradient 를 사용하여 Diffusion Sampling Process를 임의의 클래스 레이블 로 가이드함.

Conditional Reverse Noising Process

• 먼저 Unconditional Reverse Noising Process 으로 시작하여 유도.

• 레이블 로 조건을 주기 위해 다음에 따라 각 Transition을 Sample로 추출.
• 여기서 는 Normalizing Constant.

• 규빈 정리
• 어떤 클래스로 Denoising 할지 조건을 주면서 노이즈가 포함된 샘플을 1 스텝만큼 Denoising 하는 프로세스 .
• 이 프로세스는 Unconditional Transition과 같은 모양의 정규분포를 따름(평균만 평행이동).
• 즉, Noising의 역함수가 맞음.

• 이 샘플링 알고리즘을 요약하면 다음과 같음.

Conditional Sampling for DDIM

• 위의 유도식은 Stochastic Diffusion Sampling Process에서만 유효하고 DDIM(Denoising Diffusion Implicit Model)과 같은 Deterministic Sampling 방법에는 적용할 수 없음.

• 이를 위해 저자는 Score-Based Conditioning 트릭을 사용. 이 트릭은 Diffusion 모델과 Score Matching 간 연결을 활용하는 것으로, 특히 본 논문의 모델 이 샘플에 추가된 노이즈를 예측한다면, 다음의 식을 사용해 Score Function을 유도할 수 있음.

• 위 식을 Score Function 로 대체.

• 마지막으로 새로운 Epsilon Prediction 를 정의. 이 함수는 Score의 결합분포를 나타냄.

• 이 예측 함수를 가지므로 DDIM과 같은 샘플링 기법을 사용할 수 있음.

• 이 알고리즘을 요약하면 다음과 같음.

Scaling Classifier Gradients

• 대규모의 생성 태스크에 Classifier Guidance를 적용하기 위해 Classification 모델을 ImageNet으로 학습시킴.

• Classifier 구조는 UNet의 Downsampling 부분이며 마지막 8x8 레이어에 Attention Pool을 적용하여 최종 결과를 냄.

• 학습 후, 분류기를 알고리즘 1에 나온 것 처럼 Diffusion 모델의 샘플링 프로세스에 통합.

• Unconditional ImageNet 모델로 실험했을 때, 1보다 큰 상수 Factor로 분류기의 Gradients를 스케일해야한다는 것을 알아냄.

• 스케일이 1인 경우, 최종 샘플에 대해 육안 검사 시 의도된 클래스와 일치하지 않았음.

• 분류기의 Gradients를 확대하면 문제가 해결되었고, 분류기의 클래스 확률이 거의 100%로 증가.

• Scaling Classifier Gradients의 효과를 이해하기 위해, 임을 알아야함.
• 여기서 는 임의 상수.

• 일 때, 이 분포는 보다 더 날카로워지는데, 이는 지수에 의해 큰 값들이 증폭되기 때문.

• 즉, 큰 Gradients Scale을 사용하는 것이 분류기의 최빈값(modes)에 더 초점을 맞추며, 이는 Higher Fidelity 샘플들(But Less Diverse)을 생성할 가능성을 높입니다.

• 위 그림을 보면 Scale Factor가 1인 경우 웰시 코기 클래스에 속한다고 확신할 수 없는 샘플들이 생성되었지만, Scale Factor가 10인 경우 확실히 웰시 코기 그림을 생성함.

• 위 테이블 4를 보면, Precision이 증가함에 따라 Recall이 감소하는 것을 보여주며, 따라서 Fidelity Versus Diversity인 Trade-Off를 도입했다는 것을 알 수 있음.

Result

• 본 논문의 모델을 Unconditional Image Generation에 대해 평가하기 위해, LSUN 데이터셋에서 3개의 분리된 Diffusion 모델(Bedroom, Horse, Cat)을 학습시킴.

• Classifier Guidance를 평가하기 위해, Conditional Diffusion 모델을 ImageNet Dataset의 128x128, 256x256, 512x512 해상도로 학습시킴.

State-of-the-art Image Synthesis

• 본 논문의 Diffusion 모델이 각 태스크에서 가장 낮은 FID를 기록했고, 하나의 태스크를 제외하고 sFID에서 가장 낮았음.

• Classifier Guidance가 본 논문의 모델이 The Best GANs를 능가하도록 돕는 것을 관측했고, 25번만의 Diffusion Steps로 더 넓은 분포 Coverage(높은 Recall)를 유지함.

• 위 그림을 보면 왼쪽 사진들은 BigGAN-deep의 샘플들이며 고품질의 이미지들이지만 각 카테고리의 이미지들이 서로 유사함.

• 중간 사진은 본 논문의 Diffusion 모델로, 고품질의 이미지이면서 다양한 모습.

• 오른쪽 사진은 학습 데이터셋 사진

Limitations and Future Work

• Diffusion 모델은 생성 모델링의 유망한 방향성을 제시하지만 샘플링 시간이 GANs보다 느림. 이는 수차례의 Denosing Step 때문.

• Single Step의 Diffusion 모델도 연구가 되고 있지만, 이는 GANs와 기존의 Single-Step Likelihood-Based 모델에 비교될만한 성능을 내지 못함.

• Classifier Guidance 기술은 라벨링 된 데이터셋에만 한정됨.

Conclusion

• Diffusion 모델(정적인 학습 목표를 가진 Likelihood 기반 모델)이 SOTA GANs보다 더 나은 품질의 샘플을 생성.

• 본 논문의 개선된 구조는 Unconditional Image Generation에서도 잘 작동하며, Classifier Guidance 기술로 Class-Conditional 태스크에서도 잘 작동.

• Classifier Gradients의 확대를 통해 Diversity For Fidelity의 Trade-Off를 조정할 수 있음.