arXiv 2021. [Paper] [Github]
Prafulla Dhariwal, Alex Nichol
OpenAI
11 May 2021
1. Introduction
기존의 diffusion 모델들이 좋은 성능을 보여주고 있지만 여전히 GAN의 성능을 따라잡지 못했다. 저자들은 diffusion과 GAN 사이의 격차가 적어도 두가지 요인에서 비롯된다고 설명한다.
- GAN은 이미 충분히 많이 탐구되고 개선되었다.
- GAN은 fidelty와 diversity 사이의 trade-off에서, diversity를 희생해 고품질의 샘플을 생성한다.
저자들은 먼저 모델 아키텍처를 개선한 다음 fidelty를 위해 diversity를 절충하는 방법을 고안한다. 이런 개선 사항을 통해 여러가지 metric과 데이터셋에서 GAN을 능가하는 새로운 모델을 만들었다.
2. Background
Diffusion model은 점진적으로 noise를 추가하는 process를 reverse하여 샘플링을 한다. Noise $x_T$에서 샘플링을 시작하여 점진적으로 noise가 덜 있는 을 샘플링하고 최종적으로 x_0을 샘플링한다. 각 timestep $t$는 특정 noise level에 해당하며, 는 t에 의해 결정되는 noise ϵ과 x_0의 합으로 생각할 수 있다.
Diffusion model은 x_t로부터 noise가 덜 있는 x_t−1을 생성하도록 학습된다. DDPM 논문에서는 모델을 함수 ϵ_θ(x_t,t)로 parameterize하여 x_t의 noise 부분을 예측하도록 하였다. 이 모델을 학습시키기 위해서는 데이터 x_0, timestep t, noise 로 랜덤하게 만들어진 샘플 t가 minibatch에 포함되어야 한다. 학습에 사용하는 목적 함수는 이고, 실제 noise ϵ과 예측된 noise ϵ_θ의 간단한 mean-squared error loss로 구성된다.
DDPM에서는 p_θ(x_{t−1}|x_t)를 가우시안 분포 로 가정하며, 평균 는 로부터 계산할 수 있고 분산 는 상수로 두거나 별도 헤드로 학습한다.
또한, 실제 variational lower bound L_vlb보다 간단한 mean-squared error 목적 함수 L_simple이 더 좋은 결과를 나타낸다.
2.1 Improvement
DDPM 논문이 나오고 여러 논문들에서 diffusion model을 개선하고자 하였다. 이 논문에서 사용한 개선점들을 설명한다.
Improved DDPM 논문은 분산 Σθ(xt,t)을 상수로 고정하는 것이 차선책이며, 다음과 같이 Σθ(xt,t)을 parameterize하는 것을 제안하였다.
v를 출력하는 모델을 두어 βt와 β~t 사이로 interpolate한다. 또한 ϵθ(xt,t)와 Σθ(xt,t)를 모두 학습하기 위하여 가중치 합 을 목적 함수로 사용한다. (λ=0.001)
DDIM 논문은 DDPM과 같은 forward marginals를 가지지만, reverse noise의 분산을 변경하여 다른 reverse sampling이 가능한 non-Markovian noising process를 제안하였다. Noise를 0으로 두면 ϵθ가 deterministic하게 latent와 이미지를 매핑하게 되어 더 적은 step으로 샘플링이 가능하다.
2.2 Sample quality metrics
(1) Inception Score(IS)
- 생성된 샘플이 얼마나 ImageNet의 다양한 클래스 분포를 잘 반영하는지 측정하는 지표
- 모델이 실제 데이터 분포를 충분히 커버하는지 (전체 데이터의 다양성을 보장하는지) 평가하지 못함
- 특정 소수의 데이터만 잘 외워서 생성하는 경우에도 높은 점수를 받을 수 있음 (과적합 가능성)
(2) Fréchet Inception Distance (FID)
- IS보다 다양성을 잘 반영하고, 사람의 직관과 더 일치한다고 주장
- Inception-V3 모델의 latent space에서 두 개의 이미지 분포 간 거리를 측정하는 방식
(3) sFID (spatial Fréchet Inception Distance)
- 공간적(spatial) 특징을 고려하여 이미지 품질을 평가
- 고차원의 구조적 일관성이 있는 이미지(ex. 얼굴, 물체 배열 등)를 평가하는 데 더 유리
(4) Precision & Recall
- Precision(정밀도)는 생성된 샘플 중 얼마나 많은 샘플이 실제 데이터 분포에 속하는 가를 측정
- 이미지의 품질(fidelity) 평가
- Recall(재현율)은 실제 데이터 중 얼마나 많은 샘플이 생성된 샘플 분포에 포함되는가를 측정
- 생성 모델이 데이터의 분포를 얼마나 잘 커버하는지(diversity)를 평가
실험에서 사용한 기준은 아래와 같다.
- FID를 기본 평가 지표로 사용 (fidelity와 diversity 모두 반영)
- Fidelity 평가: Precision or IS
- Diversity 평가: Recall
3. Architecture Improvements
저자들은 다양한 변경사항을 실험하여 최적의 모델 구조를 찾고자 했다.
DDPM은 UNet 아키텍처를 사용하였다. UNet은 residual layer와 downsampling convolution을 사용하여 특징을 추출하고 이후 upsampling convolution을 사용해 해상도를 복원하며, 크기가 같은 layer끼리 skip connection으로 연결한다. 추가로, 16×16 feature map에는 단일 헤드 global attention을 사용하였으며, 타임 스텝 정보를 각 residual block에 투영한다.
저자들은 다음의 아키텍처의 변화를 실험하였다.
- 모델 크기를 일정하게 유지하면서 깊이 증가
- Attention head의 수 증가
- Attention을 뿐만 아니라 32×32와 8×8에서도 사용
- upsampling과 downsampling에 BigGAN의 residual block을 사용
- Residual connection을 로 rescale
위 항목에 대하여 ImageNet 128×128에서 배치 사이즈 256, 샘플링 step 250으로 학습한 모델의 FID는 아래 표와 같다.
Residual connection을 rescale하는 것을 제외하고는 모두 성능이 개전되었으며, 같이 사용하였을 때 긍정적인 효과를 보여주었다.
위 그래프에서 깊이가 증가하면 성능에 도움이 되지만, 학습 시간이 늘어나고 더 넓은 모델과 동일한 성능에 도달하는 데 시간이 더 오래 걸리므로 추가 실험에서 이 변경 사항을 사용하지 않았다.
저자들은 attention 구성에 대해서도 연구했다. 이를 위해 attention head를 상수로 고정하거나 head당 채널 수를 고정하는 실험을 했다. 나머지 아키텍처에서는 128개의 기본 채널, resolution당 2개의 residual block, multi-resolution attention, BigGAN up/downsampling을 사용하고 70만 iteration에 대해 모델을 학습한다.
위 표를 통해 head가 많거나, head 당 채널 수가 적은 경우 FID가 향상됨을 알 수 있다. 위 그래프에서 64개 채널이 wall-clock time에 가장 적합하다는 것을 알 수 있으므로 head당 64개 채널을 기본값으로 사용한다.
3.1 Adaptive group normalization
저자들은 Adaptive Group Normalization(AdaGN)이라는 layer로도 실험을 진행하였다. 이 layer는 group normalization 연산 후 각 residual block에 timestep embedding과 class embedding을 AdaIN과 같은 방법으로 결합한다.
는 residual block 중간의 activation이다. y=[ys,yb]는 timestep embedding과 class embedding의 linear projection으로 구한다.
위 표에서 AdaGN이 실제로 FID를 개선하는 것을 볼 수 있다.
논문의 나머지 부분에서는 다음의 모델 아키텍처를 기본값으로 사용한다
- Resolution당 2개의 residual block
- 가변 너비
- Head당 64개 채널이 있는 multiple head
- 32, 16, 8 resoltion에서의 attention
- BIGAN의 residual block으로 up/downsampling
- Residual block에 timestep embedding과 class embedding을 주입하기 위한 AdaGN
4. Classifier Guidance
잘 설계된 아키텍처를 사용하는 것 외에도 조건부 이미지 합성을 위한 GAN은 클래스 레이블을 많이 사용한다. GAN에서는 classifier처럼 동작하도록 설계된 head가 있는 discriminator를 사용하거나 class-conditional normalization 방법을 사용한다. 이를 통해 클래스 정보가 중요하다는 것을 확인했다.
이를 diffusion model에서도 적용하기 위해 classifier guidance를 활용하여 특정 클래스 정보를 더 잘 따르게 유도한다. 먼저 AdaGN을 통해 클래스 정보를 준다. 여기에서 diffusion generator를 개선하기 위해 classifier를 활용하는 다른 접근 방식이 필요하다.
noise가 있는 이미지 x_t에 대해 classifier p_ϕ(y|x_t,t)를 훈련한 다음, gradient ∇x_tlogp_ϕ(y|x_t,t)를 사용하여 임의의 class y로 샘플링 프로세스를 guide 할 수 있다. 간결화를 위해 p_ϕ(y|x_t,t)=p_ϕ(y|x_t), ϵ_θ(x_t,t)=ϵ_θ(x_t)로 사용한다.
4.1 Conditional reverse noising process
Diffusion model을 레이블 y로 컨디셔닝하면 다음과 같이 샘플링하는 것으로 충분하다.
이 분포는 perturbed Gaussian distribution으로 근사할 수 있다. 아래는 유도 과정이다.
Diffusion model은 x_t+1로부터 x_t를 가우시안 분포를 사용해서 예측한다.
logp_ϕ(y|x_t)가 Σ^−1보다 낮은 곡률을 가진다고 가정할 수 있다. 이 가정은 diffusion step의 수가 무한에 가까워질 때 ||Σ||→0이 되어 합리적이게 된다. 이 경우, x_t=μ 근처에서 테일러 전개를 사용하여 logp_ϕ(y|x_t)를 근사할 수 있다.
여기서 C_1은 상수이다. 이를 대입하면 아래와 같다.
4는 정규화 상수 Z에 해당하기 때문에 무시할 수 있다. 따라서 conditional transition operator는 unconditional transition operator와 유사하지만, 평균이 만큼 이동한 정규분포로 근사될 수 있음을 확인한다. Algorithm 1은 해당 샘플링 알고리즘을 요약한다.
4.2 Conditional sampling for DDIM
조건부 샘플링에 대한 위의 유도는 stochastic한 샘플링에만 유효하며 DDIM과 같은 deterministic한 샘플링 방법에는 적용할 수 없다. 이를 위해 diffusion model과 score matching 사이의 연결을 활용하는 score 기반 컨디셔닝 트릭을 사용한다. 특히 샘플에 추가된 noise를 예측하는 모델 ϵ_θ(x_t)가 있는 경우 다음과 같이 score function을 도출하는 데 사용할 수 있다.
위 식을 score function에 대입하면 아래와 같다.
이에 joint distribution에 대응되는 새로운 epsilon을 아래와 같이 정의할 수 있다.
샘플링 과정에서 ϵ_θ(x_t)대신에 위에서 정의한 새로운 epsilon을 사용하면 DDIM에서 사용하는 샘플링 과정을 사용할 수 있다. 이 샘플링 과정은 아래와 같다.
4.3 Scaling classifier gradients
저자들은 대규모 생성 task에 classifier guidance를 적용하기 위해 ImageNet에서 classification model을 학습시켰다. Classifier의 아키텍처는 최종 출력을 생성하기 위해 8x8 레이어에 attention pool이 있는 UNet 모델의 downsampling 부분을 사용하였다. 해당 diffusion model과 동일한 noising 분포에서 이러한 classifier를 학습시키고 과적합을 줄이기 위해 random crop을 추가하였다. 학습 후 Algorithm 1에 설명된 대로 classifier를 diffusion model의 샘플링 프로세스에 통합하였다.
저자들은 unconditional ImageNet model을 사용한 초기 실험에서는 classifier gradient를 1보다 큰 상수로 scaling해야 한다는 것을 발견했다. Sclae로 1을 사용하면 classifier가 합리적인 확률(약 50%)을 원하는 클래스에 할당한다. 그러나 실제로 확인해보면 이러한 샘플은 의도한 클래스와 일치하지 않는다. Classifier gradient를 scaling하면 이 문제가 해결되었고 classifier의 클래스 확률이 거의 100%로 증가한다. 아래 그림은 이 효과의 예를 보여준다.
위 그림에서 왼쪽은 classifier scale로 1.0을 사용했고, 오른쪽은 10.0을 사용했다. FID는 왼쪽이 33.0, 오른쪽이 12.0으로 클래스에 더 일치하는 이미지가 생성되었다.
Scaling classifier gradient의 효과는
로부터 이해할 수 있다.
s>1일 때 지수에 의해 값이 증폭되기 때문에 p(y|x)^s는 p(y|x)보다 더 뾰족해져 특정 클래스에 대한 확률이 더욱 강하게 강조된다. 즉, 더 큰 scale을 사용하면 classifier의 모드들에 더 초점을 맞추며, 결과적으로, 생성된 샘플이 더 높은 품질(Higher Fidelity)을 갖지만, 다양성(Diversity)은 줄어든다.
위의 식들에서 기존에는 unconditional 모델이었지만, conditional 모델도 학습할 수 있으며, 동일한 방식으로 Classifier Guidance를 적용할 수 있다. 아래 표는 unconditional model과 conditional model 모두의 샘플 품질이 classifier guidance에 의해 크게 향상될 수 있음을 보여준다.
클래스 레이블을 사용하여 직접 학습시키는 것이 여전히 도움이 되지만 충분히 높은 scale로 guide된 unconditional model이 guide되지 않은 conditional model의 FID에 상당히 근접할 수 있음을 알 수 있다. 물론 conditional model을 guide하면 FID가 더욱 향상된다.
또한 classifier guidance가 recall을 희생시키면서 precision를 향상시켜 샘플의 fidelity와 diversity 사이의 trade-off를 도입함을 보여준다. 아래 그래프에서 gradient scale에 따라 이 trade-off가 어떻게 달라지는지 볼 수 있다.
Gradient Scaling을 1.0 이상으로 조정하면 Recall(다양성 측정 지표)이 줄어들고, Precision과 IS(품질 측정 지표)이 증가하는 Trade-off가 발생한다는 것을 확인했다. FID와 sFID는 다양성과 품질을 모두 반영하는 지표이므로, 가장 좋은 값은 적절한 중간값에서 얻어진다.
또한 아래 그림에서 BigGAN의 truncation trick과 classifier guidance를 비교한다.
Classifier Guidance는 FID와 IS 간의 Trade-off에서 BigGAN-deep보다 확실히 더 우수한 성능을 보였다.
다만, Precision과 Recall 간의 Trade-off는 더 명확하지 않았으며, Precision이 낮을 때는 더 좋은 Recall을 제공하지만, Precision이 높아지면 Recall에서 더 이상 BigGAN-deep보다 우수한 성능을 보이지 않았다.
5. Results
5.1 State-of-the-art Image Synthesis
위 표의 결과를 보면 본 논문이 제시한 모델이 모든 테스크에서 가장 낮은 FID를 기록한 것을 확인할 수 있다. sFID도 거의 모든 테스크에서 최고 성능을 보였다. Diffusion Model은 GAN과 비슷한 수준의 Perceptual Quality(시각적 품질)를 유지하면서도, 더 다양한 샘플을 생성할 수 있음을 알 수 있다.
위 사진은 BigGAN-deep과 diffusion model이 생성한 샘플들을 비교한 것이다. 샘플의 시각적 품질은 비슷하지만, diffusion이 GAN보다 더 다양한 샘플을 포함하는 것을 확인할 수 있다.
5.2 Comparison to upsampling
Two-Stage Upsampling은 먼저 저해상도 이미지를 생성한 후, 업샘플링 모델을 통해 고해상도로 변환하는 방식이다. 이는 ImageNet 256×256에서 FID를 크게 향상시켰지만, BigGAN-deep 수준에는 도달하지 못했다.
표는 upsampling model과 guidance model을 비교한 표이다. (Upsampling model은 Improved DDPM의 upsampling stack을 ADM에 적용한 것으로, ADM-U라 표기한다. Classifier guidance와 upsampling을 결합하는 경우 해상도가 더 낮은 모델만 guide하였다.)
표에서 guidance와 upsampling이 서로 다른 방향으로 샘플의 품질을 향상시키는 것을 알 수 있다. Upsampling은 recall을 높게 유지한 채로 precision을 개선하며, guidance는 훨씬 더 높은 precision을 위해 diversity를 절충할 수 있도록 한다.