[논문 리뷰] Diffusion Models Beat GANs on Image Synthesis

[Paper] [Github] 
Prafulla Dhariwal, Alex Nichol 
OpenAI 
11 May 2021

 

Abstract

우리는 diffusion models가 현재 최첨단 생성 모델보다 우수한 이미지 샘플 품질을 달성할 수 있음을 보인다. 우리는 일련의 ablation을 통해 더 나은 아키텍처를 찾아 unconditional 이미지 합성에서 이를 달성한다. 조건부 이미지 합성에서는 classifier guidance를 활용하여 샘플 품질을 더욱 향상시키는데, 이는 분류기의 gradient를 사용하여 다양성과 정확성 간의 균형을 조절하는 단순하고 계산 효율적인 방법이다. 우리는 ImageNet 128×128에서 FID 2.97, ImageNet 256×256에서 FID 4.59, ImageNet 512×512에서 FID 7.72를 달성하며, 샘플당 25번의 forward pass만으로도 BigGAN-deep과 동등한 성능을 보이면서도 분포의 더 넓은 영역을 포괄한다. 마지막으로, 우리는 classifier guidance가 upsampling diffusion models와 잘 결합되어, ImageNet 256×256에서 FID 3.94, ImageNet 512×512에서 FID 3.85로 추가적인 성능 향상을 가져옴을 발견한다.

---

1  Introduction

 

지난 몇 년 동안 생성 모델은 인간과 유사한 자연어를 생성하고, 무한한 고품질의 합성 이미지를 만들어내며, 매우 다양한 인간의 음성과 음악을 생성하는 능력을 갖추게 되었다. 이러한 모델은 텍스트 프롬프트로부터 이미지를 생성하거나, 유용한 특징 표현을 학습하는 등 다양한 방식으로 활용될 수 있다. 이러한 모델들이 이미 현실적인 이미지와 사운드를 생성할 수 있음에도 불구하고, 현재의 최첨단 기술을 넘어서는 개선의 여지가 여전히 많으며, 더 나은 생성 모델은 그래픽 디자인, 게임, 음악 제작을 비롯한 수많은 분야에 광범위한 영향을 미칠 수 있다.

GAN는 FID, Inception Score, Precision과 같은 샘플 품질 지표를 기준으로 대부분의 이미지 생성 작업에서 최첨단 성능을 기록하고 있다. 그러나 이러한 지표 중 일부는 다양성을 완전히 반영하지 못하며, 연구에 따르면 GANs는 최첨단 likelihood 기반 모델보다 낮은 다양성을 학습하는 것으로 나타났다. 또한, GANs는 훈련이 어려워 신중하게 선택된 하이퍼파라미터와 정규화 기법이 없으면 쉽게 붕괴할 수 있다. 이러한 단점으로 인해 GANs는 확장성이 낮고 새로운 도메인에 적용하기 어려운 경우가 많다.

 

이러한 이유로, likelihood 기반 모델을 사용하여 GAN과 유사한 샘플 품질을 달성하려는 많은 연구가 이루어졌다 [51, 25, 42, 9]. 이러한 모델들은 더 높은 다양성을 포착하고, 일반적으로 GANs보다 훈련과 확장이 용이하지만, 여전히 시각적 샘플 품질에서는 부족한 면이 있다. 또한, VAEs를 제외하면, 이러한 모델들은 wall-clock time 기준으로 GANs보다 샘플링 속도가 느리다.

 

Diffusion models은 likelihood 기반 모델의 한 종류로, 최근 연구에서 높은 품질의 이미지를 생성할 수 있음이 입증되었다. 이 모델들은 분포 커버리지, 정적인 훈련 목표, 쉬운 확장성 등의 장점을 제공한다. Diffusion models은 신호에서 점진적으로 노이즈를 제거하여 샘플을 생성하며, 그 훈련 목표는 reweighted variational lower-bound으로 표현될 수 있다.

 

이 모델들은 CIFAR-10 데이터셋에서 최첨단 성능을 기록하고 있지만, LSUN 및 ImageNet과 같은 어려운 생성 데이터셋에서는 여전히 GANs보다 성능이 떨어진다. Nichol과 Dhariwal [43]은 이러한 모델들이 계산량을 증가시킬수록 안정적으로 개선되며, 업샘플링 스택을 활용하면 어려운 ImageNet 256×256 데이터셋에서도 고품질 샘플을 생성할 수 있음을 발견했다. 그러나 해당 모델의 FID는 여전히 이 데이터셋에서 최첨단 성능을 기록하고 있는 BigGAN-deep에 비해 경쟁력이 부족하다.

