[21’ PMLR] IDDPM: Improved Denoising Diffusion Probabilistic Models

Date: 2024.03.28 Updated: 2024.03.28

카테고리: Generation

0. Abstract

DDPM에서 log-likelihood가 낮았던 문제를 몇 가지 방법을 도입하여 해결하였다.
특히 Variance $\sigma$를 학습하는 것이 고정하는 것보다 유리함을 보였다. 특히 sampling step이 짧은 경우 효과적이었다.
결과적으로 IDDPM의 몇 가지 개선점에 의해 모델이 scalable해졌다. 즉 모델의 크기가 늘어남에 따라 모델의 성능이 상승했다.

1. Improving the Log-likelihood

DDPM은 뛰어난 성능을 보여줬지만, likelihood-based model과 비교해서 log-likelihood가 낮았다. 일반적으로 log-likelihood가 높다는 것은 데이터 분포의 모든 mode를 잘 포착한다는 것을 의미하기에 log-likelihood가 낮다는 것은 치명적인 단점이다. DDPM 저자들은 이 문제가 사람이 거의 인지하지 못하는 정도의 차이라고 주장하여 정당화했지만, 이 논문에서는 이 문제를 해결하기 위해 몇 가지 방법을 제안한다. 참고로 DDPM과는 달리 여기서는 $T=4000$을 사용했다.

1.1. Learning $\Sigma_{\theta}(x_t, t)$

DDPM의 한 가지 찝찝한 부분 중 하나는 아래와 같은 $L_{\text{simple}}$은 오로지 mean $\mu_{\theta}(x_t, t)$만을 학습하고, variance $\Sigma_{\theta}(x_t, t)$는 $\sigma_t^2 \mathbf{I}$로 고정했다는 점이다.

\[L_{\text{simple}} = \mathbb{E}_ {t, x_0, \epsilon} \left[ \Vert \epsilon - \epsilon_\theta(x_t, t) \Vert^2 \right]\]

그리고 DDPM에서는 $\sigma^ 2 _ t$를 $\beta_t$로 고정하든 $\tilde{\beta}_ t$로 고정하든 학습에 큰 차이가 없다고 했는데, 의미적으로 각각은 variance의 upper bound, lower bound를 의미한다고 한다. 양 극단의 값을 선택했는데도 별로 차이가 없었던 이유는 무엇일까? (참조: DDPM 논문 리뷰)

이유는 diffusion step의 길이 $T$가 충분히 길면 $\beta_t$와 $\tilde{\beta}_t$의 차이가 무의미해지기 때문이다. 위 그래프에서 $t=0$ 부근에서는 차이가 나지만, 이미 이 부근에서는 imperceptible detail만 처리하기 때문에 큰 문제가 없다. 따라서 $T$가 큰 경우에는 $\sigma_t$를 어떻게 설정하든 sample quality에 큰 영향을 미치지 않는다.

그러나 log-likelihood의 측면에서 보면 그렇지 않다. 왜냐하면 위 그래프에서 볼 수 있듯이 $t=0$ 부근의 일부 diffusion step이 log-likelihood에 큰 영향을 미치기 때문이다. 따라서 sample quality 측면에서는 $\sigma_t$를 어떻게 설정하든 큰 차이가 없지만, log-likelihood 측면에서는 $\sigma_t$를 학습하는 것이 더 좋다.

문제는 $\sigma_t$를 학습하기 어렵다는 것이다. DDPM에서도 $\Sigma_{\theta}(x_t, t)$를 직접적으로 학습하려 했지만, 학습이 잘 되지 않았다. 따라서 이 논문에서는 다음과 같은 interpolation 방식의 parameterization을 사용한다.

\[\Sigma_{\theta}(x_t, t) = \exp(v \log \beta_t + (1-v) \log \tilde{\beta}_ t)\]

그런데 여기서 learning parameter $v$의 값에 제한을 두지 않았다는 점이 주목할 만하다. 일반적으로 interpolation이면 $0 \leq v \leq 1$이지만, 여기서는 $v$의 값이 어떤 값이든 상관없다. 그래도 실제로는 $v$의 값이 잘 조절되었다고 한다.

그리고 parameter $v$를 학습하기 위해서는 $L_{\text{simple}}$에 추가적인 term을 더해야 한다. 따라서 원 DDPM에서 유도하였던 $L_{\text{vlb}}$를 도입했다. $L_{\text{vlb}}$는 다음과 같다.

\[\begin{aligned} L_{\text{vlb}} &= L_0 + L_1 + \cdots + L_{T-1} + L_T \\ L_0 &= - \log p_\theta (x_0 \vert x_1) \\ L_{t-1} &= D_{KL} (q(x_{t-1} \vert x_t, x_0) \Vert p_\theta (x_{t-1} \vert x_t)) \\ L_T &= D_{KL} (q(x_T \vert x_0) \Vert p(x_T)) \\ \end{aligned}\]

