본문 바로가기

인공지능 논문 정리/Diffusiion

[Diffusion 5] DDPM 논문 완전 정복 (하): Reverse Process와 학습

" The joint distribution $p_θ(x_{0:T})$ is called the reverse process, and it is defined as a Markov chain with learned Gaussian transitions starting at $p(x_T) = N(x_T; 0, I)$. "
 Ho et al., Denoising Diffusion Probabilistic Models (2020)

 

4장에서 우리는 Forward Process의 수학을 완전히 이해했습니다(라고 믿고있습니다.). 데이터 $x_0$가 어떻게 점진적으로 노이즈 $x_T$로 변환되는지, 그 과정의 수학적 구조와 closed-form 표현을 유도했습니다. 이제 핵심 질문이 남았습니다. 어떻게 노이즈에서 다시 데이터를 복원할 수 있을까요? 이번 장에서는 Reverse Process의 수학을 다룹니다. 베이즈 정리를 통해 역방향 전이의 이론적 형태를 유도하고, 신경망이 이를 어떻게 근사하는지 살펴봅니다. 그리고 가장 중요한 질문 "왜 단순히 노이즈를 예측하는 것이 최적의 학습 전략인가?" 에 대한 수학적 답을 제시합니다.

이 장의 목표: Reverse Process의 이론적 형태를 베이즈 정리로 유도합니다. Variational Lower Bound(VLB)를 분해하여 각 항의 의미를 이해합니다. 손실 함수가 어떻게 단순한 노이즈 예측 문제로 귀결되는지 증명합니다. 노이즈 예측, $x_0$ 예측, score function의 세 가지 관점 간의 등가성을 확립합니다.

5.1 Reverse Process의 필요성과 정의

Forward Process는 데이터를 파괴하는 과정입니다. $x_0 \to x_1 \to \cdots \to x_T$의 여정에서 원본 정보는 점차 사라지고, 최종적으로 $x_T \approx \mathcal{N}(0, \mathbf{I})$에 도달합니다. 생성 모델의 목표는 이 과정을 역전시키는 것입니다. 순수한 노이즈 $x_T$에서 시작하여 점진적으로 구조를 복원하고, 최종적으로 실제 데이터 분포에서 온 것 같은 $x_0$를 생성해야 합니다.

5.1.1 역방향 전이의 이론적 형태

놀라운 사실이 있습니다. Forward Process $q(x_t|x_{t-1})$이 충분히 작은 스텝으로 정의되어 있다면, 역방향 전이 $q(x_{t-1}|x_t)$ 또한 가우시안에 가까워집니다. 이 결과는 Feller (1949)에서 처음 증명되었으며, 확산 과정의 시간 역전에 관한 Anderson (1982)의 고전적 연구에서 더욱 정교화되었습니다.

직관적으로 이해해 봅시다. Forward Process의 각 스텝에서 추가되는 노이즈가 작다면($\beta_t$가 작다면), 한 스텝에서의 변화는 미세합니다. 이 미세한 변화는 거의 선형적이고, 따라서 역방향 전이 또한 가우시안으로 근사할 수 있습니다. 이것이 DDPM에서 $T = 1000$처럼 많은 스텝을 사용하는 이유 중 하나입니다.

5.1.2 문제점: $q(x_{t-1}|x_t)$는 계산 불가능

이론적으로 역방향 전이 $q(x_{t-1}|x_t)$가 가우시안에 가깝다는 것을 알지만, 실제로 이를 계산하는 것은 불가능합니다. 베이즈 정리를 적용하면:

$$ q(x_{t-1}|x_t) = \frac{q(x_t|x_{t-1}) q(x_{t-1})}{q(x_t)} $$

여기서 문제가 발생합니다. $q(x_{t-1})$과 $q(x_t)$는 데이터 분포 $q(x_0)$에서 출발한 주변 분포(marginal distribution)입니다:

$$ q(x_t) = \int q(x_t|x_0) q(x_0) \, dx_0 $$

데이터 분포 $q(x_0)$는 미지의 복잡한 분포이므로, 이 적분을 해석적으로 계산할 수 없습니다. 따라서 $q(x_{t-1}|x_t)$를 직접 계산하는 것은 불가능합니다.

5.1.3 해결책: 신경망을 통한 근사

DDPM의 핵심 아이디어는 $q(x_{t-1}|x_t)$를 직접 계산하는 대신, 신경망 $p_\theta(x_{t-1}|x_t)$로 근사하는 것입니다. 역방향 전이가 가우시안에 가깝다는 이론적 결과를 활용하여, 신경망의 출력을 가우시안의 파라미터로 설정합니다:

$$ p_\theta(x_{t-1}|x_t) = \mathcal{N}(x_{t-1}; \mu_\theta(x_t, t), \Sigma_\theta(x_t, t)) $$