[9] Rewon Child. Very deep vaes generalize autoregressive models and can outperform them on images. arXiv:2011.10650, 2021. 
[25] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. arXiv:2006.11239, 2020. 
[42] Charlie Nash, Jacob Menick, Sander Dieleman, and Peter W. Battaglia. Generating images with sparse representations. arXiv:2103.03841, 2021. 
[43] Alex Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. arXiv:2102.09672, 2021.
[51] Ali Razavi, Aaron van den Oord, and Oriol Vinyals. Generating diverse high-fidelity images with VQ-VAE-2. arXiv:1906.00446, 2019. 
[57] Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv:2010.02502, 2020.

 

우리는 diffusion models과 GAN 사이의 성능 차이가 적어도 두 가지 요인에서 비롯된다고 가정한다. 첫째, 최근 GAN 연구에서 사용된 모델 아키텍처는 광범위하게 탐색되고 최적화되었다. 둘째, GAN는 다양성을 희생하여 샘플 충실도를 높이는 방식으로 동작하며, 이를 통해 높은 품질의 샘플을 생성할 수 있지만 전체 분포를 충분히 커버하지는 못한다. 본 연구에서는 이러한 이점을 diffusion models에도 적용하고자 하며, 먼저 모델 아키텍처를 개선한 후, 다양성과 샘플 충실도 간의 균형을 조정하는 방안을 고안한다. 이러한 개선을 통해 우리는 여러 지표와 데이터셋에서 GAN를 능가하는 새로운 최첨단 성능을 달성한다.

 

본 논문의 구성은 다음과 같다.

• Section 2에서는 Ho et al. [25]을 기반으로 diffusion models의 기본 개념을 간략히 설명하고, Nichol과 Dhariwal [43], 그리고 Song et al. [57]의 개선점을 다룬다. 또한, 우리의 평가 설정에 대해 설명한다.
• Section 3에서는 FID를 크게 향상시키는 간단한 아키텍처 개선 방안을 소개한다.
• Section 4에서는 분류기의 gradient를 활용하여 diffusion model의 샘플링을 유도하는 방법을 설명한다. 우리는 단 하나의 하이퍼파라미터, 즉 분류기 gradient의 크기를 조정하여 다양성과 충실도 간의 균형을 조절할 수 있음을 발견했으며, 이 gradient 크기 배율을 한 자리 수 정도 증가시켜도 적대적 예제를 생성하지 않는다는 점을 확인하였다.
• Section 5에서는 개선된 아키텍처를 적용한 모델이 비조건부 이미지 합성 작업에서 최첨단 성능을 달성하며, classifier guidance를 사용하면 조건부 이미지 합성에서도 최첨단 성능을 기록함을 보인다. Classifier guidance를 사용할 경우, 단 25회의 forward pass만으로도 BigGAN과 유사한 FID를 유지하면서 샘플링이 가능함을 확인하였다. 또한, 우리는 개선된 모델과 업샘플링 스택을 비교하여 두 접근법이 상호 보완적인 개선을 제공하며, 이를 결합할 경우 ImageNet 256×256 및 512×512에서 최상의 성능을 달성함을 입증한다.

---

2  Background

이 절에서는 diffusion models에 대한 간략한 개요를 제공한다. 보다 자세한 수학적 설명은 Appendix B를 참고하기 바란다. 개념적으로(on a high level), diffusion models은 점진적인 노이즈 추가 과정을 역추적하여 분포에서 샘플을 생성한다. 구체적으로, 샘플링 과정은 초기 노이즈 \( x_T \)에서 시작하여 점진적으로 노이즈가 줄어든 샘플 \( x_{T-1}, x_{T-2}, \dots \)를 생성하며, 최종적으로 \( x_0 \)에 도달한다. 각 타임스텝 \( t \)는 특정 노이즈 수준(noise level)에 해당하며, \( x_t \)는 신호 \( x_0 \)과 일부 노이즈 \( \epsilon \)의 혼합으로 해석될 수 있으며, 신호 대 노이즈 비율은 타임스텝 \( t \)에 의해 결정된다. 본 논문에서는 노이즈 \( \epsilon \)이 diagonal Gaussian distribution,에서 샘플링된다고 가정하는데, 이는 자연 이미지에서 효과적으로 작동하며 수식 전개를 단순화하는 장점이 있다.  