두 objective를 더해 다음과 같은 hybrid objective를 사용했다. 이때 $L_{\text{simple}}$에 큰 영향을 주지 않기 위해 $\lambda = 0.001$을 사용하였고, $L_{\text{vlb}}$에 의해 $\mu_{\theta}(x_t, t)$는 학습되지 않도록 했다. 결국 $L_{\text{simple}}$에 의해 $\mu_{\theta}(x_t, t)$만 학습되고, $L_{\text{vlb}}$에 의해 $\Sigma_{\theta}(x_t, t)$만 학습되도록 했다.

\[L_{\text{hybrid}} = L_{\text{simple}} + \lambda L_{\text{vlb}}\]

1.2. Improving the Noise Schedule

DDPM에서는 linear noise schedule을 사용하였다. 즉 $\beta_0$와 $\beta_T$를 정하고 그 사이를 linear interpolation하여 사용하였다. 그러나 이는 high resolution image에서는 괜찮지만, $64 \times 64$, $32 \times 32$ 같은 low resolution image에서는 잘 작동하지 않았다. 저자들이 이유를 분석한 결과, diffusion step이 큰 경우 끝 부분에서의 이미지가 이미 너무 noisy하여 끝 부분의 forward process는 sample quality에 영향을 줄 수 없기 때문이라고 추측했다.

실제로 ImageNet $64 \times 64$에서 reverse diffusion의 앞쪽을 제거하자 FID가 크게 감소하였다. 즉 뒷부분(끝 부분의 diffusion step)은 크게 영향을 미치지 못하고 앞부분에 sample quality가 의존하는 것을 보여준다. 이를 해결하기 위해 cosine noise schedule을 사용하였다.

식은 다음과 같다. 여기서 small offset $s$는 $t=0$ 근처에서 variance $\beta_t$가 너무 작은 경우 모델이 $\epsilon$을 잘 예측하지 못하는 문제가 있어 추가하였다고 한다. 그리고 $t=T$ 근처에서 $\beta_t \leq 0.999$이도록 clipping하였다. $\beta_t$는 DDPM과 같이 $1 - \frac{\bar{\alpha}_ t}{\bar{\alpha}_ {t-1}}$로 계산할 수 있다.

\[\bar{\alpha}_ t = \frac{f(t)}{f(0)} \quad f(t) = \cos \left( \frac{t/T + s}{1 + s} \cdot \frac{\pi}{2} \right) ^ 2\]

결과로 Figure 3에서 noise가 더 천천히 부여되는 모습을 볼 수 있고, Figure 4에서 앞쪽 diffusion step 의존도도 낮아진 것을 알 수 있다.

1.3. Reducing Gradient Noise

저자들은 $L_{\text{vlb}}$를 도입하여 실험을 진행했으나 실제로는 최적화가 어렵다는 문제를 발견하였다. $L_{\text{vlb}}$의 gradient가 너무 noisy하기 때문이다. 이러한 noise를 줄이기 위해 importance sampling을 도입하였다. Loss가 높은 term일수록 더 많이 sampling하여 학습하는 것이다.

\[L_{\text{vlb}} = \mathbb{E}_ {t \sim p_t} \left[ \frac{L_t}{p_t} \right] \text{ where } p_t \propto \sqrt{\mathbb{E} \left[ L_t ^2 \right] }, \sum p_t =1\]

이 결과를 $L_{\text{vlb}}$ (resampled)로 나타내었고 훨씬 안정적으로 학습이 진행됨을 알 수 있다. 그러나 덜 noisy한 $L_{\text{hybrid}}$에 이러한 resampling을 도입할 때는 별로 도움이 되지 않아 사용하지 않았다고 한다. 결국 ablation study 에서 $L_{\text{vlb}}$를 사용할 때만 사용한 것이다.

1.4. Results and Ablations

결과를 보면 $L_{\text{hybrid}}$ 및 cosine schedule을 사용하면 FID는 비슷하게 유지하면서 log-likelihood를 개선한 것을 알 수 있다. 그리고 $L_{text{vlb}}$만 사용한 경우 FID는 높아지지만 log-likelihood는 개선되었다. Sample quality를 희생하고 log-likelihood를 개선하는 것보다는 sample quality를 유지하면서 log-likelihood를 개선하는 것이 더 중요하기에 앞으로는 $L_{\text{hybrid}}$를 사용할 것이다.

아래는 다른 likelihood-based model과 NLL(negative log-likelihood)을 비교한 결과이다. IDDPM은 다른 CNN 기반 모델보다는 낮은 NLL을 보여주었지만, Transformer 기반 모델보다는 높은 NLL을 보여주었다.

