[21’ ICLR] DDIM: Denoising Diffusion Implicit Models

Date: 2024.03.26 Updated: 2024.03.26

카테고리: Generation

0. Abstract

DDPM을 확장하여, non-markovian diffusion process를 사용하는 DDIM 모델을 제안한다.
본 모델은 deterministic generative process를 가지고 있어 sampling 과정이 DDPM보다 10~50배 빠르다.

1. Introduction

DDPM과 같은 diffusion-based generative model은 sampling 과정이 매우 느리다. 예를 들어 DDPM에서 $32 \times 32$ 이미지 50,000개를 생성하기 위해서는 20시간 정도가 걸리지만 동일한 환경에서 GAN을 사용하면 1분도 걸리지 않는다. 이러한 문제를 해결하기 위해 DDIM(Denoising Diffusion Implicit Models)을 제안한다.

DDIM에서는 non-markovian diffusion process를 사용하는데, 이를 사용하면 sampling 과정이 deterministic하게 진행되어 더 짧은 generative Markov chains를 거치게 되고, 따라서 sampling 속도가 빨라진다. 그리고 적절한 유도를 통해 DDPM과 완전히 동일한 학습 목표(objective)를 가진 채 학습이 가능하다. 즉 DDPM으로 학습하고 DDIM으로 샘플링하는 것이 가능하다.

DDIM은 DDPM에 비해 몇 가지 이점을 가진다. 먼저 샘플의 퀄리티가 좋고, 더 빠르게 샘플링이 가능하다. 그리고 sampling 과정이 deterministic하기에 동일한 latent variable로 시작한다면 유사한 high-level feature를 가지는 이미지를 생성해낸다. DDPM의 경우 sampling 과정이 stochastic하기 때문에 이러한 특징을 가지지 못한다. 이를 활용하면 DDIM에서는 initial latent variable로 interpolation을 수행할 수 있다. DDPM에서는 이것이 불가능했다.

2. Background

DDPM과 notation이 다른 관계로 식만 간단하게 정리하자.

DDPM의 reverse process를 다음과 같이 나타낼 수 있다.

\[p_\theta (x_ 0) := \int p_\theta (x_ {0:T}) dx_ {1:T} \\ p_\theta (x_ {0:T}) := p(x_ T) \prod _ {t=1} ^T p_\theta^{(t)}(x_ {t-1} \vert x_ t)\]

그리고 forward process는 다음과 같이 나타낼 수 있다.

\[q(x_ {1:T} \vert x_ 0) := \prod _ {t=1} ^T q(x_ t \vert x_ {t-1}) \\ q(x_ t \vert x_ {t-1}) := \mathcal{N}(x_ t ; \sqrt{1-\beta_t} x_ {t-1}, \beta_t \mathbf{I})\]

이때 $\alpha_t$를 다음과 같이 정의하면, (DDPM에서의 $\bar{\alpha}_ t$와 동일) $q(x_ t \vert x_ {t-1})$를 다음과 같이 나타낼 수 있다.

\[\alpha_ t := \prod _ {s=1} ^t (1-\beta_s) \\ q(x_ t \vert x_ {t-1}) := \mathcal{N}(x_ t ; \sqrt{\frac{\alpha_ t}{\alpha_{t-1}}} x_ {t-1}, \left(1 - \frac{\alpha_ t}{\alpha_{t-1}}\right) \mathbf{I})\]

Forward process의 특징은 다음과 같다.

\[q(x_ t \vert x_ 0) = \mathcal{N}(x_ t ; \sqrt{\alpha_ t} x_ 0, (1-\alpha_ t) \mathbf{I})\]

Reparametrization trick을 사용하면 다음과 같이 나타낼 수 있다.

\[x_ t = \sqrt{\alpha_ t} x_ 0 + \sqrt{1-\alpha_ t} \epsilon \quad \text{where} \quad \epsilon \sim \mathcal{N}(0, \mathbf{I})\]

Simplified objective는 다음과 같다. 실제로 DDPM에서는 $\gamma = 1$로 설정하고 학습을 진행했다.

