DDIM (Denoising Diffusion Implicit Models)

📄 관련 논문: [ICLR 2021] Denoising Diffusion Implicit Models

헷갈릴 수 있는 용어와 수식에 대해서 먼저 정리하겠다.

생성 모델 논문에서 $q$로 표기한 분포는 수학적으로 계산한 진짜 확률 분포 (Ground Truth), $p$로 표기한 분포는 신경망이 예측한 확률 분포를 의미한다.

• Forward Process (Transition Distribution)

  \[q(\vx_t \mid \vx_{t-1})\]

  이전 상태 $\vx_{t-1}$에 약간의 노이즈를 추가하여 다음 상태인 $\vx_t$를 만드는 수학적 규칙을 의미한다. (DDPM의 Markov 과정 기준)

• Marginal Distribution

  \[q(\vx_t \mid \vx_0)\]

  원본 $\vx_0$에서 $\vx_{t-1}$까지의 중간 단계를 전부 생략하고, 한 번의 계산으로 단번에 $t$ 시점의 상태 $\vx_t$를 만들어내는 분포를 의미한다.

• Forward Process Posterior

  \[q(\vx_{t-1} \mid \vx_t, \vx_0)\]

  최종 원본 이미지 $\vx_0$가 무엇인지 이미 알고 있을 때, $\vx_t$에서 노이즈를 한 step 제거한 $\vx_{t-1}$이 될 수학적인 진짜 확률 (정답 분포)을 의미한다.

  노이즈를 제거하는 역방향의 수식임에도 이름에 Forward가 들어가 있는데, 이는 진행 방향이 정방향이라는 뜻이 아니다.
  노이즈를 주입하는 Forward Process의 수학적 공식을 베이즈 정리를 통해 뒤집어서 얻어낸 사후 확률(Posterior)이라는 태생적 의미를 나타낸 것이다.
  즉, 모델이 백지에서부터 무언가를 새롭게 생성하는 능동적인 과정이 아니라, 이미 알고 있는 원본 $\vx_0$를 바탕으로 Forward Process의 수학적 궤적을 단순히 되감기 한 정답 분포를 의미한다.

• Reverse Process (Generative Process)

  \[p_\theta(\vx_{t-1} \mid \vx_t)\]

  원본 $\vx_0$를 전혀 모르는 상태에서, 노이즈 낀 이미지 $\vx_t$만 보고 신경망이 $\vx_{t-1}$를 예측하는 실제 생성 과정을 의미한다.

DDPM의 한계

DDPM은 $x_{t-1}$을 만들기 위해 반드시 바로 직전 단계인 $x_t$ 정보만 이용하는 Markov Chain을 가정한다.

즉, 하나의 이미지를 생성할 때 Forward process에서 거쳐온 $T$단계를 그대로 한 스텝씩 다 밟아야 하기 때문에, 이미지 생성 속도가 매우 느리다는 단점이 있다.

DDIM의 가정

아래의 단순화된 DDPM 손실 함수를 보면, 손실 함수를 계산하기 위해 $t$ 시점에서의 이미지 $\vx_t$가 필요하다는 것을 확인할 수 있다.

\[L_{\text{simple}}= \mathbb{E}_{t,\mathbf x_0,\boldsymbol\epsilon} \bigg[\lVert\boldsymbol\epsilon-\boldsymbol\epsilon_\theta(\vx_t,t)\rVert^2\bigg] \tag{1}\]

하지만 DDPM에서 증명된 성질에 의해, $\vx_t$는 아래의 수식으로 한 번에 계산된다.

\[\vx_t=\sqrt{\bar{\alpha}_t}\vx_0+\sqrt{1-\bar{\alpha}_t}\boldsymbol\epsilon \tag{2}\]

즉, 손실 함수를 계산하기 위해 중간 단계의 이미지들 $\lbrace\vx_1,\vx_2, \dots,\vx_{t-1}\rbrace$이 전혀 필요하지 않다.

결과적으로 DDIM 저자들은 Marginal 분포 $q(\vx_t\mid\vx_0)$만 DDPM과 똑같이 유지한다면, $\vx_0$에서 $\vx_t$로 가는 중간 경로 (Joint 분포)는 굳이 Markov Chain일 필요가 없다고 가정하였다.

즉, Non-Markovian process를 도입해도 이미 학습된 모델을 그대로 쓸 수 있다는 뜻이다.

DDIM Process

Non-Markovian Forward Process

DDIM에서는 기존 DDPM와 동일한 Marginal 분포 $q(\vx_t\mid\vx_0)=\mathcal{N}(\sqrt{\bar{\alpha}_t}\vx_0, (1-\bar{\alpha}_t)\vI)$를 유지하면서도, Markov 가정을 깨뜨린 새로운 Joint 분포를 정의하였다.