신경망 $\mu_\theta$는 주어진 $x_t$와 시간 $t$에서 역방향 전이의 평균을 예측하고, $\Sigma_\theta$는 분산을 예측합니다. Ho et al. (2020)의 원본 DDPM 논문에서는 분산을 학습하지 않고 고정된 값 $\sigma_t^2 \mathbf{I}$를 사용하는 단순화를 채택했습니다. 이 선택의 이유는 후에 자세히 설명합니다.

5.1.4 Reverse Process의 완전한 정의

Reverse Process는 Forward Process와 마찬가지로 마르코프 체인으로 정의됩니다:

$$ p_\theta(x_{0:T}) = p(x_T) \prod_{t=1}^{T} p_\theta(x_{t-1}|x_t) $$

여기서 $p(x_T) = \mathcal{N}(x_T; 0, \mathbf{I})$는 순수한 가우시안 노이즈입니다. 샘플링 과정은 다음과 같습니다:

1. $x_T \sim \mathcal{N}(0, \mathbf{I})$에서 순수 노이즈를 샘플링합니다.
2. $t = T, T-1, \ldots, 1$에 대해 $x_{t-1} \sim p_\theta(x_{t-1}|x_t)$를 반복적으로 샘플링합니다.
3. 최종적으로 $x_0$를 얻습니다.

그림 5.1.  Forward Process와 Reverse Process의 대칭 구조. Forward는 데이터에서 노이즈로, Reverse는 노이즈에서 데이터로 진행합니다.

5.2 Posterior $q(x_{t-1}|x_t, x_0)$의 유도

$q(x_{t-1}|x_t)$는 직접 계산할 수 없지만, 원본 데이터 $x_0$가 주어진 조건부 분포 $q(x_{t-1}|x_t, x_0)$는 해석적으로 유도할 수 있습니다. 이 분포는 신경망 학습의 목표를 설정하는 데 핵심적인 역할을 합니다.

5.2.1 베이즈 정리 적용

베이즈 정리를 적용하여 $q(x_{t-1}|x_t, x_0)$를 유도합니다:

$$ q(x_{t-1}|x_t, x_0) = \frac{q(x_t|x_{t-1}, x_0) \cdot q(x_{t-1}|x_0)}{q(x_t|x_0)} $$

Forward Process의 마르코프 성질에 의해 $q(x_t|x_{t-1}, x_0) = q(x_t|x_{t-1})$입니다. 4장에서 유도한 결과들을 대입하면:

$$ q(x_t|x_{t-1}) = \mathcal{N}(x_t; \sqrt{\alpha_t} \, x_{t-1}, (1-\alpha_t)\mathbf{I}) $$ $$ q(x_{t-1}|x_0) = \mathcal{N}(x_{t-1}; \sqrt{\bar{\alpha}_{t-1}} \, x_0, (1-\bar{\alpha}_{t-1})\mathbf{I}) $$ $$ q(x_t|x_0) = \mathcal{N}(x_t; \sqrt{\bar{\alpha}_t} \, x_0, (1-\bar{\alpha}_t)\mathbf{I}) $$

여기서 $\alpha_t = 1 - \beta_t$이고 $\bar{\alpha}_t = \prod_{s=1}^{t} \alpha_s$입니다.

5.2.2 가우시안 분포의 곱

세 분포 모두 가우시안이므로, 3장에서 배운 가우시안의 곱 공식을 적용할 수 있습니다. 계산의 핵심 아이디어는 다음과 같습니다. 가우시안 분포의 로그를 취하면 $x$에 대한 이차식이 됩니다:

$$ \log q(x_{t-1}|x_t, x_0) = \log q(x_t|x_{t-1}) + \log q(x_{t-1}|x_0) - \log q(x_t|x_0) $$

우변에서 $x_{t-1}$에 의존하는 항만 추출하면 됩니다. $\log q(x_t|x_0)$는 $x_{t-1}$을 포함하지 않으므로 상수로 취급합니다.

각 가우시안의 지수부를 전개합니다. $q(x_t|x_{t-1})$에서:

$$ -\frac{1}{2(1-\alpha_t)} \|x_t - \sqrt{\alpha_t} x_{t-1}\|^2 = -\frac{1}{2(1-\alpha_t)} \left( \alpha_t \|x_{t-1}\|^2 - 2\sqrt{\alpha_t} x_t^\top x_{t-1} + \|x_t\|^2 \right) $$

$q(x_{t-1}|x_0)$에서:

$$ -\frac{1}{2(1-\bar{\alpha}_{t-1})} \|x_{t-1} - \sqrt{\bar{\alpha}_{t-1}} x_0\|^2 = -\frac{1}{2(1-\bar{\alpha}_{t-1})} \left( \|x_{t-1}\|^2 - 2\sqrt{\bar{\alpha}_{t-1}} x_0^\top x_{t-1} + \bar{\alpha}_{t-1}\|x_0\|^2 \right) $$