Diffusion model은 \( x_t \)에서 조금 더 "denoised"된 \( x_{t-1} \)을 생성하도록 학습된다. Ho et al. [25]은 이 모델을 함수 \( \epsilon_{\theta}(x_t, t) \)로 매개변수화하여, 노이즈가 추가된 샘플 \( x_t \)의 노이즈 성분을 예측하도록 설계하였다. 모델을 학습하기 위해, 각 미니배치 샘플은 무작위로 선택된 데이터 샘플 \( x_0 \), 타임스텝 \( t \), 그리고 노이즈 \( \epsilon \)을 사용하여 노이즈가 추가된 샘플 \( x_t \)을 생성하는 방식으로 구성된다:
\[
x_t = \sqrt{\bar{\alpha}_t} x_0 + \sqrt{1 - \bar{\alpha}_t} \epsilon
\]
여기서 \( \bar{\alpha}_t \)는 누적 스케일링 계수를 의미한다. 이후 학습 목표는 실제 노이즈와 예측된 노이즈 간의 단순한 평균 제곱 오차(MSE) 손실을 최소화하는 것이다:
\[
L_{\text{simple}} = \mathbb{E}_{x_0, t, \epsilon} \left[ ||\epsilon_{\theta}(x_t, t) - \epsilon||^2 \right]
\]
노이즈 예측 함수 \( \epsilon_{\theta}(x_t, t) \)를 통해 샘플링을 수행하는 방법은 직관적으로 명확하지 않다. Diffusion 모델의 샘플링 과정은 \( x_T \)에서 시작하여 반복적으로 \( x_{t-1} \)을 예측하는 방식으로 진행된다. Ho et al. [25]는 합리적인 가정하에서 \( x_t \)가 주어졌을 때 \( x_{t-1} \)의 분포 \( p_{\theta}(x_{t-1} | x_t) \)를 대각 가우시안 분포로 모델링할 수 있음을 보였다:
\[
p_{\theta}(x_{t-1} | x_t) = N(x_{t-1}; \mu_{\theta}(x_t, t), \Sigma_{\theta}(x_t, t))
\]
여기서 평균 \( \mu_{\theta}(x_t, t) \)는 예측된 노이즈 \( \epsilon_{\theta}(x_t, t) \)의 함수로 계산될 수 있다:
\[
\mu_{\theta}(x_t, t) = \frac{1}{\sqrt{\alpha_t}} \left( x_t - \frac{1 - \alpha_t}{\sqrt{1 - \bar{\alpha}_t}} \epsilon_{\theta}(x_t, t) \right)
\]
분산 \( \Sigma_{\theta}(x_t, t) \)는 고정된 상수로 설정할 수도 있고 [25], 별도의 neural network head를 사용하여 학습할 수도 있다 [43]. 두 방법 모두 diffusion 단계 수 \( T \)가 충분히 크면 고품질 샘플을 생성할 수 있다.  

Ho et al. [25]는 간단한 평균 제곱 오차 손실 \( L_{\text{simple}} \)이, denoising diffusion model을 VAE로 해석하여 유도할 수 있는 실제 variational lower bound \( L_{\text{vlb}} \)보다 실험적으로 더 잘 동작한다는 점을 발견하였다. 또한, 이 손실 함수를 사용한 학습과 대응하는 샘플링 절차는 Song과 Ermon [58]의 denoising score matching 모델과 동등하며, 이 모델은 여러 노이즈 수준에서 학습된 denoising model을 사용하여 Langevin dynamics를 통해 고품질 이미지를 샘플링하는 방식이다. 따라서 우리는 "diffusion models"이라는 용어를 이러한 두 가지 모델 클래스를 포함하는 의미로 사용할 것이다.

 

• [58] Yang Song and Stefano Ermon. Improved techniques for training score-based generative models. arXiv:2006.09011, 2020.

 

2.1  Improvements

Song과 Ermon [58], 그리고 Ho et al. [25]의 획기적인 연구 이후, 최근 여러 논문에서 diffusion models을 개선하는 다양한 방법이 제안되었다. 본 절에서는 본 연구에서 활용하는 몇 가지 주요 개선점을 소개한다.  

Nichol과 Dhariwal [43]은 Ho et al. [25]에서 사용된 방식처럼 분산 \( \Sigma_{\theta}(x_t, t) \)을 고정된 상수로 설정하는 것이 적은 diffusion step에서 샘플링할 때 최적이 아님을 발견하고, 이를 신경망으로 매개변수화하는 방안을 제안하였다. 이때, 분산 \( \Sigma_{\theta}(x_t, t) \)는 다음과 같이 보간된(interpolated) 출력 \( v \)로 정의된다:
\[
\Sigma_{\theta}(x_t, t) = \exp(v \log \beta_t + (1 - v) \log \tilde{\beta}_t)
\]
여기서 \( \beta_t \)와 \( \tilde{\beta}_t \) (Equation 19)는 Ho et al. [25]에서 정의된 역방향 과정의 분산에 대한 상한 및 하한에 해당한다.  