2. Improving Sampling Speed

DDPM은 sampling 속도가 느리다는 단점이 있다. 이를 해결하기 위해 더 짧은 sampling 단계를 만들어 sampling을 진행하였다. 이는 DDIM과 어느 정도 유사하긴 하지만, 실제로는 deterministic process를 고안한 DDIM과 달리 DDPM의 방식을 거의 그대로 가져온 것이다. 즉 저자들은 새로운 sampling 방식을 고안한 것이 아니며, 오로지 더 짧은 sampling 단계를 사용했을 때 IDDPM이 DDPM보다 우수하다는 것을 보이기 위해 실험을 진행한 것이다.

그 결과 DDPM의 경우에 비해 IDDPM의 objective $L_{\text{hybrid}}$가 짧은 sampling step에서도 FID를 잘 유지하는 것을 볼 수 있다. 그리고 DDIM과도 비교한 결과 $S=50$ 이하의 경우에는 DDIM이 더 좋지만, 그 이상에서는 IDDPM이 더 좋음을 보였다. 다만 $L_{\text{hybrid}}$를 사용하여 훈련하고 DDIM으로 sampling한 경우도 비교해줬으면 좋았을 것 같아 아쉽다. 어쨌든 결론은 sampling step이 짧아도 IDDPM은 DDPM과 달리 sample quality를 유지할 수 있다는 것이다. 이유는 variance learning 때문일 것이다.

3. Comparison to GANs

SOTA인 GAN들과도 비교했는데, GAN은 보통 class-conditional하기에 IDDPM도 class-conditional하게 만들기 위해 class embedding을 timestep embedding에 더해 injection하는 방식을 택하였다. 약식 버전의 CFG(Classifier Free Guidance)가 도입된 셈이다. 여기서는 precision과 recall을 비교했다.

생성모델에서 precision은 fidelity, recall은 diversity를 의미한다고 생각하면 된다. 위 그림을 보면 직관적으로 이해될 것이다.

GAN보다 IDDPM의 precision이 낮긴 하지만 recall은 높다. 이는 IDDPM 모델이 diversity를 높게 유지한다, 즉 데이터 분포의 mode를 커버하는 능력이 뛰어나다고 볼 수 있다. 결국 초반부에 언급했던 데이터 분포의 mode에 대한 문제를 GAN과의 비교를 통해 해결했음을 보여주는 것이라 이해할 수 있다.

4. Scaling Model Size

마지막으로 모델이 scalable하다는 것을 보여주고자 한다. 지금까지 사용한 모델의 first layer depth, 즉 channel 수는 128이었는데, 이를 64부터 192까지 설정하여 실험을 진행했다. 이때 한 가지 디테일은 channel이 늘어나면 weight initialization 시 weight이 작아지므로 learning rate도 다음과 같이 조정했다는 점이다.

\[\tau = (1 / \sqrt{\text{channel multiplier}}) \tau_{128} \quad \text{where } \tau_{128} = 0.0001\]

결과를 보면 FID와 NLL 모두 scalable하다. 모델 크기가 커지면 성능이 커지는 직관적인 결과를 얻은 것이다.

💡 Summary

DDPM에서 낮은 log-likelihood 문제를 variance learning, variance cosine scheduling을 통해 해결하여 sample quality는 유지하면서 log-likelihood를 개선하였다.
Sampling step을 줄여도 DDPM과는 달리 sample quality가 크게 떨어지지 않았으며, GAN과의 비교를 통해 log-likelihood 개선 결과 mode diversity가 충분한 수준임을 보여주었다.
IDDPM의 모델 크기가 커지는 것에 따라 sample quality, log-likelihood 모두 개선되면서 모델이 scalable하다는 것을 보여주었다.
Opinion: DDIM이 등장한 이후의 논문임에도 불구하고 sampling method에 DDIM을 적용한 IDDPM 결과를 보여주지 않은 것이 아쉽다. 또한 SOTA GAN과의 비교에서는 $T=4000$이라는 긴 sampling step을 사용하였기에 diffusion model이 확실한 SOTA가 되기 위해서는 더 짧은 sampling step에서도 좋은 성능을 보여줄 수 있도록 추후 연구가 필요해 보인다.
Review: 본 논문은 ICLR 2021에서 reject되었다. 이유는 FID가 별로 개선되지 않았다는 점, log-likelihood의 개선이 실제 generative task에서 크게 중요하지 않다는 점, FID와 NLL의 trade-off를 충분히 설명하지 못했다는 점, 다른 모델들 및 데이터셋에 대한 실험 부족 등이 있었다.

I'm rubatoyeong