\[q_\sigma(\vx_{1:T} \mid \vx_0) := q_\sigma(\vx_T \mid \vx_0) \prod_{t=2}^{T} q_\sigma(\vx_{t-1} \mid \vx_t, \vx_0) \tag{3}\]

여기서 아래 첨자 $\sigma=[\sigma_1,\dots,\sigma_T]=\in\mathbb{R}^T$는 각 step의 stochasticity를 조절하는 표준편차 벡터를 의미한다.

위의 식에서 모든 $t$에 대한 Marginal 분포 $q_\sigma(\vx_t \mid \vx_0)$를 DDPM과 동일하게 유지하기 위해, $q_\sigma(\vx_{t-1} \mid \vx_t, \vx_0)$는 아래와 같은 가우시안 분포로 유도된다.

\[q_\sigma(\vx_{t-1} \mid \vx_t, \vx_0) = \mathcal{N}\bigg( \sqrt{\bar\alpha_{t-1}}\, \vx_0 + \sqrt{1 - \bar\alpha_{t-1} - \sigma_t^2} \cdot \frac{\vx_t - \sqrt{\bar\alpha_t} \vx_0}{\sqrt{1 - \bar\alpha_t}}, \sigma_t^2 \vI \bigg) \tag{4}\]

즉 $\vx_{t-1}$을 결정할 때 현재 상태 $\vx_t$뿐만 아니라 원본 $\vx_0$의 정보를 직접 참조하며, 만약 $\sigma_t=0$으로 설정하면 확률적 요소가 사라져 $\vx_t$에서 $\vx_{t-1}$로 가는 경로가 고정된다.

이제 2개의 Marginal 분포와 식 (4)를 이용해 베이즈 정리를 적용하면, 노이즈를 주입하는 Forward Process의 한 step을 다음과 같이 유도할 수 있다.

\[q_\sigma(\vx_t \mid \vx_{t-1}, \vx_0) = \frac{q_\sigma(\vx_{t-1} \mid \vx_t, \vx_0) \, q_\sigma(\vx_t \mid \vx_0)}{q_\sigma(\vx_{t-1} \mid \vx_0)} \tag{5}\]

수식을 보면 $\vx_t$가 $\vx_{t-1}$뿐만 아니라 $\vx_0$에도 의존하므로, Forward Process는 더이상 Markovian이 아니다.

Reverse Process (Sampling)

반면, DDIM의 Reverse Process는 여전히 Markovian이다.

헷갈릴 수 있는 개념

Markov는 반드시 1칸씩 촘촘하게 가야 한다는 뜻이 아니다.
다음 목적지를 정할 때, 과거를 돌아보지 않고 오직 현재 위치만 본다는 뜻이다.

DDPM의 Markov는 1000 $\to$ 999 $\to$ 998… 이렇게 1칸씩 뛴다.
즉, 999에서 998로 뛸 때 현재 위치인 999만 본다.

반면 DDIM의 Markov는 1000 $\to$ 900 $\to$ 800… 이렇게 여러 칸씩 뛴다.
DDIM 역시 900에서 800으로 뛸 때, 과거는 보지 않고 오직 현재 위치인 900만 보고 800으로 뛰기 때문에 Markov 성질을 만족한다.

DDIM은 샘플링 때 아래의 수식을 이용해 노이즈를 매 step 제거해 나가면서, 최종적으로 깨끗한 이미지를 생성한다.

\[\vx_{t-1} = \sqrt{\bar\alpha_{t-1}} \underbrace{ \left(\frac{\vx_t - \sqrt{1 - \bar\alpha_t} \, \vepsilon_\theta^{(t)}(\vx_t,t) }{\sqrt{\bar\alpha_t}}\right) }_{\hat{\vx}_0\text{ (predicted } \vx_0)} + \underbrace{ \sqrt{1 - \bar\alpha_{t-1} - \sigma_t^2} \cdot \vepsilon_\theta^{(t)}(\vx_t,t) }_{\text{direction pointing to } \vx_t} + \underbrace{\sigma_t \vepsilon_t}_{\text{random noise}} \tag{6}\]

각 항의 의미는 다음과 같다.

• $\hat{\vx}_0$ (predicted $\vx_0$): 모델이 현재 step의 노이즈 이미지 $\vx_t$를 보고, 한 번에 노이즈를 제거하여 예측한 가상의 원본 이미지
• Direction pointing to $\vx_t$: 예측된 원본 $\hat{\vx}_0$에서 다시 현재 상태 $\vx_t$를 향하는 방향을 나타냄
• Random noise: 도착한 $\vx_{t-1}$에서 약간의 무작위성을 부여함 (생성 결과를 다양하게 하는 역할)

식 (6)을 보면, 랜덤한 부분이 $\sigma_t \vepsilon_t$밖에 없기 때문에 $\sigma_t$ 값에 따라서 생성 과정이 결정된다.