추가적으로, Nichol과 Dhariwal [43]은 \( \epsilon_{\theta}(x_t, t) \)과 \( \Sigma_{\theta}(x_t, t) \)을 함께 학습하기 위해 가중합 손실 함수 \( L_{\text{simple}} + \lambda L_{\text{vlb}} \)을 사용하는 하이브리드 목표를 제안하였다. 이 방법을 사용하면 역방향 과정의 분산을 학습하면서도 샘플 품질을 크게 저하시키지 않고 적은 step으로 샘플링이 가능하다. 본 연구에서는 이러한 손실 함수와 분산 매개변수화를 채택하여 실험 전반에서 사용한다.  

Song et al. [57]은 DDIM을 제안하였는데, 이는 DDPM과 동일한 forward 마진널을 가지지만, 비마르코프(non-Markovian) 노이즈 추가 과정을 통해 역방향 노이즈의 분산을 변경하여 다양한 샘플링 방식을 구현할 수 있도록 한다. 특히, 역방향 노이즈를 0으로 설정하면, 어떤 모델 \( \epsilon_{\theta}(x_t, t) \)이든 잠재 변수(latent)에서 이미지로의 결정론적 매핑으로 변환할 수 있다. 이를 통해 샘플링 단계 수를 줄이는 또 다른 방법을 제공하며, Nichol과 Dhariwal [43]의 연구에서도 이러한 방식이 50개 이하의 샘플링 step을 사용할 때 유용함을 확인하였다. 따라서 본 연구에서는 50개 이하의 샘플링 step을 사용할 경우, DDIM을 적용하는 방식을 채택하였다.

 

2.2  Sample Quality Metrics

모델 간 샘플 품질을 비교하기 위해, 본 연구에서는 다음과 같은 지표를 활용하여 정량적 평가를 수행한다. 이러한 지표들은 실무에서 널리 사용되며 인간의 평가와 상당히 일치하지만, 완벽한 대체 지표는 아니며 샘플 품질 평가를 위한 더 나은 지표를 찾는 것은 여전히 열린 연구 문제로 남아 있다.  

Inception Score (IS)는 Salimans et al. [54]에 의해 제안되었으며, 모델이 전체 ImageNet 클래스 분포를 얼마나 잘 포착하는지, 그리고 개별 샘플이 특정 클래스의 신뢰할 수 있는 예제로 얼마나 잘 생성되는지를 측정한다. 그러나 이 지표는 전체 분포를 포괄하는 정도나 클래스 내 다양성을 충분히 반영하지 못하는 한계를 가진다. 예를 들어, 데이터셋의 일부만 암기한 모델도 여전히 높은 IS 값을 기록할 수 있다 [3].  

이러한 한계를 보완하여, Fréchet Inception Distance (FID)가 Heusel et al. [23]에 의해 제안되었다. 연구자들은 FID가 Inception Score보다 인간의 평가와 더 잘 일치한다고 주장하였다. FID는 Inception-V3 [62]의 잠재 공간에서 두 이미지 분포 간 거리를 측정하는 대칭적인 지표를 제공한다.  

최근에는 sFID가 Nash et al. [42]에 의해 제안되었는데, 이는 기존의 풀링된 특징 대신 공간적(spatial) 특징을 사용하는 FID의 변형이다. 연구자들은 sFID가 공간적 관계를 더 잘 포착하며, 일관된 고수준 구조를 가진 이미지 분포를 보다 정확하게 평가할 수 있음을 발견하였다.  

마지막으로, Kynkäänniemi et al. [32]는 Improved Precision and Recall 지표를 제안하였다. 이 지표는 샘플 충실도를 모델 샘플이 데이터 매니폴드(manifold)에 속하는 비율(precision)로 측정하고, 샘플 다양성을 데이터 샘플이 모델 샘플 매니폴드에 속하는 비율(recall)로 측정한다.  

본 연구에서는 FID를 기본 지표로 사용하여 전체적인 샘플 품질을 비교한다. FID는 샘플의 다양성과 충실도를 모두 반영하며, 최첨단 생성 모델 연구에서 사실상 표준적인 지표로 활용되고 있다. 또한, 샘플의 충실도를 평가하기 위해 Precision 또는 IS, 다양성 및 분포 커버리지를 평가하기 위해 Recall을 사용한다.  