$x_{t-1}$에 대한 이차항과 일차항의 계수를 정리하면, $q(x_{t-1}|x_t, x_0)$가 다음 형태의 가우시안임을 알 수 있습니다:

정리 5.1 (Posterior 분포). Forward Process의 posterior 분포 $q(x_{t-1}|x_t, x_0)$는 다음과 같은 가우시안입니다: $$ q(x_{t-1}|x_t, x_0) = \mathcal{N}(x_{t-1}; \tilde{\mu}_t(x_t, x_0), \tilde{\beta}_t \mathbf{I}) $$ 여기서 posterior 평균과 분산은 다음과 같이 주어집니다: $$ \tilde{\mu}_t(x_t, x_0) = \frac{\sqrt{\alpha_t}(1-\bar{\alpha}_{t-1})}{1-\bar{\alpha}_t} x_t + \frac{\sqrt{\bar{\alpha}_{t-1}} \beta_t}{1-\bar{\alpha}_t} x_0 $$ $$ \tilde{\beta}_t = \frac{(1-\bar{\alpha}_{t-1}) \beta_t}{1-\bar{\alpha}_t} $$

5.2.3 Posterior 평균의 의미 해석

정리 5.1의 posterior 평균 $\tilde{\mu}_t(x_t, x_0)$를 자세히 살펴봅시다. 이 평균은 $x_t$와 $x_0$의 가중 평균입니다:

$$ \tilde{\mu}_t = \underbrace{\frac{\sqrt{\alpha_t}(1-\bar{\alpha}_{t-1})}{1-\bar{\alpha}_t}}_{\text{$x_t$의 가중치}} x_t + \underbrace{\frac{\sqrt{\bar{\alpha}_{t-1}} \beta_t}{1-\bar{\alpha}_t}}_{\text{$x_0$의 가중치}} x_0 $$

두 가중치의 합이 1에 가깝다는 것은 (정확히 1은 아님) $\tilde{\mu}_t$가 $x_t$와 $x_0$ 사이의 보간(interpolation)에 가깝다는 것을 의미합니다. 시간 $t$가 클 때(Forward Process 후반부), $x_t$의 가중치가 더 크고, $t$가 작을 때는 $x_0$의 가중치가 더 커집니다.

이 결과의 직관적 의미는 명확합니다. Reverse step에서 "한 스텝 전으로 돌아가기" 위해서는, 현재 상태 $x_t$에서 출발하되 원본 데이터 $x_0$ 방향으로 조금 이동해야 합니다. 그리고 그 이동량은 $x_0$를 얼마나 아는가에 의존합니다.

5.2.4 Posterior 분산의 특성

Posterior 분산 $\tilde{\beta}_t$는 시간 $t$에만 의존하고 데이터에는 의존하지 않습니다. 이는 학습해야 할 파라미터가 아니라 노이즈 스케줄 $\{\beta_t\}$에서 미리 결정되는 상수입니다.

$\tilde{\beta}_t$의 크기를 분석해 봅시다. $t$가 작을 때 (Forward Process 초반), $\bar{\alpha}_{t-1} \approx 1$이므로 $\tilde{\beta}_t \approx \beta_t$입니다. $t$가 클 때 (Forward Process 후반), $(1-\bar{\alpha}_{t-1})$이 커지면서 $\tilde{\beta}_t$는 $\beta_t$보다 작아집니다.

Ho et al. (2020)은 실험에서 신경망의 분산을 $\tilde{\beta}_t$ 또는 $\beta_t$로 고정했을 때 비슷한 성능을 얻었다고 보고했습니다. 이는 Reverse Process에서 분산의 정확한 값보다 평균의 정확한 예측이 더 중요하다는 것을 시사합니다.

 

그림 5.2.  Posterior 분포 $q(x_{t-1}❘x_t, x_0)$의 시각화. $x_t$와 $x_0$의 가중 평균을 중심으로 한 가우시안입니다.

5.3 평균의 파라미터화: 세 가지 관점

신경망 $p_\theta(x_{t-1}|x_t) = \mathcal{N}(x_{t-1}; \mu_\theta(x_t, t), \sigma_t^2 \mathbf{I})$의 핵심은 평균 $\mu_\theta$를 어떻게 파라미터화하느냐입니다. Ho et al. (2020)은 posterior 평균 $\tilde{\mu}_t$의 구조를 활용한 영리한 파라미터화를 제안했습니다.

5.3.1 $x_0$ 예측을 통한 파라미터화

Posterior 평균 $\tilde{\mu}_t(x_t, x_0)$에서 $x_0$만 알면 $\tilde{\mu}_t$를 완전히 결정할 수 있습니다. 따라서 신경망이 $x_0$를 예측하도록 학습시킬 수 있습니다.