\[L_ \gamma (\epsilon_ \theta) := \sum _ {t=1} ^T \gamma_t \mathbb{E}_ {x_ 0 \sim q(x_ 0), \epsilon \sim \mathcal{N}(0, \mathbf{I})} \left[ \left\Vert \epsilon_ \theta ^{(t)}(\sqrt{\alpha_ t} x_ 0 + \sqrt{1-\alpha_ t} \epsilon) - \epsilon_t \right\Vert ^2 \right]\]

만약 $\alpha_T \rightarrow 0$이면, $q(x_ T \vert x_ 0) $는 standard Gaussian에 수렴하고 따라서 $p_\theta(x_ T) := \mathcal{N}(0, \mathbf{I})$로 두는 것이 자연스럽다. 그러나 이것은 forward process의 길이 $T$가 충분히 길 때만 가능하다. 따라서 DDPM의 sampling은 필연적으로 오래 걸릴 수밖에 없다.

3. Variational Inference for non-Markovian Forward Process

DDIM에서의 key idea는 DDPM objective $L_ \gamma$가 joint distribution $q(x_{1:T} \vert x_0)$에 의존하는 것이 아니라 marginal distribution $q(x_t \vert x_0)$에 의존한다는 것이다. 따라서 동일한 marginal distribution을 가지는 다른 joint distribution을 사용하면 동일한 objective를 가지면서도 다른 방법으로 sampling을 할 수 있다.

3.1. Non-Markovian forward process

따라서 DDIM에서는 그 하나의 방법으로 non-Markovian forward process를 사용한다. 즉, 아래 그림과 같이 $x_t$가 $x_{t-1}$과 $x_0$에 의존한다고 생각하는 것이다.

한편 reverse process는 다음과 같이 정의할 수 있다. 이때 $\mathcal{Q}$는 inference distribution family이며, real vector $\sigma \in \mathbb{R}^ T _ {\geq 0}$에 의해 forward process의 stochasticity를 조절한다. 즉, 이제 $x_t$로부터 $x_{t-1}$을 얻는 Markovian process가 아니라 $x_t$및 $x_0$로부터 $x_{t-1}$을 얻는 non-Markovian process이기에 더 deterministic하다고 이해하면 된다. 예로 만약 $\sigma \rightarrow 0$이면 $x_0$과 $x_t$를 아는 상황에서 $x_{t-1}$은 고정된다. 자세한 내용은 후술하겠다.

\[q_\sigma (x_ {1:T} \vert x_ 0) := q_\sigma (x_T \vert x_0) \prod _ {t=2} ^T q_\sigma (x_ {t-1} \vert x_ t, x_ 0) \\ q_\sigma (x_ T \vert x_ 0) = \mathcal{N}(\sqrt{\alpha_ T} x_ 0, (1-\alpha_ T) \mathbf{I}) \\\]

DDPM과 동일한 marginal을 가지기 위해서는 다음과 같은 분포여야 한다. 즉, 아래와 같은 식을 만족할 때 $q_\sigma(x_t \vert x_0)$는 $\mathcal{N}(\sqrt{\alpha_t} x_0, (1-\alpha_t) \mathbf{I})$가 된다. 유도 과정은 본 논문의 Appendix에 자세히 나와있다. 여기서 증명할 것은 아니기에 다음 식을 받아들이자. 이러한 정의를 통해 objective를 동일하게 맞출 수 있다는 것이 중요하다.

\[q_\sigma (x_ {t-1} \vert x_ t, x_ 0) = \mathcal{N}(\sqrt{\alpha_ {t-1}} x_0 + \sqrt{1-\alpha_ {t-1} - \sigma_ t ^ 2} \cdot \frac{x_ t - \sqrt{\alpha_ t} x_ 0}{\sqrt{1-\alpha_ t}}, \sigma_ t ^ 2 \mathbf{I})\]

Forward process는 Bayes’ rule에 의해 다음과 같이 정해질 수 있다. 역시 Gaussian이 되는데, 별로 중요하지 않다. 어차피 DDPM과 동일한 objective로 학습을 진행할 것이기 때문이다.