다른 방법과 비교할 때, 가능한 경우 공개된 샘플 또는 모델을 사용하여 지표를 다시 계산한다. 이는 두 가지 이유 때문이다. 첫째, 일부 연구 는 쉽게 접근할 수 없는 임의의 학습 데이터 하위 집합을 기준으로 비교를 수행한다. 둘째, 구현상의 미묘한 차이가 FID 값에 영향을 미칠 수 있다. 일관된 비교를 보장하기 위해, 본 연구에서는 전체 학습 데이터를 기준 배치(reference batch)로 사용하며, 모든 모델에 대해 동일한 코드베이스를 사용하여 평가를 수행한다.

---

3  Architecture Improvements

이 절에서는 diffusion models의 샘플 품질을 최적화할 수 있는 모델 아키텍처를 찾기 위해 여러 아키텍처에 대한 ablation study를 수행한다.  

Ho et al. [25]는 diffusion models을 위해 UNet 아키텍처를 도입하였으며, Jolicoeur-Martineau et al. [26]는 이 아키텍처가 이전의 denoising score matching 모델에서 사용된 아키텍처들 [58, 33]보다 샘플 품질을 크게 향상시킨다는 것을 발견하였다. UNet 모델은 잔차(residual) layer와 다운샘플링 컨볼루션을 포함하는 스택과, 업샘플링 컨볼루션을 포함하는 잔차 layer 스택을 사용하며, 동일한 공간 크기를 갖는 layer 간에 skip connection을 적용한다. 또한, 16×16 해상도에서 단일 head를 사용하는 글로벌 attention layer를 도입하고, 각 residual block에 timestep embedding을 투영(projection)하는 방식을 사용하였다.  

Song et al. [60]은 CIFAR-10 [31] 및 CelebA-64 [34] 데이터셋에서 UNet 아키텍처의 추가적인 개선이 성능을 향상시킨다는 점을 발견하였다. 본 연구에서는 ImageNet 128×128에서도 동일한 결과를 확인하였으며, 더 높은 해상도를 갖는 대규모 및 다양한 데이터셋에서도 아키텍처가 샘플 품질에 상당한 영향을 미친다는 점을 보여준다.

 

본 연구에서는 다음과 같은 아키텍처 변경을 탐색하였다:

• 모델 크기를 유지하면서 깊이(depth)와 너비(width)를 조정
• attention head 수 증가
• 32×32, 16×16, 8×8 해상도에서 attention 적용 (기존에는 16×16에서만 사용)
• BigGAN의 잔차 블록(residual block)을 사용하여 업샘플링 및 다운샘플링 수행 (Song et al. [60] 방식)
• 잔차 연결을 \( \sqrt{1/2} \)로 리스케일링 (Song et al. [60], Dhariwal et al. [27, 28] 방식)

 

모든 실험에서 ImageNet 128×128을 사용하여 batch size 256으로 모델을 학습하고, 250개의 샘플링 step을 사용하여 평가하였다. Table 1에서 FID 값을 기준으로 아키텍처 변경에 따른 성능 차이를 비교하였으며, residual connection을 리스케일링하는 경우를 제외한 모든 변경이 성능을 향상시키며, 이러한 개선이 복합적으로 작용할 때 성능이 더욱 향상됨을 확인하였다.

 

Figure 2에서, 모델의 깊이를 증가시키면 성능이 향상되지만 학습 시간이 길어지고 동일한 성능을 달성하는 데 더 오랜 시간이 걸린다는 점을 관찰하였다. 따라서 추가 실험에서는 깊이를 증가시키는 변경을 적용하지 않았다.  

또한, 본 연구에서는 Transformer 아키텍처과 더 잘 맞는 attention configuration을 연구하였다. 이를 위해, attention head 수를 일정하게 유지하는 방법, head당 채널 수를 일정하게 유지하는 방법을 실험하였다. 이 외의 아키텍처 설정은 기본 채널 수 128, 각 해상도에서 2개의 residual block, 다중 해상도 attention, BigGAN 업샘플링 및 다운샘플링 적용으로 고정한 채, 총 70만(700K) iteration 동안 학습하였다. Table 2는 이러한 실험 결과를 보여주며, attention head 수를 늘리거나 head당 채널 수를 줄이는 것이 FID를 개선하는 데 효과적임을 확인하였다. Figure 2에서는 head당 64개 채널을 사용하는 것이 벽시계 시간(wall-clock time) 기준으로 가장 효율적이라는 점을 보여주었으며, 따라서 본 연구에서는 기본 설정으로 head당 64개 채널을 사용하기로 결정하였다. 이 설정은 최신 Transformer 아키텍처와의 일관성을 유지하면서, 최종 FID 성능에서도 기존 설정과 동등한 수준의 성능을 달성하였다.

 

3.1  Adaptive Group Normalization