$\sigma_t=\sqrt{\frac{1-\bar\alpha_{t-1}}{1-\bar\alpha_t}}\sqrt{1-\frac{\bar\alpha_t}{\bar\alpha_{t-1}}}$으로 설정하면 DDPM과 동일한 확률적 경로를 따르며, $\sigma_t=0$인 경우 랜덤 노이즈 항이 사라져 노이즈와 결과물이 1:1로 매핑되는 결정론적 (Deterministic) 모델이 된다.

DDPM vs DDIM (가속 sampling)

DDIM은 학습할 때 사용했던 1000 step을 모두 밟을 필요 없이, 전체 step 중 일부만 골라 건너뛰며 샘플링해도 좋은 품질의 이미지를 생성할 수 있다.

이렇게 step을 건너뛸 수 있는지 여부는 DDPM과 DDIM이 수식을 전개한 순서에서 갈린다.

DDPM의 전개 순서 (정방향 $\to$ 역방향)

DDPM은 자연계의 확산 과정을 모방하기 위해, 1step의 노이즈 주입 규칙부터 고정해 두고 시작했다.

1. Forward Process 정의: 노이즈를 1step 주입하는 규칙 $q(\vx_t \mid \vx_{t-1})$을 가우시안 분포를 따르는 Markov Chain으로 정의함
2. Marginal Distribution 도출: 수식을 전개해 보니, 임의의 $t$ 시점에서의 상태 $q(\vx_t \mid \vx_0)$도 가우시안 분포로 도출됨

3. Forward Process Posterior 도출: 베이즈 정리를 이용해, Reverse Process의 정답 분포 $q(\vx_{t-1} \mid \vx_t, \vx_0)$를 아래와 같은 가우시안 분포의 형태로 유도함

   \[q(\vx_{t-1} \mid \vx_t, \vx_0)=\mathcal{N}(\vx_{t-1}; \tilde{\vmu}_t(\vx_t, \vx_0), \tilde{\beta}_t \vI)\]

DDPM에서는 처음부터 1step 노이즈 주입을 가우시안 분포로 고정해버린 탓에, 역방향 경로 또한 가우시안 분포를 유지하려면 매 step의 노이즈 보폭이 충분히 작아야 한다는 제약이 생겼다.
즉, 역방향 수식 역시 단일 스텝 계수인 $\alpha_t$에 묶여버렸기 때문에,
따라서 step을 건너뛰려 하면 가우시안 전제가 깨지면서 모델이 경로를 이탈하게 되므로, 무조건 $T$ step을 다 밟아야만 한다.

DDIM의 전개 순서 (역방향 $\to$ 정방향)

DDIM은 처음부터 중간 과정을 건너뛴다는 목적을 가지고 결과를 미리 고정한 뒤 과정을 역으로 끼워 맞췄다.

1. Marginal Distribution 고정: 기존 학습된 가중치를 그대로 재사용하기 위해, DDPM에서 쓰던 Marginal 분포 $q(\vx_t \mid \vx_0)$를 목적지로 설정함
2. Forward Process Posterior 설계: Marginal 분포를 유지하면서 Non-Markovian 목표를 만족하기 위해, Reverse Process의 정답 분포 $q(\vx_{t-1} \mid \vx_t, \vx_0)$가 어떤 모양이어야 하는지 수학적으로 역산해서 식 (4)를 도출함
3. Forward Process 확인: Marginal 분포 $q(\vx_t \mid \vx_0)$, $q(\vx_{t-1} \mid \vx_0)$ 2개와 식 (4)를 베이즈 정리에 대입해 보니, 새롭게 정의한 확률적 프레임워크 내에서는 노이즈를 주입하는 과정이 식 (5) 같은 Non-Markovian 형태를 띠게됨을 사후적으로 Forward Process를 확인함

DDIM의 Reverse Process 수식 (6)을 보면, 개별 step인 $\alpha_t$가 아니라 누적 step 계수인 $\bar{\alpha}_t$만으로 구성되어 있다.
덕분에 100 step 앞의 $\bar{\alpha}_{t-100}$ 값을 수식에 밀어 넣어도 분포가 깨지지 않는다.

DDIM Inversion

DDIM Inversion은 간단히 말해 이미지를 다시 초기 노이즈 상태로 되돌리는 과정이다.

DDPM은 샘플링 과정에서 매 step마다 랜덤 노이즈를 주입하기 때문에, 같은 시작점에서 출발해도 다른 결과가 나온다.
즉, 이미지를 만드는 과정에 우연성이 개입되어 있어 원래의 초기 노이즈를 정확히 역추적하는 것이 불가능하다.

반면, DDIM은 $\sigma_t=0$일 때 결정론적 경로를 따른다.
따라서 이미지가 생성되는 과정을 수학적으로 정확히 거슬러 올라가, 원본 이미지를 생성했던 초기 노이즈를 완벽하게 찾아낼 수 있다.