\[q_\sigma (x_ t \vert x_ {t-1}, x_ 0) = \frac{q_\sigma (x_ {t-1} \vert x_ t, x_ 0) q_\sigma (x_ t \vert x_ 0)}{q_\sigma (x_ {t-1} \vert x_ 0)}\]

3.2. Generative process and unified variational inference objective

Generative process의 $p_\theta^{(t)}(x_{t-1} \vert x_t)$는 $q_\sigma(x_{t-1} \vert x_t, x_0)$로부터 얻을 수 있다. Generative process에서는 $x_t$만 알고 $x_0$는 모르는 상황이므로 먼저 $x_t$를 이용해 $x_0$를 예측하고, 이를 이용하여 $q_\sigma$로부터 $x_{t-1}$을 샘플링한다. 먼저 $x_0$를 예측하기 위해 reparametrization trick 식을 다시 보면 다음과 같다.

\[x_ t = \sqrt{\alpha_ t} x_ 0 + \sqrt{1-\alpha_ t} \epsilon_t \quad \text{where} \quad \epsilon_t \sim \mathcal{N}(0, \mathbf{I}) \\ \therefore x_ 0 = (x_t - \sqrt{1-\alpha_ t} \epsilon_t) / \sqrt{\alpha_ t}\]

모델 $\epsilon_\theta^{(t)} (x_t)$가 $x_t$로부터 $\epsilon_t$를 예측할 수 있도록 학습한다면, $x_0$의 예측값 $f_\theta^{(t)}(x_t)$는 다음과 같이 나타낼 수 있다. 이를 denoised observation이라 한다.

\[f_\theta^{(t)}(x_t) = (x_t - \sqrt{1-\alpha_ t} \epsilon_\theta^{(t)}(x_t)) / \sqrt{\alpha_ t}\]

이를 $q_\sigma(x_{t-1} \vert x_t, x_0)$에 대입하여 $x_{t-1}$을 샘플링한다. 즉, 다음과 같이 나타낼 수 있다.

\[p_\theta^{(t)}(x_{t-1} \vert x_t) = \begin{cases} \mathcal{N}(f_\theta^{(1)}(x_1), \sigma_ 1 ^ 2 \mathbf{I}) & \text{if } t = 1 \\ q_\sigma(x_{t-1} \vert x_t, f_\theta^{(t)}(x_t)) & \text{otherwise}, \end{cases}\]

이제 DDIM의 objective가 DDPM의 objective와 같아질 수 있음을 보일 것이다. 먼저 variational inference objective를 다시 정리하면 다음과 같다.

Variational inference objective $J_\sigma$의 정의를 보면 $\sigma$가 다른 모델의 경우 서로 다른 objective를 가지기 때문에 각각 $\sigma$에 대해 학습을 진행해야 하는 것처럼 보인다. 하지만 저자들이 보인 Theorem 1에 의해 그럴 필요가 없다.

Theorem 1. 모든 $\sigma \gt 0$에 대하여 $J_\sigma = L_\gamma + C$를 만족하는 $\gamma \in \mathbb{R}^ T _ {\geq 0}, C \in \mathbb{R}$가 존재한다.

만약 DDPM에서 모델 $\epsilon_\theta ^{(t)}$에서 서로 다른 $t$에 대하여 parameter $\theta$를 공유하지 않는다면, $\epsilon_\theta$의 최적값은 $\gamma$와 무관하다. 왜냐하면 $\gamma$와 관계없이 각 항에 대해서 최적화하면 되기 때문이다.

다른 말로 하면 어떤 $\gamma$를 쓰든 최적해는 동일하기 때문에 $L_1$, 즉 DDPM에서 사용한 objective를 사용해도 동일하다. 그런데 $J_\sigma$가 $L_\gamma$와 상수 차이 외에는 동일하고, $L_\gamma$의 최적값과 $L_1$의 최적값이 같으므로 $J_\sigma$ 대신 $L_1$을 사용하여 학습해도 된다. 즉 계속 언급했던 DDPM의 objective를 사용할 수 있다는 것이 정당화되었다. 물론 의문이 드는 점도 있다. 실제로 DDPM 구현에서는 $t$에 따라 parameter $\theta$가 공유되는 것으로 알고 있어 지식 사이의 충돌이 있다. 이 부분은 논문의 구현 파트를 살펴봐야 할 것 같다.

