[2021 ICML Spotlight] Improved Denoising Diffusion Probabilistic Models

[Jasonlee1995]

DDPM

저자들이 DDPM에서 중요하게 봤던 점들은 다음과 같음

fix forward variance : variance를 학습하는 것이 성능에 안좋음
simplified loss 사용 : simplified loss의 sample quality가 variational lower bound보다 좋음
DDPM의 단점 : 다른 likelihood based models보다 likelihood가 안좋음 & generation 속도가 느림

제목 그대로 DDPM을 개선한 논문으로, 해당 논문의 4가지 main contribution은 다음과 같음

1. DDPM의 단점이었던 log-likelihood를 개선할 수 있음
   large T 사용, reverse diffusion process의 variance 학습, cosine noise schedule 사용
2. DDPM의 단점이었던 generation 속도를 높일 수 있음
3. DDPM은 GAN보다 target distribution을 잘 cover함
4. DDPM은 scalable함
   model capacity, training compute가 커질수록 sample quality, likelihood가 좋아짐

중요하다고 생각되는 부분만 간단히 요약

1. Improving the Log-likelihood

1. increase diffusion steps
   T = 1000 → T = 4000
2. learn forward process variance
   use hybrid loss $L_{\mathrm{hybrid}} = L_{\mathrm{simple}} + \lambda L_{\mathrm{vlb}}$ to optimize model
3. change to cosine noise schedule
   linear noise schedule → cosine noise schedule
4. reduce gradient noise
   timestep t uniform sampling on training → importance sampling
   (only helpful when using $L_{\mathrm{vlb}}$ as training objective)

2. learn forward process variance details

DDPM에서 $\sigma_t^2$으로 $\beta_{t}$를 사용하던 $\tilde{\beta}_{t}$를 사용하던 별 차이가 없다고 했음
$\beta_{t}$는 upper bound이고 $\tilde{\beta}_{t}$는 lower bound로 opposite extremes임
저자들은 왜 이 값들이 sample quality에 영향을 미치지 않는지가 의문이었다고 함

Figure 1
$\beta_{t}$와 $\tilde{\beta}_{t}$는 t=0 근처의 구간을 제외하고 거의 똑같음
→ choice of $\sigma_t$ might not matter at all for sample quality

Figure 2
diffusion process의 first few step이 NLL에 가장 큰 영향을 미침
→ log-likelihood를 개선하기 위해서 variance를 학습해야함

결론 : sample quality를 생각하면 variance를 학습할 필요가 없지만, log-likelihood를 생각하면 variance를 학습해야함
lower bound와 upper bound가 거의 같기에, variance를 학습하는 것은 매우 까다로울 수 있음
따라서, variance를 직접적으로 학습하지 않고 reparameterize하여 학습
$\beta_{t}$와 $\tilde{\beta}_{t}$를 log domain에서 interpolation하는 scalar를 학습

DDPM에서 사용한 simplified loss에는 variance가 들어있지 않기에 variance를 학습할 수 없음
variance를 학습하기 위해 simplified loss와 variational lower bound loss를 섞은 hybrid loss로 모델 학습
결국 sample quality가 중요하기에 simplified loss term이 variational lower bound term보다 중요함
→ $\lambda = 0.001$ 사용, variational lower bound term으로는 mean 학습 X
(stop-gradient to the mean for the variational lower bound)
즉, simplified loss로만 mean을 학습해서 sample quality를 좋게 유지하고, variational lower bound로만 variance를 학습하여 log-likelihood를 최대한 좋게 만들자...라는 취지

3. change to cosine noise schedule details

linear noise schedule의 문제점 : forward noising process의 end가 too noisy함 → sample quality에 기여 X
linear noise schedule로 학습하면 reverse diffusion process의 초반을 skip해도 sample quality가 괜찮음

diffusion process가 linear drop-off하도록 cosine schedule을 설계
offset s를 사용한 이유 : $\beta_t$가 t=0 근처에서 너무 작아지는 것을 방지하기 위해
(beginning of the process에서 tiny amounts of noise를 사용하면 모델이 noise를 잘 못맞춤)
$\sqrt{\beta_0}$가 pixel bin size인 2/255보다 slightly smaller하도록 설정 → s=0.008
$\mathrm{cos}^2$을 사용한 이유는 그냥 원하던 모양이라서 그런거지, 다른 함수를 사용해도 됨

4. reduce gradient noise details

이론적으로는 $L_{\mathrm{vlb}}$로 학습한 모델이 $L_{\mathrm{hybrid}}$보다 log-likelihood가 좋아야됨

Figure 6
실제로는 $L_{\mathrm{hybrid}}$로 학습한 모델의 log-likelihood가 더 좋았음
그렇다면 왜 $L_{\mathrm{vlb}}$로 학습하는 모델의 log-likelihood가 더 안좋은가?