본 연구에서는 adaptive group normalization (AdaGN)이라 부르는 layer을 실험적으로 적용하였다. 이 layer는 group normalization 연산 이후에 timestep 및 class embedding을 각 residual block에 주입하는 방식으로 동작하며, 이는 adaptive instance normalization (AdaIN) [27] 및 FiLM [48]과 유사한 개념이다.  

AdaGN layer는 다음과 같이 정의된다:
\[
\text{AdaGN}(h, y) = y_s \cdot \text{GroupNorm}(h) + y_b
\]
여기서, \( h \)는 residual block의 첫 번째 convolution 이후의 중간 활성값(intermediate activations)이다. \( y = [y_s, y_b] \)는 timestep 및 class embedding의 선형 투영(linear projection)으로부터 얻어진다.  

 

우리는 AdaGN이 초기 diffusion models의 성능을 향상시키는 것을 이미 확인하였으며, 따라서 본 연구의 모든 실험에서 기본적으로 포함된 설정으로 사용하였다. Table 3에서는 AdaGN 사용 여부에 대한 ablation 실험을 수행하였으며, 이 layer가 실제로 FID를 개선한다는 점을 확인하였다.  

모든 비교 실험에서 다음과 같은 동일한 모델 설정을 사용하였다:

• 기본 채널 수 128
• 각 해상도에서 2개의 residual block
• head당 64개 채널을 갖는 multi-resolution attention
• BigGAN 방식의 업샘플링 및 다운샘플링
• 총 70만(700K) iteration 학습

 

본 연구의 나머지 실험에서는 최종적으로 개선된 모델 아키텍처를 기본 설정으로 사용한다:

• 가변적인 너비(variable width)
• 각 해상도에서 2개의 residual block
• 64 채널을 갖는 multi-head attention
• 32, 16, 8 해상도에서 attention 적용
• BigGAN residual blocks을 사용한 업샘플링 및 다운샘플링
• adaptive group normalization(AdaGN)을 사용하여 residual blocks에 timestep 및 class embedding을 주입

---

4  Classifier Guidance

조건부 이미지 합성을 위한 GAN는 잘 설계된 아키텍처뿐만 아니라 클래스 레이블(class labels)을 적극적으로 활용한다. 이러한 방식은 주로 클래스-조건부 정규화 통계(class-conditional normalization statistics) 또는 \( p(y|x) \)와 같은 확률을 예측하도록 설계된 판별자(discriminator heads)와 같은 형태로 적용된다. Lucic et al. [36]은 이러한 모델들이 제한된 레이블(label-limited) 환경에서 synthetic labels(합성 레이블)을 생성하는 것이 유용할 수 있다는 점을 발견하며, 이는 클래스 정보가 모델의 성능 향상에 핵심적인 요소임을 추가적으로 시사한다.  

GAN에서의 이러한 관찰을 바탕으로, 본 연구에서는 diffusion models을 클래스 레이블에 조건화하는 방법을 탐색한다. 우리는 이미 정규화 layer에서 클래스 정보를 활용하고 있다 (Section 3.1). 본 절에서는 또 다른 접근 방식으로, 분류기 \( p(y|x) \)를 활용하여 diffusion 생성기를 개선하는 방법을 연구한다.  

Sohl-Dickstein et al. [56] 및 Song et al. [60]은 사전 학습된 diffusion model이 분류기의 gradient를 이용해 조건부 샘플링을 수행할 수 있음을 보였다. 구체적으로, 노이즈가 추가된 이미지 \( x_t \)를 학습 데이터로 사용하여 분류기 \( p_{\phi}(y | x_t, t) \)를 학습한다. 이후, 분류기의 gradient \( \nabla_{x_t} \log p_{\phi}(y | x_t, t) \)를 활용하여 diffusion 샘플링 과정에서 특정 클래스 레이블 \( y \)를 향해 샘플링이 유도되도록 한다.  

 

• [56] Jascha Sohl-Dickstein, Eric A. Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. arXiv:1503.03585, 2015.
• [60] Yang Song, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv:2011.13456, 2020.

 

본 절에서는 먼저 분류기를 활용한 두 가지 조건부 샘플링 방법을 검토한 후, 분류기들을 실험적으로 적용하여 샘플 품질을 개선하는 방법을 설명한다.  

본 연구에서는 표기법을 단순화하기 위해 \( p_{\phi}(y | x_t, t) \)를 \( p_{\phi}(y | x_t) \), \( \epsilon_{\theta}(x_t, t) \)를 \( \epsilon_{\theta}(x_t) \)로 표현한다. 하지만, 이 함수들은 각 타임스텝 \( t \)마다 서로 다른 함수에 해당하며, 학습 시에는 반드시 입력 \( t \)에 대해 조건화되어야 함을 유의해야 한다.

 