4. Sampling from Generalized Generative Process

위의 결론으로부터 pretrained DDPM을 사용하고, generative process만 DDIM 방식으로 하는 것이 가능하다. 이제 $\sigma$만 바꿔가며 sampling을 수행하면 된다. Sampling은 다음과 같은 방식으로 진행된다.

4.1. Denoising diffusion implicit models

다음과 같은 식으로 $x_t$로부터 $x_{t-1}$을 샘플링할 수 있다.

위 식은 다음 분포로부터 얻은 것이다.

이제 동일한 모델 $\epsilon_\theta$를 사용하여 여러 $\sigma$에 대하여 sampling을 진행할 수 있다. 이때 두 가지 특수한 경우가 있다. 먼저 $\sigma$를 다음과 같이 정의하면 forward process가 Markovian이 되며, 따라서 생성 과정은 DDPM과 동일해진다.

\[\sigma_ t = \sqrt{\frac{1- \alpha_ {t-1}}{1- \alpha_ t}} \sqrt{1 - \frac{\alpha_ t}{\alpha_ {t-1}}}\]

그리고 $\sigma_t =0$이면 forward process가 deterministic하며 생성 과정에서 latent variable로부터 sample이 fixed procedure를 통해 생성된다. 따라서 이 모델을 DDIM(Denoising Diffusion Implicit Models)이라고 부른다. DDPM objective로 훈련된 implicit probabilistic model이라는 뜻이다.

4.2. Accelerated generation processes

이 부분이 DDIM의 하이라이트이다. 이제 forward process를 $x_{1:T}$에 대해 모두 정의하는 것이 아니라 그 부분집합에서 정의한다고 생각해보자. 즉, $[1, 2, \cdots, T]$ 중 길이가 $S$인 sub-sequence $\tau$에 대해 다음과 같은 subset $\lbrace x_{\tau_1}, x_{\tau_2}, \cdots, x_{\tau_S} \rbrace$을 만들고 이것에 대한 forward process를 정의한다. 어차피 forward process는 $q(x_{\tau_i} \vert x_0)$가 이미 정의되어 있으므로 쉽게 가능하고, 이를 거꾸로 하여 reverse process를 정의할 수 있다. 이제 DDPM의 문제였던 길이 $T$의 긴 generative process 대신 짧은 길이 $S$의 generative process로 대체가 가능한 것이다. 실제로는 DDPM 전체에 대해 학습한 뒤에, sampling할 때는 그 중 몇 개의 subset만 사용하여 sampling을 진행할 수 있고, 이를 통해 sampling의 속도를 빠르게 할 수 있다.

4.3. Relevance to Neural ODEs

위 식을 이제 DDIM의 방식으로 $\sigma_t =0$이라 하고 $x_{t-1}$ 대신 $x_{t - \Delta t}$를 샘플링하는 식으로 바꾸어 보기 좋게 정리하면 다음과 같다.

여기서 $(\sqrt{1-\alpha} / \sqrt{\alpha})$를 $\sigma$로, $(x / \sqrt{\alpha})$를 $\bar{x}$로 reparametrize하면 다음과 같은 ODE에서 Euler method를 적용한 것과 같은 형태임을 알 수 있다.

\[d \bar{x} (t) = \epsilon _\theta ^{(t)} \left(\frac{\bar{x} (t)}{\sqrt{\sigma^2 +1}}\right) d \sigma(t)\]

즉, DDPM과 달리 DDIM에서는 generation process와 그 역인 forward process가 자유자재로 변환이 가능하여 $x_0$의 latent representation, 혹은 encoding of observation을 $x_T$의 형태로 얻을 수 있다. 그리고 저자들은 위 ODE가 DDPM의 VE-SDE를 ODE의 형태로 바꾼 probability flow ODE(Song et al., 2020)와 동일함을 보였다. 다음 그림은 이 과정을 직관적으로 나타낸다(CS236n Lecture 14 참조).