Figure 7
log-likelihood가 안좋은 이유 → $L_{\mathrm{vlb}}$의 gradient가 $L_{\mathrm{hybrid}}$보다 훨씬 더 noisy하기 때문
즉, $L_{\mathrm{vlb}}$의 gradient가 noisy해서 optimize하는 것이 quite difficult하다는 것

논문에서의 내용과 별개로 high-level 관점에서 보면 왜 $L_{\mathrm{vlb}}$가 학습이 어려운 지를 알 수 있음
$L_{\mathrm{vlb}}$는 timestep t가 뭐가 됐던지 간에 동일하게 보겠다는 loss임
timestep t에 따라 task의 난이도가 달라지는데 이를 반영해주지 않음 → 모델 입장에서는 뭐를 더 우선적으로 학습해야할지 모름
그에 반해 $L_{\mathrm{simple}}$은 timestep t가 클수록 더 높은 가중치를 주는 loss이기에 더 학습이 쉬움

또 한가지 생각은 $L_{\mathrm{vlb}}$로 학습하게 되면 mean, variance를 동시에 학습하는 건데 이게 어려울지도?라는 생각
($L_{\mathrm{hybrid}}$는 $L_{\mathrm{simple}}$로만 mean 학습하고 $L_{\mathrm{vlb}}$로만 variance 학습하기에)

저자들은 timestep t를 uniform sampling해서 모델을 학습하는 것이 unnecessary noise를 야기한다고 hypothesize
이를 극복하기 위해 importance sampling을 해줌
$p_{t}$ 분포를 알 수 없기에, 각 timestep마다 previous 10 loss values를 유지하여 구함
(초기에는 uniform sample하다가, 학습하면서 loss 쌓이면 $p_{t}$에서 sample하되 dynamically update해줌)

주목할 점 : importance sampling이 $L_{\mathrm{vlb}}$에는 큰 도움을 줬지만 $L_{\mathrm{hybrid}}$에는 별 도움이 안됨
(found that the importance sampling technique was not helpful when optimizing the less-noisy $L_{\mathrm{hybrid}}$ objective)

Results and Ablations

diffusion steps 늘리기 → log-likelihood, sample quality 모두에 도움됨
forward process variance 학습하기 → log-likelihood에 도움되나, sample quality는 약간 떨어짐
cosine noise schedule 사용하기 → log-likelihood, sample quality 모두에 도움됨
importance sampling → log-likelihood에 도움됨
(그러나 importance sampling $L_{\mathrm{vlb}}$로 학습한 모델의 sample quality는 다른 모델들에 비해 많이 떨어짐)

결론 : sample quality를 좋게 유지하면서 log-likelihood를 높이기 위해서 $L_{\mathrm{hybrid}}$ 사용하자

diffusion 모델도 이제는 다른 likelihood models와 competitive한 log-likelihood를 가짐
물론 여기서 사용된 모델은 $L_{\mathrm{hybrid}}$가 아니라 $L_{\mathrm{vlb}}$로 구한 best log-likelihood임

2. Improving Sampling Speed

결론 : $L_{\mathrm{hybrid}}$ pre-trained model을 어떠한 fine-tuning 없이 fewer diffusion steps로 high-quality samples 생성 가능
(DDIM과 비슷한 느낌이라고 생각하면 됨)

K size subsequence 뽑는 방법 : 1, T를 포함하여 균등 간격의 실수 K개 선택 → nearest integer로 rounding

Figure 8

$L_{\mathrm{simple}}$ with fixed $\sigma_{t}^2 = \beta_{t}$ : sampling step이 줄어들면 sample quality가 확 안좋아짐
$L_{\mathrm{simple}}$ with fixed $\sigma_{t}^2 = \tilde{\beta}_{t}$ : sampling step이 줄어들어도 어느정도 괜찮음
$L_{\mathrm{DDIM}}$ : sampling step이 줄어들어도 어느정도 괜찮음
$L_{\mathrm{hybrid}}$ : sampling step이 줄어들어도 어느정도 괜찮음

결론이 되게 애매한게, sampling step이 극단적으로 작아지면 $L_{\mathrm{DDIM}}$이 좋음
그리고 sampling step이 적당히 많으면 $L_{\mathrm{simple}}$ with fixed $\sigma_{t}^2 = \tilde{\beta}_{t}$ 써도 괜찮음
$L_{\mathrm{hybrid}}$로 fewer diffusion steps 사용하면서 high-quality samples를 생성할 수 있다는 점에 focus하는게 좋을듯?
(어차피 $L_{\mathrm{hybrid}}$는 log-likelihood가 좋으니)