4.1  Conditional Reverse Noising Process

우리는 먼저 unconditional 노이즈 제거 과정 \( p_{\theta}(x_t | x_{t+1}) \)을 갖는 diffusion model에서 시작한다. 이를 특정 클래스 레이블 \( y \)에 조건화하기 위해서는 각 전이(transition)를 다음과 같이 샘플링하면 된다:

여기서 \( Z \)는 정규화 상수이며, 이에 대한 증명은 Appendix H에 포함되어 있다.  

이 분포에서 정확하게 샘플링하는 것은 일반적으로 계산이 불가능하지만, Sohl-Dickstein et al. [56]은 이를 교란된 가우시안 분포(perturbed Gaussian distribution)로 근사할 수 있음을 보였다. 다음은 해당 유도 과정에 대한 개요이다.  

조건부 전이 확률의 근사
Diffusion model은 timestep \( x_{t+1} \)에서 이전 timestep \( x_t \)을 예측하는데, 이는 가우시안 분포를 따른다:

이 분포의 로그 확률은 다음과 같이 표현된다:

여기서 \( C \)는 상수이다.  

또한, 우리는 \( \log p_{\phi}(y | x_t) \)의 곡률(curvature)이 \( \Sigma^{-1} \)에 비해 낮다고 가정할 수 있다. 이는 diffusion step 수가 무한대에 가까워질 때 \( ||\Sigma|| \to 0 \)이 되므로 합리적인 가정이다. 이 경우, Taylor 전개를 사용하여 \( x_t = \mu \) 부근에서 \( \log p_{\phi}(y | x_t) \)를 근사할 수 있다:

여기서 \( g = \nabla_{x_t} \log p_{\phi}(y | x_t) |_{x_t = \mu} \)과 \( C_1 \)은 상수이다. 이를 이용하면, 조건부 확률의 로그를 다음과 같이 근사할 수 있다:

여기서 \( C_4 \)는 정규화 계수 \( Z \)에 해당하는 상수이며, 무시할 수 있다. 따라서, 조건부 전이 연산자(transition operator)는 비조건부 전이 연산자와 유사한 가우시안 분포로 근사할 수 있으며, 단 평균값이 \( \Sigma g \)만큼 이동된 형태가 된다.  

Algorithm 1에서는 이 샘플링 알고리즘을 요약하고 있으며, 본 연구에서는 추가적으로 gradient scale factor \( s \)를 도입하여 이를 조절할 수 있도록 하였다. 자세한 내용은 Section 4.3에서 다룬다.

 

요약

• 원래 diffusion 과정에서 \( x_t \)의 전이 확률은 가우시안 분포 \( \mathcal{N}(\mu, \Sigma) \)를 따름.
• 하지만 특정 조건 \( y \)가 주어지면, 이 정보를 반영해서 \( x_t \)의 평균이 \( \Sigma g \) 만큼 이동한 새로운 가우시안 분포로 바뀌게 됨.
• 여기서 \( g \)는 조건부 확률 \( p_{\phi}(y | x_t) \)의 로그 확률에 대한 그래디언트인데, 이는 조건을 반영하여 데이터를 이동시키는 방향을 결정함.
• 즉, Diffusion 모델에서 조건부 확률을 반영하면, 단순히 평균이 이동한 가우시안 분포로 근사할 수 있음

 

4.2  Conditional Sampling for DDIM

위에서 유도한 조건부 샘플링(conditional sampling) 방식은 확률적(stochastic) diffusion 샘플링 과정에서만 유효하며, DDIM과 같은 결정론적(deterministic) 샘플링 방법에는 직접 적용할 수 없다. 이를 해결하기 위해, 본 연구에서는 Song et al. [60]의 score-based conditioning trick을 적용하며, 이는 diffusion models과 score matching [59] 간의 관계를 활용하는 기법이다. 

 

• [59] Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. arXiv:arXiv:1907.05600, 2020.
• [60] Yang Song, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv:2011.13456, 2020.

 

Score Function을 활용한 조건부 샘플링
노이즈 성분을 예측하는 모델 \( \epsilon_{\theta}(x_t) \)가 주어졌을 때, 이를 이용하여 score function을 다음과 같이 유도할 수 있다:

이를 \( p_{\theta}(x_t) p_{\phi}(y|x_t) \)의 score function에 대입하면:

이를 통해, 조건부 joint distribution의 score에 해당하는 새로운 epsilon prediction을 정의할 수 있다:

이제 수정된 노이즈 예측 값 \( \hat{\epsilon}_{\theta}(x_t) \)를 사용하여 기존 DDIM 샘플링 절차와 동일한 방식으로 샘플링을 수행할 수 있다. 즉, 결정론적 샘플링에서도 classifier guidance를 적용할 수 있도록 변형한 방식이다. Algorithm 2에서는 이 방법을 요약한 샘플링 알고리즘을 제시한다.

 

4.3  Scaling Classifier Gradients

 

대규모 생성 작업에서 classifier guidance를 적용하기 위해, 본 연구에서는 ImageNet에서 분류 모델을 학습하였다. 사용한 분류기 아키텍처는 UNet 모델의 downsampling trunk를 기반으로 하며, 8×8 해상도에서 attention pool을 사용하여 최종 출력을 생성한다. 분류 모델은 diffusion 모델과 동일한 노이즈 분포에서 학습되며, 과적합을 방지하기 위해 랜덤 크롭(random crops)을 추가하였다. 훈련이 완료된 후, Equation 10을 사용하여 Algorithm 1에 따라 diffusion 모델의 샘플링 과정에 분류기를 통합하였다.

 

초기 실험에서 비조건부 ImageNet 모델을 사용하였을 때, classifier gradient를 1보다 큰 상수 배율로 스케일링할 필요가 있음을 발견하였다. 스케일링이 1일 때, 최종 샘플에 대해 분류기가 원하는 클래스에 대해 약 50%의 확률을 할당하는 것을 관찰했으나, 시각적으로 확인했을 때, 생성된 샘플이 원하는 클래스와 일치하지 않았다. classifier gradient를 증가시키면 이 문제가 해결되었으며, 샘플의 class probability가 100%에 가까워졌다. Figure 3에서 이 효과를 확인할 수 있다.

 

Classifier Gradient Scaling의 효과 분석하기 위해 classifier gradient를 scale factor \( s \)로 조정하면, 다음과 같은 관계가 성립한다: \(s \cdot \nabla_x \log p(y | x) = \nabla_x \log \frac{1}{Z} p(y | x)^s\). 여기서 \( Z \)는 정규화 상수이다. 따라서, gradient를 스케일링하면 확률 분포가 재정규화된 형태로 유지되며, \( p(y | x)^s \)에 비례하는 분포로 변환됨을 알 수 있다. \( s > 1 \)일 때, \( p(y | x) \)의 큰 값들이 더욱 강조되면서 분포가 더 날카로워진(sharpened)다. 이는 classifier의 특정 모드를 더 집중적으로 따르게 하며, 샘플 품질(충실도, fidelity)을 향상시키지만, 다양성(diversity)은 감소할 수 있다. 

 

 

위의 유도 과정에서는 비조건부 diffusion 모델 \( p(x) \)을 가정하였지만, 조건부 diffusion 모델 \( p(x | y) \)을 학습한 경우에도 동일한 방식으로 classifier guidance를 적용할 수 있다. Table 4에서는 classifier guidance가 비조건부 및 조건부 diffusion 모델 모두에서 샘플 품질을 크게 향상시킴을 보여준다. Gradient scale이 충분히 크다면, guided unconditional 모델의 FID가 unguided conditional 모델의 FID와 거의 유사한 수준까지 도달할 수 있다. 그러나 클래스 레이블을 직접 사용하여 모델을 학습하는 것이 여전히 성능을 높이는 데 도움이 됨을 확인하였다. 조건부 모델에 classifier guidance를 추가하면 FID가 더욱 향상된다.

 

 

Table 4는 classifier guidance가 precision(충실도)을 향상시키지만, recall(다양성)이 감소한다는 점을 보여준다. 즉, 샘플 품질과 다양성 사이의 trade-off가 존재한다. Figure 4에서는 gradient scale이 증가할수록 recall(다양성)이 점진적으로 감소하고, precision 및 IS(충실도)는 증가하는 양상을 확인할 수 있다. FID와 sFID는 충실도와 다양성을 모두 고려하는 지표이므로, 최적의 성능을 달성하는 scale factor는 중간값에 위치한다.

 

 

Figure 5에서는 classifier guidance를 BigGAN의 truncation trick과 비교하였다. FID와 Inception Score(IS) 간의 trade-off를 고려했을 때, classifier guidance가 BigGAN-deep보다 항상 우수함을 확인하였다. 그러나 precision/recall trade-off에서는 특정 precision 임계값을 초과하면 classifier guidance가 더 높은 precision을 달성할 수 없었다. 즉, classifier guidance는 일정 수준까지는 효과적이지만, 극단적으로 precision을 높이려는 경우에는 한계가 존재한다.