$x_\theta(x_t, t)$를 $x_0$의 추정치라 하면:

$$ \mu_\theta(x_t, t) = \frac{\sqrt{\alpha_t}(1-\bar{\alpha}_{t-1})}{1-\bar{\alpha}_t} x_t + \frac{\sqrt{\bar{\alpha}_{t-1}} \beta_t}{1-\bar{\alpha}_t} x_\theta(x_t, t) $$

이 파라미터화에서 신경망은 노이즈가 섞인 $x_t$로부터 원본 이미지 $x_0$를 직접 예측합니다.

5.3.2 노이즈 예측을 통한 파라미터화

4장의 Forward Process closed-form을 상기합시다:

$$ x_t = \sqrt{\bar{\alpha}_t} \, x_0 + \sqrt{1-\bar{\alpha}_t} \, \epsilon, \quad \epsilon \sim \mathcal{N}(0, \mathbf{I}) $$

이 식을 $x_0$에 대해 풀면:

$$ x_0 = \frac{1}{\sqrt{\bar{\alpha}_t}} \left( x_t - \sqrt{1-\bar{\alpha}_t} \, \epsilon \right) $$

이를 posterior 평균에 대입하고 정리하면:

$$ \tilde{\mu}_t = \frac{1}{\sqrt{\alpha_t}} \left( x_t - \frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}} \epsilon \right) $$

이 결과는 매우 중요합니다. Posterior 평균은 현재 상태 $x_t$에서 노이즈 $\epsilon$의 스케일된 버전을 빼는 것으로 표현됩니다. 따라서 신경망이 노이즈 $\epsilon$을 예측하도록 학습시킬 수 있습니다.

$\epsilon_\theta(x_t, t)$를 노이즈의 추정치라 하면:

핵심 파라미터화 (노이즈 예측): $$ \mu_\theta(x_t, t) = \frac{1}{\sqrt{\alpha_t}} \left( x_t - \frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}} \epsilon_\theta(x_t, t) \right) $$

이것이 Ho et al. (2020)이 채택한 파라미터화입니다. 신경망 $\epsilon_\theta$는 $x_t$에 추가된 노이즈를 예측하고, 이 예측을 사용하여 $\mu_\theta$를 구성합니다.

5.3.3 Score Function을 통한 파라미터화

세 번째 관점은 score function $\nabla_{x} \log p(x)$와의 연결입니다. 이 관점은 DDPM을 score-based 생성 모델과 연결하는 다리 역할을 하며, 6장에서 더 깊이 다룹니다. 여기서는 연결의 핵심만 언급합니다.

Score function은 확률 밀도 함수의 로그 그래디언트입니다:

$$ s(x) = \nabla_x \log p(x) $$

$q(x_t|x_0) = \mathcal{N}(x_t; \sqrt{\bar{\alpha}_t} x_0, (1-\bar{\alpha}_t)\mathbf{I})$의 score를 계산하면:

$$ \nabla_{x_t} \log q(x_t|x_0) = -\frac{x_t - \sqrt{\bar{\alpha}_t} x_0}{1-\bar{\alpha}_t} = -\frac{\epsilon}{\sqrt{1-\bar{\alpha}_t}} $$

마지막 등호는 $x_t = \sqrt{\bar{\alpha}_t} x_0 + \sqrt{1-\bar{\alpha}_t} \epsilon$를 사용했습니다. 따라서:

$$ \epsilon = -\sqrt{1-\bar{\alpha}_t} \cdot \nabla_{x_t} \log q(x_t|x_0) $$

이 관계는 노이즈 예측과 score 추정이 본질적으로 같은 문제임을 보여줍니다. 노이즈를 예측하는 것은 score function의 스케일된 버전을 예측하는 것과 동치입니다.

세 가지 파라미터화의 등가성: $$ \text{노이즈 예측 } \epsilon_\theta \;\Longleftrightarrow\; \text{원본 예측 } x_\theta \;\Longleftrightarrow\; \text{Score 예측 } s_\theta $$ 이 세 관점은 수학적으로 동치이며, 하나를 알면 나머지를 즉시 계산할 수 있습니다.

5.4 Variational Lower Bound의 유도

이제 핵심 질문에 답할 차례입니다. 신경망 $\epsilon_\theta$를 어떻게 학습시켜야 할까요? 답은 Variational Lower Bound (VLB), 또는 동치인 Evidence Lower Bound (ELBO)를 최대화하는 것입니다. 3장에서 배운 ELBO의 개념을 DDPM에 적용합니다.

5.4.1 로그 가능도의 하한

생성 모델의 목표는 데이터의 로그 가능도 $\log p_\theta(x_0)$를 최대화하는 것입니다. 그러나 이를 직접 계산하기 위해서는 모든 잠재 변수 $x_{1:T}$에 대해 적분해야 합니다:

$$ p_\theta(x_0) = \int p_\theta(x_{0:T}) \, dx_{1:T} $$

이 적분은 계산 불가능합니다. 대신, Jensen's inequality를 사용하여 하한을 유도합니다:

$$ \log p_\theta(x_0) = \log \int p_\theta(x_{0:T}) \, dx_{1:T} = \log \int \frac{p_\theta(x_{0:T})}{q(x_{1:T}|x_0)} q(x_{1:T}|x_0) \, dx_{1:T} $$ $$ \geq \int q(x_{1:T}|x_0) \log \frac{p_\theta(x_{0:T})}{q(x_{1:T}|x_0)} \, dx_{1:T} = \mathbb{E}_{q(x_{1:T}|x_0)} \left[ \log \frac{p_\theta(x_{0:T})}{q(x_{1:T}|x_0)} \right] $$

이 하한을 ELBO라 부르며, 음의 ELBO가 VLB입니다. 이제 이를 더 유용한 형태로 분해합니다.

5.4.2 VLB의 분해

로그 비율을 전개합니다:

$$ \log \frac{p_\theta(x_{0:T})}{q(x_{1:T}|x_0)} = \log p(x_T) + \sum_{t=1}^{T} \log \frac{p_\theta(x_{t-1}|x_t)}{q(x_t|x_{t-1})} $$

Forward Process의 마르코프 성질을 활용하여 $q(x_{1:T}|x_0) = \prod_{t=1}^{T} q(x_t|x_{t-1})$를 대입하고, $t \geq 2$에 대해 $q(x_t|x_{t-1})$를 베이즈 정리로 변환합니다:

$$ q(x_t|x_{t-1}) = \frac{q(x_{t-1}|x_t, x_0) q(x_t|x_0)}{q(x_{t-1}|x_0)} $$

이를 대입하고 텔레스코핑(telescoping)을 적용하면, VLB는 다음과 같이 분해됩니다:

정리 5.2 (VLB의 분해). Variational Lower Bound는 다음과 같이 세 종류의 항으로 분해됩니다: $$ L_{\text{VLB}} = L_T + \sum_{t=2}^{T} L_{t-1} + L_0 $$ 여기서 각 항은: $$ L_T = D_{\text{KL}}(q(x_T|x_0) \| p(x_T)) $$ $$ L_{t-1} = D_{\text{KL}}(q(x_{t-1}|x_t, x_0) \| p_\theta(x_{t-1}|x_t)), \quad t = 2, \ldots, T $$ $$ L_0 = -\log p_\theta(x_0|x_1) $$

5.4.3 각 항의 의미

$L_T$ (Prior Matching Term): 이 항은 Forward Process의 마지막 분포 $q(x_T|x_0)$와 prior $p(x_T) = \mathcal{N}(0, \mathbf{I})$ 사이의 KL divergence입니다. 4장에서 보았듯이, 적절한 노이즈 스케줄과 충분히 큰 $T$에 대해 $q(x_T|x_0) \approx \mathcal{N}(0, \mathbf{I})$이므로, 이 항은 거의 0에 가깝습니다. 또한 이 항은 학습 가능한 파라미터 $\theta$를 포함하지 않으므로 최적화에서 무시합니다.

$L_{t-1}$ (Denoising Matching Terms): 핵심적인 항들입니다. 각 $L_{t-1}$은 이론적 posterior $q(x_{t-1}|x_t, x_0)$와 학습된 전이 $p_\theta(x_{t-1}|x_t)$ 사이의 KL divergence입니다. 신경망이 각 스텝에서 정확한 denoising을 수행하도록 학습됩니다.

$L_0$ (Reconstruction Term): 마지막 스텝에서 $x_0$를 얼마나 잘 복원하는지를 측정합니다. 연속 데이터의 경우 가우시안으로 모델링하고, 이산 데이터(예: 픽셀 값)의 경우 discretized 가우시안을 사용합니다.

5.5 KL Divergence의 해석적 계산

VLB의 핵심 항인 $L_{t-1}$을 계산합니다. $q(x_{t-1}|x_t, x_0)$와 $p_\theta(x_{t-1}|x_t)$가 모두 가우시안이므로, 두 가우시안 사이의 KL divergence 공식을 적용할 수 있습니다.

5.5.1 가우시안 KL Divergence

두 가우시안 $\mathcal{N}(\mu_1, \sigma^2 \mathbf{I})$와 $\mathcal{N}(\mu_2, \sigma^2 \mathbf{I})$ 사이의 KL divergence는 (같은 분산을 가정):

$$ D_{\text{KL}}(\mathcal{N}(\mu_1, \sigma^2 \mathbf{I}) \| \mathcal{N}(\mu_2, \sigma^2 \mathbf{I})) = \frac{1}{2\sigma^2} \|\mu_1 - \mu_2\|^2 $$