3. Comparisons to GANs

결론 : GAN과 비교했을 때 diffusion model은 competitive FID, better mode coverage를 가짐

Table 4

likelihood는 mode-coverage에 대한 good proxy이나, GAN은 likelihood를 측정할 수 없기에 비교하기 어려움
이를 극복하기 위해 precision and recall을 이용하여 mode-coverage 측정
GAN은 대부분 class-conditional models이기에, diffusion도 class-conditional하게 설계하여 학습해줌
(class information을 timestep t와 같이 넣어주는데, 자세한건 생략)

Table 4
diffusion 모델은 GAN보다 mode coverage가 좋음
큰 모델을 사용하면 diffusion 모델이 GAN보다 FID도 좋음

4. Scaling Model Size

결론 : diffusion model은 scalable함
(학습 시간, 모델 크기가 커지면 NLL, FID가 좋아짐)

Figure 10

$L_{\mathrm{hybrid}}$로 학습 + depth를 이용하여 model capacity 조절
FID curve looks approximately linear on a log-log plot → FID scales according to a power law
NLL curve does not fit a power law as cleanly

저자들은 NLL curve가 scale에 대해 less favorable한 2가지 이유를 기술
→ unexpectedly high irreducible loss, overfitting to training distribution
(unexpectedly high irreducible loss : diffusion model 자체의 특성으로 인해 loss가 더 안줄어든다는 의미)

저자들은 $L_{\mathrm{vlb}}$로 모델을 학습하는 것이 아니기에, optimal log-likelihood를 달성하는게 어렵다라고도 언급

5. Appendix

A. Hyperparameters

single-head attention → multi-head attention (4 attention heads)
attention on 16 x 16 → attention at 16 x 16, 8 x 8
change on timestep embedding : GroupNorm(h + v) → GroupNorm(h)(w+1) + b
(v, w, b are conditioning vectors)

dataset별 셋업은 생략
주요 포인트만 짚자면 residual block 수, channel 수, dropout rate 정도?
Adam, batch size 128, lr 1e-4, EMA with 0.9999
$\beta_1 = 0.0001/4 , \beta_T = 0.02 / 4, T = 4000$

D. Combining $L_{\mathrm{hybrid}}$ and $L_{\mathrm{vlb}}$ Models

Figure 13
(세로 축이 hybrid / vlb임)
$L_{\mathrm{hybrid}}$ : diffusion process의 middle에서 좋음
$L_{\mathrm{vlb}}$ : diffusion process의 start, end에서 좋음
→ $L_{\mathrm{vlb}}$ focusing more on imperceptible details, hence the lower sample quality

Figure 14
t가 [100, T-100)에서는 $L_{\mathrm{hybrid}}$, 다른데는 $L_{\mathrm{vlb}}$ 사용하도록 ensemble
ensemble 모델은 FID에서는 hybrid 모델보다 좀 떨어지나, NLL에서는 둘보다 좋음

E. Log-likelihood with Fewer Diffusion Steps

Figure 15
constant striding을 이용한 subsequence sampling을 사용해도 FID는 괜찮았으나 log-likelihood에는 꽤 영향을 줌
이를 해결하기 위해 uniform strided schedule에다가 (1, T/K) 구간에서의 모든 t를 사용하도록 살짝 바꿔줌
해당 schedule을 통해 fewer forward step으로도 log-likelihood를 잘 유지할 수 있음

F. Overfitting on CIFAR-10

Figure 16
저자들은 CIFAR-10에서 실험한 모든 모델들이 overfit하지만 similar optimal FID를 가진다는 것을 발견함
이에 대해 저자들은 cosine schedule을 사용한 모델들은 linear schedule로 학습한 모델들보다 optimal performance를 먼저 도달하며, 그 이후로 overfit한다라고 분석
overfit을 방지하기 위해 dropout rate을 더 높임

overfit이 일어나는 이유에 대해서 저자들은 2가지 이유를 제시함
less noise in cosine schedule provide less regularization, cosine schedule makes optimization and thus overfitting easier

G. Early stopping for FID

Figure 17
CIFAR-10에서와 같이, class-conditional 64 x 64 ImageNet에서도 overfitting이 발견됨
(CIFAR-10에서와 같이 training을 오래할수록 FID가 오히려 안좋아짐)
이를 방지하기 위해 EMA hyperparmeter, dropout rate tuning을 해줘야함
(0.9999 or 0.99995 for EMA, 0.1 or 0.3 for dropout)
dropout을 사용하면 best attainable FID 자체는 좋지만, 해당 성능을 내기 위해 학습을 오래해야함
그리고 dropout을 걸어줘도 결국은 overfit하게 됨
저자들은 model size를 차라리 키우고 early stopping을 해라...라고 결론내림