[Paper Review] I-DDPM: Improved Denoising Diffusion Probabilistic Models 논문 리뷰

• Paper: Improved Denoising Diffusion Probabilistic Models (PMLR 2021): arxiv, code

1. Improving the Log-Likelihood

• DDPM 모델은 FID나 Inception Score 에서는 좋은 score를 냈지만, log-likelihood score는 좋지 않았음
• 논문에서는 비교 지표로 NLL (Negative log-likelihood) 를 사용
  • NLL:
  • generative model이 얼마나 data distribution의 mode에 대한 정보를 잘 담고 있는지를 측정해주는 metric이다. 만약 mode coverage가 안 좋은 모델이라면(mode collapse), 이 score가 안좋게 나온다.
  • 최근 연구에 따르면 log-likelihood score에서의 향상이 sample quality나 learnt feature representation의 발전도 나타낸다고 한다.

본 논문은 DDPM에서 왜 log-likelihood score가 안좋은지를 분석하고 이를 개선하고자 한다. (이 metric score가 안좋다는건 DDPM의 mode coverage 성능이 안좋다는 것) 저자들은 이를 위해 ImageNet이나 CIFAR-10과 같은 diversity가 큰 데이터셋을 가지고 실험을 해보았으며, DDPM의 모델을 수정하여 diversity와 resolution 모두를 잡은 모델을 만들고자 하였다.

1.1 Learning

• (복습) DDPM
  • DDPM의 저자들은 을 로 정의하고, 를 학습하지 않는 fixed variable로 두었다.
  • 본래 는 로 정의하는 것이 이론적으로 옳지만,
  • 저자들은 실험을 하다보니 를 로 단순하게 정의한 후 모델을 학습해도 sample quality 측면에서 를 사용할 때와 비슷하다는 것을 밝혀내었고, 실제 실험에서는 를 사용하였다.
• (Figure 1) diffusion step 이 커지면 와 는 거의 유사해진다.
  • 즉, model distribution을 추정하는데에는 variance 보다 mean 값이 훨씬 중요함을 알 수 있다.

---

• (복습) DDPM
  • DDPM의 loss: variational lower bound(vlb)

  

• (Figure2)는 diffusion process에서 처음 몇 step이 variational lower bound를 결정하는데 매우 중요한 역할을 함을 보여준다. 저자들은 이 원인이 variance 때문이라고 생각하였고, model의 variance인 을 학습시키면 좀 더 좋은 결과를 낼 수 있을 것이라고 판단하였다.
  • Figure2의 결과는 훈련이 다 된 모델에 대해서 t diffusion step의 를 계산해본 것이다. 결과를 보면 초기 step에서의 loss가 유독 큰 것을 확인할 수 있는데, 이렇게 되면 이들의 summation인 최종 loss, 가 제대로 학습되지 않을 수도 있다.
    • 이미지가 어느정도 나왔을 때에 대해서만 (few steps of diffusion process) denoising network가 잘 학습될 수 있음
    • 따라서 처음 몇 step에서의 loss 값을 조정해줄 필요가 있다. 저자들은 이러한 현상의 원인이 variance 때문이라고 판단하였고 이 값을 학습하고자 하였다.

• 저자들은 을 neural network를 통해 학습시켰으며, 이 variance 값을 와 를 interpolation 한 값으로 굉장히 rough하게 정의내렸다.

  

  • step이 커지면 와 중 어떤 걸 사용하더라도 학습이 잘 되니까.. 대략 이 값들과 유사한 값으로 간단하게 정의를 내리고, 처음 몇 step 에서만 detail 하게 값을 잡아주려고 interpolation 기법을 사용한 것 같다
  • 이때, 의 값이 너무 작아서 log scale을 사용했다고 하며, 는 모델에서 예측하고자 하는 output 값으로, variance를 예측해주는 역할을 한다.
  • 학습하는 variable은 vector

• 최종 objective function으로는 hybrid objective를 사용한다.

  

  • : variance에 관계없는 loss term으로, DDPM에서 단순화한 objective function이다.

    

    • 만 사용해서 학습을 하니까 성능이 안나와서 simplifed loss term (를 학습시키는 loss)을 사용했다고 함 (절 참고)

  • : 주의할 점은 여기서 mean 값을 예측하는 gradient는 꺼버리고, 오직 variance만 학습하도록 바꿔주었다고 한다.

    • 가 알아서 denoising 하는 쪽은 잘 학습시키니까 는 variance를 어떤 값으로 정할지만 학습하도록 한 듯?

1.2 Improving the Nosing Schedule

• (복습) DDPM에서는 variance schedule을 10^(-4)부터 2*10^2로 linear하게 증가시켜가며 학습했었다.
• (Figure3) 만 봐도 상단의 row는 너무나도 빠르게 noisy 해지는 걸 확인할 수 있다.
• (Figure4) linear하게 scheduling을 한 경우에는, 실제로 reverse diffusion process에서 20% 정도를 skip하고 이미지를 생성해도 비슷한 결과가 나온다 → 맨 마지막 step의 이미지들이 너무 과하게 noise하다는 것..
  • diffusion step의 마지막 20% step 정도는 전부 gaussian noise를 따른다고 해석할 수도 있다. (T step에서만 (0,1)의 gaussian distribution을 따르는게 아니라 0.8T step에서도 대략 (0,1)의 gaussian distribution을 따른다고 볼 수 있다.)