5. Experiments

결과를 보기에 앞서 어떤 방식으로 이미지를 생성했는지 간단히 보자. 먼저 DDPM과 동일하게 $T=1000$, $L_1$을 이용하여 훈련한 모델을 사용하였으며 DDIM에서 바뀐 것은 오로지 샘플링 과정이다. 이때 변수는 두 가지로 $\tau$ 즉 얼마나 빠르게 샘플을 얻을 것인지, 그리고 $\sigma$ 즉 deterministic DDIM과 stochastic DDPM 중 어느 비율을 더 늘릴 것인지이다. $\sigma$의 경우 다음과 같이 간단하게 정의하였다.

\[\sigma_ {\tau_i} (\eta) = \eta \sqrt{\frac{1-\alpha_{\tau_i-1}}{1-\alpha_{\tau_i}}} \sqrt{1 - \frac{\alpha_{\tau_i}}{\alpha_{\tau_i-1}}}\]

따라서 $\eta = 0$이면 deterministic DDIM이 되고, $\eta = 1$이면 stochastic DDPM이 된다. 참고로 $\hat{\sigma}$라는 것도 정의하였는데 orginal DDPM에서 CIFAR10 dataset에 대해 더 큰 standard deviation으로 실험했던 결과이다. 즉 $\sigma(1)$보다 더 큰 SD로 실험한 결과라고 생각하면 된다.

5.1. Sample quality and efficiency

Table 1에서는 Sample quality와 computational cost 간의 trade-off를 잘 보여준다. $S = \dim(\tau)$가 커질수록 sample quality는 높아지지만 time step이 많아지므로 computational cost가 증가한다. 한편 $S \leq 100$에서는 $\eta = 0$일 때의 성능이 가장 높은 반면 $S = 1000$에서는 $\hat{\sigma}$의 성능이 가장 좋았다. 즉 충분한 길이의 trajectory에서는 stochastic DDPM의 성능이 좋지만, 짧은 trajectory에서는 deterministic DDIM의 성능이 우수하다는 것이다.

5.2. Sample consistency in DDIMs

DDIM은 deterministic generative process를 가지고 있으므로 initial latent variable $x_T$가 동일하다면 trajectory가 달라도 high-level feature가 유사한 모습을 보여준다. 여기서 trajectory가 다르다는 것은 number of steps가 다르다는 것이다. 이는 interpolation에 유용하게 사용될 수 있다.

5.3. Interpolation in deterministic generative process

따라서 위와 같이 $x_T$로부터 interpolation이 가능하다. 반면 DDPM은 stochastic generative process를 가지고 있으므로 동일한 $x_T$를 가지고 있어도 다양한 $x_0$를 얻게 된다. 따라서 DDPM에서 interpolation을 진행하기 위해서는 아래와 같이 중간 단계의 feature $x_t$를 가져왔어야 했다.

5.4. Reconstruction from latent space

여기서는 reconstruction error를 확인하여 latent space $x_T$에 $x_0$가 잘 mapping되어 있는지 확인하였다. 당연한 결과로 $S$가 클수록 reconstruction error가 작아지는 것을 확인할 수 있었다. 이는 neural ODEs, normalizing flows와 일맥상통하는 결과이다. 그러나 stochastic한 DDPM은 이러한 reconstruction이 불가능하다.

💡 Summary

DDIM은 non-markovian diffusion process를 사용하여 sampling 속도를 빠르게 하였으며, sampling 과정을 deterministic하게 만들었다.
DDIM은 DDPM과 완전히 동일한 학습 목표(objective)를 가진 채 학습이 가능하여 pretrained DDPM을 사용할 수 있다.
Deterministic한 DDIM의 특성으로 latent space에서의 interpolation 및 reconstruction 등이 자유롭게 가능하다.

📃 Reference

[21’ ICLR] DDIM: Denoising Diffusion Implicit Models

I'm rubatoyeong