이 포스트는 DDPM을 설명한 글입니다. VAE를 먼저 읽으시면 이해가 더 잘 됩니다. :)
복습
VAE를 다음과 같이 표현할 수 있다.
Latent Vector인 \(x_T\)에서 synthetic image인 \(x_0\)를 sampling하는 과정이며, 여기서 \(p,q\)를 training을 통해서 구하는 것이다.
VAE to DDPM
Size
그럼 일단 첫 번째로 Latent vector \(x_T\in \mathbb{R}^{H\times W\times 3}\)가 synthetic image와 동일한 크기라고 가정을 하자.
Hierarchical VAE
그 다음으로 VAE가 여러 스탭으로 진행된다고 생각해보자.
ELBO
아직 Diffusion의 D자도 들어가지는 않았지만, hierarchical VAE로 만들기만 하면 논문에 나온 공식을 계산 할 수 있다. 공식 이해는 implementation에서 그리 중요하지는 않지만 이거 아니면 논문을 읽을 필요가 없으므로 꼭 이해해 두자.
$$ \begin{aligned} \log{p(x_0)} &\geq -KL(q(x_T|x_0) || p(x_T)) +\mathbb{E}_{x_{1} \sim q(x_{1}|x_0)} \bigg[ {\log p(x_0 | x_{1})} \bigg] - \mathbb{E}_{x_{1:T} \sim q(x_{1:T}|x_0)} \sum^{T-1}_{i=1} \bigg[ KL(q(x_i|x_{i+1}, x_0) || p(x_i | x_{i+1})) \bigg]\\ Loss = \text{위의 negative} &= KL(q(x_T|x_0) || p(x_T)) -\mathbb{E}_{x_{1} \sim q(x_{1}|x_0)} \bigg[ {\log p(x_0 | x_{1})} \bigg] + \mathbb{E}_{x_{1:T} \sim q(x_{1:T}|x_0)} \sum^{T-1}_{i=1} \bigg[ KL(q(x_i|x_{i+1}, x_0) || p(x_i | x_{i+1})) \bigg] \end{aligned} $$
우리가 증명할 목표는 다음과 같은데, 결국 위 부등식을 증명하게 된다. 여기서 첫번째 항은 \(p(x_T\)와 \(q(x_T|x_0)\)이 비슷함을 원하는 것이기 때문에 Regularization term이라고 생각하면 좋고, 두번째는 \(p(x_0|x_1)\)이 커지는 것을 원하기 때문에, Reconstruction Error로 보면 좋다. 마지막은 Diffusion Loss로 각 스텝마다 \(p\)와 \(q\)가 서로의 반대가 되는 것을 유도하는 항이라고 보면 된다.
위 부등식의 step-by-step 증명은 이 페이지를 참조하면 좋다. [2]
Diffusion
VAE의 경우 \(p\)를 찾는 것을 목표로 하기 위해 \(q\)를 같이 training 하는 방법이라면, diffusion은 고정된 \(q\)를 가지고 \(q\)의 역함수인 \(p\)를 찾는 것이다. (정확히 말하면 역함수라는 표현은 잘못 된 것임.)
Diffusion은 이 논문[1]에서 처음 소개 된 것 같은데, ‘확산’이라는 의미에 맞게 가면 갈 수록 random noise를 섞어주는 역할을 한다. 즉 이 VAE에서는 \(q\)가 random noise를 섞고, 트레이닝을 통해 \(p\)가 그것의 반대를 approximate하게 하는 것을 찾는 것이 목표이다.
Forward Process
Forward Process는 다음과 같다.
$$ q(x_t | x_{t-1}) \coloneqq \mathcal{N}(x_t ; \sqrt{1-\beta_t}x_{t-1}, \beta_t I) $$
즉 매 스탭마다 값을 점점 줄이고, \(\beta_t\)만큼 variance를 주어서 결국 noise를 섞어 주는 것이다. 이 논문에서는 \(\beta_t\)는 고정된 상수를 이용했다고 한다. 따라서 이 논문에서는 forward process에는 learnable parameter이 아예 존재하지 않는다고 보면 된다. Diffusion의 장점 중 하나는 hierarchical VAE에서 중간 값을 closed-form으로 계산이 가능하다는 점이 있다.
즉, \(x_0\)을 주었을 때, \(x_t\)를 계산하면,
$$ \begin{aligned} q(x_t | x_0) &= \mathcal{N}(x_t; \sqrt{\bar{\alpha_t}}x_0, (1-\bar{\alpha_t})I) \\ \alpha_t &\coloneqq 1-\beta_t \\ \bar{\alpha_t} &\coloneqq \prod^t_{s=1}\alpha_s \end{aligned} $$ 로 계산 된다. 증명은 다음을 참고하자 [3].
여기서 봐야될 것은 만약 \(t\rightarrow\infty \)으로 가게 된다면 \(\bar{\alpha_t} \rightarrow 0 \)이 되면서, \(q(x_t | x_0) = \mathcal{N}(x_t; \sqrt{\bar{\alpha_t}}x_0, (1-\bar{\alpha_t})I) \rightarrow \mathcal{N}(x_t; 0, I) \)인 가우시안 분포가 되어 즉 \(T\)가 크기만 하다면 거의 완벽한 랜덤인 분포로 바뀌게 되는 것이다. 그리고, 이후 스텝으로 갈수록 점 점 variance를 크게 주는 데,
$$ \begin{aligned} \beta_1 &< \beta_2 < \beta_3 < \cdots < \beta_T \\ \bar{\alpha_1} &> \bar{\alpha_2} > \bar{\alpha_3} > \cdots > \bar{\alpha_T} \\ \end{aligned} $$ 이렇게 두는 것이 좋다고 하는 데 그 이유는 잘 모르겠다.
Reverse Process
Forward Process가 원래 이미지에서 점점 노이즈를 섞는 것이라고 하면 Reverse Process는 노이즈를 점점 제거하여 image를 generation 하는 것이라고 보면 된다.
\(\theta\)로 learnable한 \(p_\theta\)인 모델이 있다고 하면, 노이즈를 한 번 제거하는 것을 \(p_\theta(x_{t-1}| x_t)\)를 의미한다. 여기서도 VAE와 동일하게 모델이 distribution의 형태를 가지고 있다고 생각하면 된다.
$$ p_\theta ( x_{t-1}|x_t) \coloneqq \mathcal{N} ( x_{t-1}; \mu_{\theta}(x_t, t) , \Sigma_\theta(x_t, t)) $$
다만 이 논문에서는 \(\Sigma_\theta(x_t, t)\)가 untrained time dependent constant로 가정하여 \(\Sigma_\theta(x_t, t)=\sigma^2_t I\)로 두고, 두 가지 가능성을 알려주는 데, (1) \(\sigma^2_t = \beta_t \), (2) \(\frac{1-\bar{\alpha}_{t-1}}{1-\bar{\alpha}_t} \beta_t\), 첫 번째는 generated image \(x_0 \sim \mathcal{N}(0,I)\)일 때 optimal하고, 두 번째는 \(x_0\)이 단 하나의 이미지일때 optimal 하다고 한다 [1]. 즉 양 극단의 경우를 두고 generated image가 단 한 장인지, 완전 랜덤인지 두 개의 경우를 이야기 하는 것 같다. 다만 논문에서는 무엇을 사용하든 비슷한 결과를 보여준다고 한다.
마지막 \(p_\theta(x_T) = p(x_T) = \mathcal{N}(0,I)\)라고 두면 된다.
Loss
위에서 증명한 식을 다시 가져오면: $$ Loss = KL(q(x_T|x_0) || p_\theta(x_T)) -\mathbb{E}_{x_{1} \sim q(x_{1}|x_0)} \bigg[ {\log p_\theta(x_0 | x_{1})} \bigg] + \mathbb{E}_{x_{1:T} \sim q(x_{1:T}|x_0)} \sum^{T}_{t=2} \bigg[ KL(q(x_{t-1}|x_{t}, x_0) || p_\theta(x_{t-1} | x_{t})) \bigg] $$
여기서 일단 첫번째 항은 사실 learnable parameter에 영향을 주지 않기 때문에 통째로 날려버리면,
$$ Loss = -\mathbb{E}_{x_{1} \sim q(x_{1}|x_0)} \bigg[ {\log p_\theta(x_0 | x_{1})} \bigg] + \mathbb{E}_{x_{1:T} \sim q(x_{1:T}|x_0)} \sum^{T}_{t=2} \bigg[ KL(q(x_{t-1}|x_{t}, x_0) || p_\theta(x_{t-1} | x_{t})) \bigg] $$
로 다시 적을 수 있다.
여기서 \(q(x_{t-1}|x_{t}, x_0)\)를 계산할 수 있어야 하는데, diffusion의 경우 이를 closed-form으로 깔끔하게 계산이 가능하다.
$$ \begin{aligned} & q(x_{t-1}|x_{t}, x_0) = \mathcal{N}(x_{t-1}; \tilde{\mu}_t(x_t, x_0), \tilde{\beta}_{t}I) \\ & \text{where }\quad \tilde{\mu}_t(x_t, x_0) \coloneqq \frac{\sqrt{\bar{\alpha}_{t-1}}\beta_t}{1-\bar{\alpha}_t}x_0 + \frac{\sqrt{\alpha_t} (1-\bar{\alpha}_{t-1})}{1-\bar{\alpha}_t}x_t \\ & \text{where }\quad \tilde{\beta}_t \coloneqq \frac{1-\bar{\alpha}_{t-1}}{1-\bar{\alpha}_t}\beta_t \end{aligned} $$
이에 대한 증명은 다음 참조 [4].
따라서 \(KL(q(x_{t-1}|x_{t}, x_0) || p_\theta(x_{t-1} | x_{t}))\)의 두 항이 전무 Gaussian 형태가 되기에, KL-발산의 계산이 간단해진다. 이것에 대한 계산은 귀찮으니까 넘어가고, 대충 MSE의 형태를 띄운다.
$$ KL(q(x_{t-1}|x_{t}, x_0) || p_\theta(x_{t-1} | x_{t})) = \frac{1}{2\sigma^2_t} || \tilde{\mu}_t(x_t, x_0) - \mu_\theta (x_t, t) ||^2 + C $$
여기서 \(C\)는 learnable-parameter과 관련없는 상수. 즉 모델이 \(x_t\) 를 받았을 때 추측한 \(x_{t-1}\)의 평균 값과 \( q(x_{t-1}|x_{t}, x_0) \)의 평균 값의 MSE가 Loss가 되는 것이다.
다만 논문에서는 neural network가 \(\mu\) (이미지)를 바로 추측하는 것 보다는 반대로 noise를 output하는 것이 더 stable 하다고 주장한다. 즉 clean image를 output하는 것보다는 noise를 output해서 input image에서 제거하는 것을 말한다. 사실 둘 다 비슷한 내용이고 후속 논문 중에서는 둘 다 크게 상관이 없다고 하는 논문들도 있지만 여기서는 noise-prediction 방법을 이용에 주로 설명한다. 이 둘은 수학적으로 동일함으로 아마도 neural network의 stability 때문에 차이가 갈리는 것 같다. 이에 대한 계산은 아래를 참조.
Reparameterization using \(\epsilon\)
위에서 \(\tilde{\mu}_t(x_t, x_0) \coloneqq \frac{\sqrt{\bar{\alpha}_{t-1}}\beta_t}{1-\bar{\alpha}_t}x_0 + \frac{\sqrt{\alpha_t} (1-\bar{\alpha}_{t-1})}{1-\bar{\alpha}_t}x_t\) 라고 했는 데, 이를 \(x_t\) 와 \(\epsilon\sim \mathcal{N}(0,1)\)으로만 표현해 주자.
앞에서 \(q(x_t|x_0) \) 계산을 통해 \(x_t = \sqrt{\bar\alpha_t}x_0 + \sqrt{1-\bar\alpha_t} \epsilon\)이라고 했기에, \( x_0 = \frac{1}{\sqrt{\bar\alpha_t}} x_t - \frac{\sqrt{1-\bar\alpha_t}}{\sqrt{\bar\alpha_t}}\epsilon \)을 넣어 치환 해주면,
$$ \begin{aligned} \tilde{\mu}_t(x_t, x_0) &= \frac{\sqrt{\bar{\alpha}_{t-1}}\beta_t}{1-\bar{\alpha}_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}(\frac{1}{\sqrt{\bar\alpha_t}} x_t - \frac{\sqrt{1-\bar\alpha_t}}{\sqrt{\bar\alpha_t}}\epsilon) + \frac{\sqrt{\alpha_t} (1-\bar{\alpha}_{t-1})}{1-\bar{\alpha}_t}x_t \\ &= \frac{1}{\sqrt{\alpha_t}} \Big( {x}_t - \frac{1 - \alpha_t}{\sqrt{1 - \bar{\alpha}_t}} {\epsilon} \Big) \end{aligned} $$
로 정리가 된다. 이제 우리가 원하는 neural network \(\epsilon_\theta \)는 위의 \(\epsilon\)을 predict하는 것이라고 두면, 본래 \(\mu_\theta\)는 같은 형태를 띄어야 되므로,
$$ \begin{aligned} \mu_\theta (x_t, t) &= \frac{1}{\sqrt{\alpha_t}} \Big( {x}_t - \frac{1 - \alpha_t}{\sqrt{1 - \bar{\alpha}_t}} \epsilon_\theta(x_t, t) \Big) \end{aligned} $$
로 잡을 수 있다. 결국,
$$ \begin{aligned} KL(q(x_{t-1}|x_{t}, x_0) || p_\theta(x_{t-1} | x_{t})) &= \frac{1}{2\sigma^2_t} || {\color{red}\tilde{\mu}_t(x_t, x_0)} - {\color{blue}\mu_\theta (x_t, t)} ||^2 + C \\ &= \frac{1}{2\sigma^2_t} || {\color{red}\frac{1}{\sqrt{\alpha_t}} \Big( {x}_t - \frac{1 - \alpha_t}{\sqrt{1 - \bar{\alpha}_t}} {\epsilon} \Big)} - {\color{blue}\frac{1}{\sqrt{\alpha_t}} \Big( {x}_t - \frac{1 - \alpha_t}{\sqrt{1 - \bar{\alpha}_t}} \epsilon_\theta(x_t, t) \Big)} ||^2 + C \\ &= \frac{1}{2\sigma^2_t} || {\frac{1}{\sqrt{\alpha_t}}\frac{1 - \alpha_t}{\sqrt{1 - \bar{\alpha}_t}} \Big( \epsilon_\theta(x_t, t) - {\epsilon} \Big)} ||^2 + C \\ &= \frac{(1 - \alpha_t)^2}{2\sigma^2_t\alpha_t(1 - \bar{\alpha}_t)} || \epsilon_\theta(x_t, t) - {\epsilon} ||^2 + C \\ \end{aligned} $$ 이고, 여기서 만약 조금 더 전개해서 \(x_t \) 역시 \(\epsilon\)을 사용했으므로,
$$ \begin{aligned} KL(q(x_{t-1}|x_{t}, x_0) || p_\theta(x_{t-1} | x_{t})) &= \frac{(1 - \alpha_t)^2}{2\sigma^2_t\alpha_t(1 - \bar{\alpha}_t)} || \epsilon_\theta(x_t, t) - {\epsilon} ||^2 + C \\ &= \frac{(1 - \alpha_t)^2}{2\sigma^2_t\alpha_t(1 - \bar{\alpha}_t)} || \epsilon_\theta(\sqrt{\bar\alpha_t}x_0 + \sqrt{1-\bar\alpha_t} \epsilon, t) - {\epsilon} ||^2 + C \\ \end{aligned} $$
로 쓸 수 있다.
Loss Simplification
여기서 조금 골 때리는 점은 이렇게 정직하게 Loss 계산을 끝냈는 데,
$$ Loss = -\mathbb{E}_{x_{1} \sim q(x_{1}|x_0)} \bigg[ {\log p_\theta(x_0 | x_{1})} \bigg] + \mathbb{E}_{x_{1:T} \sim q(x_{1:T}|x_0)} \sum^{T}_{t=2} \bigg[ KL(q(x_{t-1}|x_{t}, x_0) || p_\theta(x_{t-1} | x_{t})) \bigg] $$
실험적으로 여기서 첫번째 항도 없이 그냥 마지막 항만 loss로 잡아도 된다고 한다. 뿐만 아니라 stochastic하게 \(t\sim[1,T]\)로 sampling해서 계산 한다면 된다는 것이다. 즉 여태까지 증명 열심히 해봐짜 결국 step-by-step noise-removal 모델을 training 하는 것과 그리 다르지 않다는 것….
따라서, 여기까지 계산이 끝났는데,
$$ \begin{aligned} KL(q(x_{t-1}|x_{t}, x_0) || p_\theta(x_{t-1} | x_{t})) &= \frac{(1 - \alpha_t)^2}{2\sigma^2_t\alpha_t(1 - \bar{\alpha}_t)} || \epsilon_\theta(x_t, t) - {\epsilon} ||^2 + C \\ &= {\color{red}\frac{(1 - \alpha_t)^2}{2\sigma^2_t\alpha_t(1 - \bar{\alpha}_t)}} || \epsilon_\theta(\sqrt{\bar\alpha_t}x_0 + \sqrt{1-\bar\alpha_t} \epsilon, t) - {\epsilon} ||^2 + C \\ \end{aligned} $$
상수는 날려버리고 빨간 색 역시 결국 learning-rate에 붙어 있는 거라고 생각하면 loss는 결국
$$ Loss_{simple} = \mathbb{E}_{t\sim [1,T], x_0\sim\text{Training Dataset}, \epsilon \sim \mathcal{N}(0,I)} || \epsilon_\theta(\sqrt{\bar\alpha_t}x_0 + \sqrt{1-\bar\alpha_t} \epsilon, t) - {\epsilon} ||^2 $$
와 동일하다.
Pseudo Code
소론
Loss Simplification에서 “수학적 계산이 의미는 있었나?” 싶을 정도로 Loss의 많은 부분을 통으로 날려버린다. 후속 논문들에서도 이런 수학적 배경을 거의 무시하면서 좋은 결과는 내는 경우가 있다. 최근 ICLR 2023에 제출 된 눈문 중에서도 아예 gaussian noise를 사용하지 않고 [5] [6] diffusion을 하는 괴상한 연구도 보여지는데, 결과적으로 machine learning이 engineering하게 바뀌면서 이런 수학적 타당성이 많이 지워지는 implementation도 많다는 점이 눈에 띈다.
References
[1] Deep Unsupervised Learning using Nonequilibrium Thermodynamics
[2] DDPM ELBO 증명
[4] DDPM tilde-mu 증명
[5] Soft Diffusion
[5] Cold Diffusion