이를 $L_{t-1}$에 적용하면:

$$ L_{t-1} = D_{\text{KL}}(q(x_{t-1}|x_t, x_0) \| p_\theta(x_{t-1}|x_t)) = \frac{1}{2\sigma_t^2} \|\tilde{\mu}_t(x_t, x_0) - \mu_\theta(x_t, t)\|^2 $$

여기서 $\sigma_t^2$는 두 분포에서 공유하는 분산입니다 (Ho et al.은 $\sigma_t^2 = \tilde{\beta}_t$ 또는 $\sigma_t^2 = \beta_t$를 사용).

5.5.2 노이즈 예측 파라미터화 대입

5.3절에서 유도한 노이즈 예측 파라미터화를 대입합니다:

$$ \tilde{\mu}_t = \frac{1}{\sqrt{\alpha_t}} \left( x_t - \frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}} \epsilon \right) $$ $$ \mu_\theta = \frac{1}{\sqrt{\alpha_t}} \left( x_t - \frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}} \epsilon_\theta(x_t, t) \right) $$

차이를 계산하면:

$$ \tilde{\mu}_t - \mu_\theta = \frac{\beta_t}{\sqrt{\alpha_t}\sqrt{1-\bar{\alpha}_t}} (\epsilon_\theta(x_t, t) - \epsilon) $$

따라서:

$$ L_{t-1} = \frac{1}{2\sigma_t^2} \cdot \frac{\beta_t^2}{\alpha_t(1-\bar{\alpha}_t)} \|\epsilon - \epsilon_\theta(x_t, t)\|^2 $$

가중치를 $\lambda_t = \frac{\beta_t^2}{2\sigma_t^2 \alpha_t(1-\bar{\alpha}_t)}$로 정의하면:

$$ L_{t-1} = \lambda_t \|\epsilon - \epsilon_\theta(x_t, t)\|^2 $$

5.5.3 단순화된 손실 함수

Ho et al. (2020)의 핵심 발견은 가중치 $\lambda_t$를 무시하고 단순히 1로 설정해도 좋은 성능을 얻을 수 있다는 것입니다. 즉:

단순화된 손실 함수: $$ L_{\text{simple}} = \mathbb{E}_{t, x_0, \epsilon} \left[ \|\epsilon - \epsilon_\theta(x_t, t)\|^2 \right] $$ 여기서 $t \sim \text{Uniform}(\{1, \ldots, T\})$, $x_0 \sim q(x_0)$, $\epsilon \sim \mathcal{N}(0, \mathbf{I})$이고, $x_t = \sqrt{\bar{\alpha}_t} x_0 + \sqrt{1-\bar{\alpha}_t} \epsilon$입니다.

이 단순화가 효과적인 이유에 대해 Ho et al.은 다음과 같은 해석을 제시했습니다. 가중치 $\lambda_t$는 작은 $t$에서 크고 큰 $t$에서 작습니다. 가중치를 무시하면, 모든 시간 스텝에서의 노이즈 예측이 동등하게 중요해지며, 이는 샘플 품질 향상에 도움이 됩니다.

5.6 Tweedie's Formula와의 연결

노이즈 예측이 최적의 전략인 이유를 더 깊이 이해하기 위해, 통계학의 고전적 결과인 Tweedie's Formula와의 연결을 살펴봅니다.

5.6.1 Tweedie's Formula

Tweedie's Formula (Robbins, 1956; Efron, 2011)는 다음과 같이 진술됩니다. $X$가 prior 분포를 따르고, 노이즈가 추가된 관측 $Y = X + \sigma Z$ ($Z \sim \mathcal{N}(0, \mathbf{I})$)가 주어졌을 때, $X$의 posterior 평균은:

$$ \mathbb{E}[X|Y] = Y + \sigma^2 \nabla_Y \log p_Y(Y) $$

여기서 $p_Y$는 $Y$의 주변 분포입니다. 이 공식은 "최적의 denoiser"가 score function과 직접 연결됨을 보여줍니다.

5.6.2 DDPM에의 적용

DDPM의 맥락에서 $Y = x_t$, $X = \sqrt{\bar{\alpha}_t} x_0$로 대응시키면, $x_t = \sqrt{\bar{\alpha}_t} x_0 + \sqrt{1-\bar{\alpha}_t} \epsilon$이므로 노이즈의 표준편차는 $\sqrt{1-\bar{\alpha}_t}$입니다.

Tweedie's Formula를 적용하면:

$$ \mathbb{E}[\sqrt{\bar{\alpha}_t} x_0 | x_t] = x_t + (1-\bar{\alpha}_t) \nabla_{x_t} \log q(x_t) $$

$x_0$에 대해 풀면:

$$ \mathbb{E}[x_0 | x_t] = \frac{1}{\sqrt{\bar{\alpha}_t}} \left( x_t + (1-\bar{\alpha}_t) \nabla_{x_t} \log q(x_t) \right) $$

이를 앞서 유도한 노이즈-score 관계 $\epsilon = -\sqrt{1-\bar{\alpha}_t} \nabla_{x_t} \log q(x_t|x_0)$와 비교하면:

$$ \mathbb{E}[x_0 | x_t] = \frac{1}{\sqrt{\bar{\alpha}_t}} \left( x_t - \sqrt{1-\bar{\alpha}_t} \cdot \mathbb{E}[\epsilon|x_t] \right) $$

이 결과는 최적의 $x_0$ 예측이 최적의 노이즈 예측과 등가임을 보여줍니다. 신경망이 $\epsilon$을 정확히 예측하면, 그것으로부터 $x_0$의 최적 추정치를 즉시 계산할 수 있습니다.

그림 5.3.  Tweedie's Formula의 기하학적 해석. 노이즈가 섞인 관측에서 score 방향으로 이동하여 원본을 추정합니다.

 

5.7 학습 및 샘플링 알고리즘

이제 이론적 내용을 구체적인 알고리즘으로 정리합니다.

5.7.1 학습 알고리즘

Algorithm 1: DDPM Training

repeat
    $x_0 \sim q(x_0)$                         // 데이터셋에서 샘플링
    $t \sim \text{Uniform}(\{1, \ldots, T\})$     // 시간 스텝 샘플링
    $\epsilon \sim \mathcal{N}(0, \mathbf{I})$                    // 노이즈 샘플링
    $x_t = \sqrt{\bar{\alpha}_t} x_0 + \sqrt{1-\bar{\alpha}_t} \epsilon$  // Forward Process (closed-form)
    gradient step on $\nabla_\theta \|\epsilon - \epsilon_\theta(x_t, t)\|^2$
until converged

알고리즘의 핵심은 4장에서 유도한 closed-form 덕분에 임의의 시간 $t$로 한 번에 도달할 수 있다는 것입니다. 이것이 학습을 효율적으로 만듭니다.

5.7.2 샘플링 알고리즘

Algorithm 2: DDPM Sampling

$x_T \sim \mathcal{N}(0, \mathbf{I})$
for $t = T, T-1, \ldots, 1$ do
    $z \sim \mathcal{N}(0, \mathbf{I})$ if $t > 1$, else $z = 0$
    $x_{t-1} = \frac{1}{\sqrt{\alpha_t}} \left( x_t - \frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}} \epsilon_\theta(x_t, t) \right) + \sigma_t z$
end for
return $x_0$

샘플링은 $T$ 스텝의 순차적 denoising입니다. 각 스텝에서 신경망 $\epsilon_\theta$가 노이즈를 예측하고, 이를 제거하여 한 스텝 덜 노이즈한 상태로 이동합니다. 마지막 스텝($t=1$)에서는 노이즈를 추가하지 않습니다.

5.7.3 분산 $\sigma_t^2$의 선택

Ho et al. (2020)은 두 가지 선택을 실험했습니다:

1. $\sigma_t^2 = \beta_t$: Forward Process의 분산을 그대로 사용
2. $\sigma_t^2 = \tilde{\beta}_t = \frac{(1-\bar{\alpha}_{t-1})\beta_t}{1-\bar{\alpha}_t}$: Posterior 분산 사용

두 선택 모두 비슷한 샘플 품질을 보였습니다. 직관적으로, $\sigma_t^2 = \beta_t$는 $x_0$가 주어지지 않았을 때의 최적 분산에 가깝고, $\sigma_t^2 = \tilde{\beta}_t$는 $x_0$가 주어졌을 때의 최적 분산입니다. 실제로 $x_0$는 주어지지 않으므로, 두 값의 중간 어딘가가 이상적일 수 있습니다.

Nichol과 Dhariwal (2021)은 $\sigma_t^2$를 학습하는 방법을 제안하여 log-likelihood를 개선했습니다. 이는 VLB의 전체 항을 최적화하는 것과 관련됩니다.

5.8 Denoising Score Matching과의 연결

DDPM의 학습 목표가 어떻게 고전적인 Denoising Score Matching과 연결되는지 살펴봅니다. 이 연결은 Ho et al. (2020) 논문의 중요한 기여 중 하나입니다.

5.8.1 Denoising Score Matching

Denoising Score Matching (Vincent, 2011)은 score function $\nabla_x \log p(x)$를 학습하는 방법입니다. 핵심 아이디어는 다음과 같습니다. 데이터 $x$에 노이즈를 추가하여 $\tilde{x}$를 얻고, 신경망 $s_\theta(\tilde{x})$가 노이즈가 섞인 분포의 score를 예측하도록 학습합니다:

$$ \mathcal{L}_{\text{DSM}} = \mathbb{E}_{x \sim p(x), \tilde{x} \sim q(\tilde{x}|x)} \left[ \|s_\theta(\tilde{x}) - \nabla_{\tilde{x}} \log q(\tilde{x}|x)\|^2 \right] $$

가우시안 노이즈 $\tilde{x} = x + \sigma \epsilon$의 경우:

$$ \nabla_{\tilde{x}} \log q(\tilde{x}|x) = -\frac{\tilde{x} - x}{\sigma^2} = -\frac{\epsilon}{\sigma} $$

따라서 Denoising Score Matching은 본질적으로 노이즈를 예측하는 것과 같습니다.

5.8.2 DDPM과의 등가성

DDPM의 단순화된 손실 함수를 다시 씁시다:

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

이는 각 노이즈 레벨 $t$에서의 Denoising Score Matching 손실을 시간에 대해 평균낸 것입니다. $\epsilon_\theta(x_t, t)$와 $s_\theta(x_t, t)$ 사이의 관계:

$$ s_\theta(x_t, t) = -\frac{\epsilon_\theta(x_t, t)}{\sqrt{1-\bar{\alpha}_t}} $$

를 사용하면, DDPM 손실은 다중 스케일 Denoising Score Matching과 등가입니다:

$$ L_{\text{simple}} = \mathbb{E}_{t, x_0, \epsilon} \left[ (1-\bar{\alpha}_t) \|s_\theta(x_t, t) - \nabla_{x_t} \log q(x_t|x_0)\|^2 \right] $$

이 연결은 DDPM이 다양한 노이즈 레벨에서 score function을 동시에 학습하는 모델로 해석될 수 있음을 보여줍니다.

5.9 요약

이번 장에서 우리는 DDPM의 Reverse Process와 학습 목표를 수학적으로 완전히 이해했습니다.

Reverse Process는 학습된 가우시안 전이 $p_\theta(x_{t-1}|x_t)$로 정의되며, 이론적 posterior $q(x_{t-1}|x_t, x_0)$를 목표로 학습됩니다. Posterior $q(x_{t-1}|x_t, x_0)$는 베이즈 정리를 통해 해석적으로 유도되며, $x_t$와 $x_0$의 가중 평균을 중심으로 한 가우시안입니다. 평균의 파라미터화는 세 가지 등가적 관점—노이즈 예측, $x_0$ 예측, score 예측—으로 표현될 수 있습니다. VLB는 KL divergence 항들의 합으로 분해되며, 각 항은 해석적으로 계산 가능합니다. 단순화된 손실 함수 $L_{\text{simple}} = \|\epsilon - \epsilon_\theta(x_t, t)\|^2$가 실제로 잘 작동하며, 이는 Denoising Score Matching과 수학적으로 연결됩니다. Tweedie's Formula는 최적의 denoiser와 score function의 관계를 명확히 합니다.

참고문헌

  1. Ho, J., Jain, A., & Abbeel, P. (2020). Denoising Diffusion Probabilistic Models. Advances in Neural Information Processing Systems, 33, 6840-6851. arXiv:2006.11239
  2. Sohl-Dickstein, J., Weiss, E., Maheswaranathan, N., & Ganguli, S. (2015). Deep Unsupervised Learning using Nonequilibrium Thermodynamics. Proceedings of the 32nd International Conference on Machine Learning (ICML), 2256-2265. arXiv:1503.03585
  3. Nichol, A., & Dhariwal, P. (2021). Improved Denoising Diffusion Probabilistic Models. Proceedings of the 38th International Conference on Machine Learning (ICML), 8162-8171. arXiv:2102.09672
  4. Anderson, B. D. O. (1982). Reverse-time diffusion equation models. Stochastic Processes and their Applications, 12(3), 313-326.
  5. Feller, W. (1949). On the theory of stochastic processes, with particular reference to applications. Proceedings of the Berkeley Symposium on Mathematical Statistics and Probability, 403-432.
  6. Hyvärinen, A. (2005). Estimation of Non-Normalized Statistical Models by Score Matching. Journal of Machine Learning Research, 6, 695-709.
  7. Vincent, P. (2011). A Connection Between Score Matching and Denoising Autoencoders. Neural Computation, 23(7), 1661-1674.
  8. Efron, B. (2011). Tweedie's Formula and Selection Bias. Journal of the American Statistical Association, 106(496), 1602-1614.
  9. Robbins, H. (1956). An Empirical Bayes Approach to Statistics. Proceedings of the Third Berkeley Symposium on Mathematical Statistics and Probability, 1, 157-163.
  10. Kingma, D. P., & Welling, M. (2014). Auto-Encoding Variational Bayes. Proceedings of the 2nd International Conference on Learning Representations (ICLR). arXiv:1312.6114