• 저자들은 이를 해결하기 위해 noise schedule을 다음과 같이 sinusoidal 함수를 이용해서 했다.
  • sinusoidal function을 이용한건 딱히 이유가 있진 않고, 그냥 수학적으로 유명한 식이기도 하고 noise schedule를 이렇게 잡으니까 variance가 잘 잡혀서 그랬다고 한다. (Figure 5 참고)

  

1.3 Reducing Gradient Noise

본래 저자들은 만을 사용해도 학습이 잘될 줄 알았는데, 실제로 ImageNet 64x64 dataset 처럼 diversity가 큰 data에 대해 학습을 하니 를 쓸 때 보다 optimize가 잘 안됐다고 한다.

• (Figure 6) 가 log-likelihood를 더 잘 학습시킴

저자들은 이 원인이 의 gradient에 너무 noise가 많이 껴서 그렇다고 추측하였고, gradient noise를 scale한 후 이를 바탕으로 의 variance를 다시 구했다.

• DDPM에서 처럼 t를 uniform 하게 sampling한 후, 임의의 t step에 대해 vlb loss를 구하면 noise가 너무 많이 껴서 loss가 널뛰기를 한다.

• 따라서 저자들은 이전의 10개의 loss 값들을 queue에 넣어둔 후 이를 바탕으로 scaling을 하도록 해서 안정적으로 loss가 학습되도록 했다고 한다. (학습 초기에는 queue가 비어있으니까 임의로 10개를 sampling해서 uniform 시켰다고 함)
  • batch training을 안정화하는 거랑 비슷한 느낌인 듯?
• (Figure 6) 그 결과 를 사용했을 때의 결과가 가장 좋았다고 한다
  • 에 noise를 줄이는 이 sampling 기법을 사용해보기도 했는데, 이 경우에는 noise가 딱히 안줄었다고 함

Results and Ablations

• cosine schedule을 하면 NLL, FID 모두 좋아짐
• + cosine schedule: baseline보다 NLL은 좋아지지만, FID가 안좋아짐
• + cosine schedule의 경우 NLL은 좋아지나, 매우매우 FID가 안좋아진다.
  • 🧐  는 mean과 variance 모두를 학습하는 loss term인데, 여기서 mean에 비해 별로 중요하지 않은 variance가 비교적 강하게 학습이 되고 있어서 성능이 잘 안나오는 것 같다. mean을 좀 더 잘 학습하도록 constraint를 걸어주면 도 잘 학습될 듯?

log likelihood를 비교해보면 Improved diffusion model이 가장 좋았다

Improving Sampling Speed

본 모델은 4k step으로 학습되었다. 훈련때 처럼 inference 때 4k step을 모두 거쳐 denoising을 하면 sampling 속도가 너무 느리다. 저자들은 기존의 pre-trained model에 대해 좀 더 빠르게 sampling을 하기 위해 diffusion step을 줄여보았다고 한다.

• 예를 들어 모델이 diffusion steps으로 훈련되었다고 하면,
• 여기서 임의의 subsequence 를 뽑아서 sampling noise schedule 를 계산하고 이를 바탕으로 sampling variances를 구한다.

• (Figure 8) 4k steps으로 훈련된 model이 있을 때, 25, 50, 100, 200, 400, 1000, 4000로 sampling한 결과
  • DDIM이나 과 비교했을 때 성능이 좋음

Comparision to GANs

저자들은 GAN model과 비교했을 때 Improved-Diffusion 모델이 얼마나 mode-coverage가 좋은지를 비교하기 위해 precision and recall metric을 사용해봤다고 한다.

• GAN: class conditional model 이용
  • BigGAN-deep with 100M param
• Improved-diffusion: timestep embedding 에 class embedding 를 추가한 후, 이 embedding vector를 residual block에 넣어 model에 추가
  • small model with 100M param for 1.7 M training steps
  • larger model with 270M param for 250K training steps

결과는 Table4 참고

• Precision (quality) & Recall (diversity)

GAN은 unsupervised 하게 model의 distribution을 real data의 distribution에 fitting하나, diffusion model은 직접적으로 모델이 real data의 distribution을 따르도록 학습하기 때문에 좀 더 diversity 적인 면에서 장점이 있다.

Scaling Model Size

최근 연구들은 model size를 키우고, 더 많이 학습시킬 수록 모델의 성능이 좋아짐을 밝혔다.

저자들은 DDPM에서의 model scale과 training compute의 영향을 조사하기 위해 실험을 해보았고, 그결 Figure 10처럼 이 둘을 키울수록 FID와 NLL이 좋아짐을 확인할 수 있었다고 한다.