들어가며...
정말 수식 증명만 조지는 내용이라 DDPM 자체에 대한 내용은 이전 글을 참고하면 좋다. 근데 지금보니까 저 글을 쓸때도 완벽하게 이해하고 쓴 건 아닌 것 같다는 생각.. DDPM 논문 링크는 여기 . 이 어려운 논문이 벌써 3년이나 됐나 싶다.
Forward and Reverse process
이미지 x 0 x_0 x 0 에 작은 variance(β \beta β )의 가우시안을 아주 조금씩 계속(T T T 만큼) 더해가다보면 최종 output x T x_T x T 은 x 0 x_0 x 0 와 같은 spatial dimension을 가지는 가우시안 분포가 된다. 논문에서는 각 process에서의 variance를 스케줄링하여 고정값으로 사용 .
Forward process q ( x t ∣ x t − 1 ) q(x_t \vert x_{t-1}) q ( x t ∣ x t − 1 )
q ( x t ∣ x t − 1 ) = N ( x t ; 1 − β t x t − 1 , β t I ) q(x_t \vert x_{t-1}) = \mathcal{N}(x_t; \sqrt{1-\beta_t}x_{t-1}, \beta_t I) q ( x t ∣ x t − 1 ) = N ( x t ; 1 − β t x t − 1 , β t I )
그런데 알다시피 실제로 딥러닝을 학습시키려면 stochastic한 샘플링은 미분이 불가능하기 때문에 reparameterization trick을 사용한다. 각 time step에서 노이즈를 더할 때 이전 step의 output을 1 − β t \sqrt{1-\beta_t} 1 − β t 만큼 scaling 해주는 걸 알 수 있는데, 이렇게 하면 변수의 variance를 1 1 1 로 유지할 수 있게됨.
( 1 − β t ) 2 + ( β t ) 2 = 1 \left(\sqrt{1-\beta_t}\right)^2+\left(\sqrt{\beta_t}\right)^2 = 1 ( 1 − β t ) 2 + ( β t ) 2 = 1
수학적 귀납법을 통해 살펴보게 되면, α t = 1 − β t \alpha_t = 1-\beta_t α t = 1 − β t 에 대해서
For ϵ ∼ N ( 0 , I ) , x t = α t x t − 1 + 1 − α t ϵ \text{For}~\epsilon \sim \mathcal{N}(0, I), x_{t} = \sqrt{\alpha_t}x_{t-1} + \sqrt{1-\alpha_t}\epsilon For ϵ ∼ N ( 0 , I ) , x t = α t x t − 1 + 1 − α t ϵ
일반적인 t t t 번째 step에서의 식이 위처럼 표현되니까
x t = α t α t − 1 x t − 2 + α t 1 − α t − 1 ϵ t − 2 + 1 − α t ϵ t − 1 x_t = \sqrt{\alpha_t}\sqrt{\alpha_{t-1}}x_{t-2} + \sqrt{\alpha_t}\sqrt{1-\alpha_{t-1}}\epsilon_{t-2} + \sqrt{1-\alpha_t}\epsilon_{t-1} x t = α t α t − 1 x t − 2 + α t 1 − α t − 1 ϵ t − 2 + 1 − α t ϵ t − 1
위와 같이 t > 1 t>1 t > 1 인 상황에서 한단계 더 확장해서 표현 가능하고, 뒤의 epsilon들을 잘 합치게 되면,
α t 1 − α t − 1 ϵ t − 2 + 1 − α t ϵ t − 1 → ( ( α t 1 − α t − 1 ) 2 + ( 1 − α t ) 2 ) ϵ \sqrt{\alpha_t}\sqrt{1-\alpha_{t-1}}\epsilon_{t-2} + \sqrt{1-\alpha_t}\epsilon_{t-1} \rightarrow \left(\left(\sqrt{\alpha_t}\sqrt{1-\alpha_{t-1}} \right)^2 + \left(\sqrt{1-\alpha_t} \right)^2\right) \epsilon α t 1 − α t − 1 ϵ t − 2 + 1 − α t ϵ t − 1 → ( ( α t 1 − α t − 1 ) 2 + ( 1 − α t ) 2 ) ϵ
그래서 이게 잘 정리가 된 것이 다음과 같은 식.
q ( x t ∣ x 0 ) = N ( x t ; α ˉ t x 0 , ( 1 − α ˉ t ) I ) q(x_t \vert x_0) = \mathcal{N}(x_t; \sqrt{\bar{\alpha}_t}x_0, (1-\bar{\alpha}_t)I) q ( x t ∣ x 0 ) = N ( x t ; α ˉ t x 0 , ( 1 − α ˉ t ) I )
이걸 또다시 reparameterization 으로 표현하면 loss에서 써먹을 수 있게된다. 논문 전개과정에 주구장창 나오는 식.
x t : = α ˉ t x 0 + 1 − α ˉ t ϵ , ϵ ∼ N ( 0 , I ) x_t := \sqrt{\bar{\alpha}_t} x_0 + \sqrt{1-\bar{\alpha}_t} \epsilon,~\epsilon \sim \mathcal{N}(0, I) x t := α ˉ t x 0 + 1 − α ˉ t ϵ , ϵ ∼ N ( 0 , I )
Reverse process p θ ( x t − 1 ∣ x t ) p_\theta(x_{t-1} \vert x_t) p θ ( x t − 1 ∣ x t )
Forward process의 posterior를 모방하도록 학습될 녀석. 분포를 예측하는 과정임. 정방향에 유사하게 역방향을 학습하고 싶다? 어디서 본 것 같은 워딩인데 자세히 노려보면 VAE에서의 접근과 같기 때문에 결국 얘도 lower bound를 최적화해야함.
p θ ( x 0 : T ) : = p ( x T ) ∏ t = 1 T p θ ( x t − 1 ∣ x t ) p_\theta(x_{0 : T}) := p(x_T) \prod_{t=1}^T p_\theta(x_{t-1} \vert x_t) p θ ( x 0 : T ) := p ( x T ) t = 1 ∏ T p θ ( x t − 1 ∣ x t )
x T x_T x T 부터 샘플링하는 과정을 위의 joint distribution probability로 표현이 가능하고,
p θ ( x t − 1 ∣ x t ) : = N ( x t − 1 ; μ θ ( x t , t ) , Σ θ ( x t , t ) ) p_\theta(x_{t-1} \vert x_t) := \mathcal{N}(x_{t-1}; \mu_{\theta}(x_t,~t), \Sigma_\theta (x_t,~t)) p θ ( x t − 1 ∣ x t ) := N ( x t − 1 ; μ θ ( x t , t ) , Σ θ ( x t , t ))
Diffusion process의 가정인 ‘아주 작은 가우시안을 더해가는 과정’의 역과정은 ‘아주 작은 가우시안을 빼가는 과정’과 동일하기 때문에 각 포인트에서 노이즈를 예측한 다음 에 빼주면서 샘플링이 가능하다.
위에서 식으로 전개한 forward process와 지금 언급하고 있는 reverse process를 ‘같게’ 만드는 것이 loss function이 될 것이다.
Loss function
Evidence lower bound(ELBO) in DDPM
E ( − log p θ ( x 0 ) ) ≤ E q ( − log p θ ( x 0 ∣ z ) − log p θ ( z ) q ϕ ( z ∣ x 0 ) ) \mathbb{E}(-\log p_\theta (x_0)) \leq \mathbb{E}_q \left( -\log p_\theta(x_0 \vert z) -\log \frac{p_\theta(z)}{q_\phi(z \vert x_0)} \right) E ( − log p θ ( x 0 )) ≤ E q ( − log p θ ( x 0 ∣ z ) − log q ϕ ( z ∣ x 0 ) p θ ( z ) )
위의 식이 VAE에서 온건데, ELBO의 앞쪽 term은 z z z 로부터 x x x 를 reconstruction하는 부분에 대한 decoder 학습을 담당하고, 뒤쪽 term은 gaussian과 같은 pre-defined 분포로 가정한 p θ ( z ) p_\theta(z) p θ ( z ) 와 encoder의 output q ϕ ( z ∣ x ) q_\phi(z \vert x) q ϕ ( z ∣ x ) 의 분포가 가까워지게끔 하는 KL divergence에 해당된다. 위의 term에서 reconstruction과 KL divergence에서의 decoder 부분을 갈아끼우게 되면,
E ( − log p θ ( x 0 ) ) ≤ E q ( − log p θ ( z ) − log p θ ( x 0 ∣ z ) q ϕ ( z ∣ x 0 ) ) \mathbb{E}(-\log p_\theta (x_0)) \leq \mathbb{E}_q \left( -\log p_\theta(z) -\log \frac{p_\theta(x_0 \vert z)}{q_\phi(z \vert x_0)} \right) E ( − log p θ ( x 0 )) ≤ E q ( − log p θ ( z ) − log q ϕ ( z ∣ x 0 ) p θ ( x 0 ∣ z ) )
이처럼 되고, latent z z z 를 time step에 대한 변수로 치환하면서 encoder의 변수를 제거하게 되면(DDPM에서는 parametric 학습을 encoder에 대해 진행하지 않기 때문) 우리가 얻고자 하는 DDPM의 ELBO를 표현할 수 있다.
E ( − log p θ ( x 0 ) ) ≤ E q ( − log p θ ( x T ) − log p θ ( x 0 : T − 1 ∣ x T ) q ( x 1 : T ∣ x 0 ) ) \mathbb{E}(-\log p_\theta (x_0)) \leq \mathbb{E}_q \left( -\log p_\theta(x_T) -\log \frac{p_\theta(x_{0:T-1} \vert x_T)}{q(x_{1:T} \vert x_0)} \right) E ( − log p θ ( x 0 )) ≤ E q ( − log p θ ( x T ) − log q ( x 1 : T ∣ x 0 ) p θ ( x 0 : T − 1 ∣ x T ) )
그리고 각 process의 joint는 Markov process이기 때문에 summation으로 풀어쓸 수 있음.
E ( − log p θ ( x 0 ) ) ≤ E q ( − log p θ ( x T ) − log ∏ t ≥ 1 p θ ( x t − 1 ∣ x t ) q ( x t ∣ x t − 1 ) ) \mathbb{E}(-\log p_\theta (x_0)) \leq \mathbb{E}_q \left( -\log p_\theta(x_T) -\log \prod_{t \ge 1} \frac{p_\theta(x_{t-1} \vert x_t)}{q(x_t \vert x_{t-1})} \right) E ( − log p θ ( x 0 )) ≤ E q ( − log p θ ( x T ) − log t ≥ 1 ∏ q ( x t ∣ x t − 1 ) p θ ( x t − 1 ∣ x t ) )
뒤쪽 log \log log 랑 production 위치를 바꾸면 summation이 됨.
E ( − log p θ ( x 0 ) ) ≤ E q ( − log p θ ( x T ) − ∑ t ≥ 1 log p θ ( x t − 1 ∣ x t ) q ( x t ∣ x t − 1 ) ) \mathbb{E}(-\log p_\theta (x_0)) \leq \mathbb{E}_q \left( -\log p_\theta(x_T) - \sum_{t \ge 1}\log \frac{p_\theta(x_{t-1} \vert x_t)}{q(x_t \vert x_{t-1})} \right) E ( − log p θ ( x 0 )) ≤ E q ( − log p θ ( x T ) − t ≥ 1 ∑ log q ( x t ∣ x t − 1 ) p θ ( x t − 1 ∣ x t ) )
그렇다면 굳이 왜 기존의 VAE랑 다르게 위치를 바꾸었냐 라고 본다면 직관적으로는 forward process와 reverse process를 같게끔 학습하기 위함이라고 말할 수 있지만 보다 엄밀하게 표현하자면 VAE에서 학습시키는 decoder는 implicit conditional probability 인 p θ ( x 0 ∣ z ) p_\theta(x_0 \vert z) p θ ( x 0 ∣ z ) 를 학습하는 것이 주된 목적이었다면 Diffusion에서 학습시키는 decoder는 diffusion process 가 나타내는 미분 방정식을 학습 하는 것이 주된 목적이 되기 때문이다.
헌데 여기서 문제가 발생하는 것은 encoder 역할을 수행했던 diffusion process q q q 는 이전 step의 sample에 대해 노이즈를 더하는 과정 q ( x t ∣ x t − 1 ) q(x_t \vert x_{t-1}) q ( x t ∣ x t − 1 ) 을 수행했기 때문에 실제로 학습해야하는 그 역과정 q ( x t − 1 ∣ x t ) q(x_{t-1} \vert x_t) q ( x t − 1 ∣ x t ) 의 분포는 알 수 없다는(intractable) 것.
q ( x t − 1 ∣ x t ) = q ( x t ∣ x t − 1 ) q ( x t − 1 ) q ( x t ) q(x_{t-1} \vert x_t) = \frac{q(x_t \vert x_{t-1})q(x_{t-1})}{q(x_{t})} q ( x t − 1 ∣ x t ) = q ( x t ) q ( x t ∣ x t − 1 ) q ( x t − 1 )
본인은 이걸 한동안 이해 못했던거 같은데 쉽게 이해하는 법은 다음과 같음.
x t − 1 x_{t-1} x t − 1 에다가 더하는 노이즈를 알고 있기 때문에 x t x_t x t 에 대한 확률 분포는 normal distribution으로 명확하게 알 수가 있음( q ( x t ∣ x t − 1 ) ) (q(x_t \vert x_{t-1})) ( q ( x t ∣ x t − 1 )) 근데 실제로 각 time step에서 노이즈가 더해진 애들의 분포 자체는 알 수 없으니까 ( q ( x t ) , q ( x t − 1 ) ) (q(x_t),~q(x_{t-1})) ( q ( x t ) , q ( x t − 1 )) Bayes 식에서 intractable한 term이 두 개나 생겨서 분포에 접근하기가 힘듦.
그렇기 때문에 위의 식에 모든 term에 대해 x 0 x_0 x 0 을 조건부에 추가하면 항상 분포를 알 수 있게 되므로(normal distribution의 누적이니까) tractable하게 바꿀 수 있다.
q ( x t − 1 ∣ x t , x 0 ) = q ( x t ∣ x t − 1 , x 0 ) q ( x t − 1 , x 0 ) q ( x t , x 0 ) = q ( x t ∣ x t − 1 , x 0 ) q ( x t − 1 ∣ x 0 ) q ( x 0 ) q ( x t ∣ x 0 ) q ( x 0 ) = q ( x t ∣ x t − 1 , x 0 ) × q ( x t − 1 ∣ x 0 ) q ( x t ∣ x 0 ) \begin{aligned}
q(x_{t-1} \vert x_t,x_0) =& \frac{q(x_t \vert x_{t-1},x_0)q(x_{t-1},x_0)}{q(x_{t},x_0)} \newline
=& \frac{q(x_t \vert x_{t-1},x_0)q(x_{t-1}\vert x_0) q(x_0)}{q(x_{t} \vert x_0) q(x_0)} \newline
=& q(x_t \vert x_{t-1},x_0) \times \frac{q(x_{t-1} \vert x_0)}{q(x_t \vert x_0)}
\end{aligned} q ( x t − 1 ∣ x t , x 0 ) = = = q ( x t , x 0 ) q ( x t ∣ x t − 1 , x 0 ) q ( x t − 1 , x 0 ) q ( x t ∣ x 0 ) q ( x 0 ) q ( x t ∣ x t − 1 , x 0 ) q ( x t − 1 ∣ x 0 ) q ( x 0 ) q ( x t ∣ x t − 1 , x 0 ) × q ( x t ∣ x 0 ) q ( x t − 1 ∣ x 0 )
x 0 x_0 x 0 에 조건부가 된 식이라서, p θ ( x t ∣ x t − 1 ) p_\theta(x_t \vert x_{t-1}) p θ ( x t ∣ x t − 1 ) 와 달라진 것 아니냐는 생각이 들 수도 있지만 위의 식은 t > 1 t > 1 t > 1 인 모든 latent에 대해 x 0 x_0 x 0 에 독립 분포를 가진다는 특징이 있다(마르코프 프로세스니까). 그래서 최종 마무리가 된 loss function을 다음 식처럼 나타낸다.
D K L ( q ( x T ∣ x 0 ) ∣ ∣ p θ ( x T ) ) + ∑ t > 1 D K L ( q ( x t − 1 ∣ x t , x 0 ) ∣ ∣ p θ ( x t − 1 ∣ x t ) ) − E q ( log p θ ( x 0 ∣ x 1 ) ) D_{KL}(q(x_T \vert x_0) \vert\vert p_\theta(x_T)) +\sum_{t > 1} D_{KL} (q(x_{t-1} \vert x_t, x_0) \vert\vert p_\theta(x_{t-1} \vert x_t)) -\mathbb{E}_q\left(\log p_{\theta}(x_0 \vert x_1) \right) D K L ( q ( x T ∣ x 0 ) ∣∣ p θ ( x T )) + t > 1 ∑ D K L ( q ( x t − 1 ∣ x t , x 0 ) ∣∣ p θ ( x t − 1 ∣ x t )) − E q ( log p θ ( x 0 ∣ x 1 ) )
유도 과정은 앞서 본 식에서 t > 1 t>1 t > 1 인 부분에 대해 posterior를 조건부로 바꿔서 tractable하게 만들 수 있으니까, 정리하게 되면 다음과 같이 나온다.
L ≤ E q ( − log ( p θ ( x T ) ) − ∑ t = 2 T log p θ ( x t − 1 ∣ x t ) q ( x t ∣ x t − 1 ) − log p θ ( x 0 ∣ x 1 ) q ( x 1 ∣ x 0 ) ) ≤ E q ( − log ( p θ ( x T ) ) − ∑ t = 2 T log ( p θ ( x t − 1 ∣ x t ) q ( x t − 1 ∣ x t , x 0 ) × q ( x t − 1 ∣ x 0 ) q ( x t ∣ x 0 ) ) − log p θ ( x 0 ∣ x 1 ) q ( x 1 ∣ x 0 ) ) ≤ ∑ t > 1 D K L ( q ( x t − 1 ∣ x t , x 0 ) ∣ ∣ p θ ( x t − 1 ∣ x t ) ) + E q ( − log ( p θ ( x T ) ) − log q ( x 1 ∣ x 0 ) q ( x T ∣ x 0 ) − log p θ ( x 0 ∣ x 1 ) q ( x 1 ∣ x 0 ) ) ≤ D K L ( q ( x T ∣ x 0 ) ∣ ∣ p θ ( x T ) ) + ∑ t > 1 D K L ( q ( x t − 1 ∣ x t , x 0 ) ∣ ∣ p θ ( x t − 1 ∣ x t ) ) − E q ( log p θ ( x 0 ∣ x 1 ) ) \begin{aligned}
\mathcal{L} \le& \mathbb{E}_q\left(-\log (p_\theta(x_T))-\sum_{t=2}^T \log \frac{p_\theta(x_{t-1} \vert x_t)}{q(x_t \vert x_{t-1})} -\log \frac{p_\theta(x_0 \vert x_1)}{q(x_1 \vert x_0)} \right) \newline
\le& \mathbb{E}_q\left(-\log (p_\theta(x_T))-\sum_{t=2}^T \log \left( \frac{p_\theta(x_{t-1} \vert x_t)}{q(x_{t-1} \vert x_t, x_0)} \times \frac{q(x_{t-1} \vert x_0)}{q(x_t\vert x_0)}\right) -\log \frac{p_\theta(x_0 \vert x_1)}{q(x_1 \vert x_0)} \right) \newline
\le& \sum_{t > 1} D_{KL} (q(x_{t-1} \vert x_t, x_0) \vert\vert p_\theta(x_{t-1} \vert x_t))+\mathbb{E}_q\left(-\log (p_\theta(x_T))-\log \frac{q(x_1 \vert x_0)}{q(x_T \vert x_0)}-\log \frac{p_\theta(x_0 \vert x_1)}{q(x_1 \vert x_0)} \right) \newline
\le& D_{KL}(q(x_T \vert x_0) \vert\vert p_\theta(x_T)) +\sum_{t > 1} D_{KL} (q(x_{t-1} \vert x_t, x_0) \vert\vert p_\theta(x_{t-1} \vert x_t)) -\mathbb{E}_q\left(\log p_{\theta}(x_0 \vert x_1) \right)
\end{aligned} L ≤ ≤ ≤ ≤ E q ( − log ( p θ ( x T )) − t = 2 ∑ T log q ( x t ∣ x t − 1 ) p θ ( x t − 1 ∣ x t ) − log q ( x 1 ∣ x 0 ) p θ ( x 0 ∣ x 1 ) ) E q ( − log ( p θ ( x T )) − t = 2 ∑ T log ( q ( x t − 1 ∣ x t , x 0 ) p θ ( x t − 1 ∣ x t ) × q ( x t ∣ x 0 ) q ( x t − 1 ∣ x 0 ) ) − log q ( x 1 ∣ x 0 ) p θ ( x 0 ∣ x 1 ) ) t > 1 ∑ D K L ( q ( x t − 1 ∣ x t , x 0 ) ∣∣ p θ ( x t − 1 ∣ x t )) + E q ( − log ( p θ ( x T )) − log q ( x T ∣ x 0 ) q ( x 1 ∣ x 0 ) − log q ( x 1 ∣ x 0 ) p θ ( x 0 ∣ x 1 ) ) D K L ( q ( x T ∣ x 0 ) ∣∣ p θ ( x T )) + t > 1 ∑ D K L ( q ( x t − 1 ∣ x t , x 0 ) ∣∣ p θ ( x t − 1 ∣ x t )) − E q ( log p θ ( x 0 ∣ x 1 ) )
ELBO into objective function
위의 식은 수식일 뿐이고 실제로 최적화 과정에서 사용되는 식은 이를 가우시안 분포에 대해 다시 정리한 식이다. 위의 ELBO를 보게 되면 크게 세 부분으로 정리되는데, 이 중 가장 앞부분의 KL divergence는 diffusion process를 통해 자연스럽게 만족되는 식 (gaussian 분포를 따라가는 것)이기 때문에 실제 loss function에서 제외된다. 원래 VAE에서는 해당 부분이 encoder의 output을 가우시안 분포로 정규화하는 KL divergence term으로 사용되는데, diffusion에서는 굳이 q q q 를 학습시킬 필요가 없다보니 빠진다고 생각할 수 있다.
D K L ( q ( x T ∣ x 0 ) ∣ ∣ p θ ( x T ) ) D_{KL}(q(x_T \vert x_0) \vert\vert p_\theta(x_T)) D K L ( q ( x T ∣ x 0 ) ∣∣ p θ ( x T ))
그리고 가장 마지막 부분인 log likelihood 식이 가장 이해하기 어려운데, 이를 먼저 언급하고 넘어가도록 하겠다.
E q ( log p θ ( x 0 ∣ x 1 ) ) \mathbb{E}_q\left(\log p_{\theta}(x_0 \vert x_1) \right) E q ( log p θ ( x 0 ∣ x 1 ) )
이 식을 다시금 이해하자면, noise가 아주 조금 더해진 x 1 x_1 x 1 에서 x 0 x_0 x 0 로 reconstruction할 때의 모든 확률을 구해야 한다. 중간 과정의 경우에는 분포 사이의 KL divergence 식이기 때문에 확률 분포를 굳이 확률값으로 매핑할 필요가 없었지만, 이 식에서는 확률값 자체를 구해야하기 때문에 적분을 취해야 한다.
우선 단순화시켜 생각하기 위해, log \log log likelihood의 likelihood 자체만 따로 빼내서 보도록 하자.
p θ ( x 0 ∣ x 1 ) p_\theta(x_0 \vert x_1) p θ ( x 0 ∣ x 1 )
p θ p_\theta p θ 는 언급했던 것과 같이 x 1 x_1 x 1 를 통해 x 0 x_0 x 0 을 구성하는 각 요소들을 추출할 확률 분포를 구하는 diffusion process의 역과정을 학습한 neural network이다.
p θ ( x 0 ∣ x 1 ) ∼ N ( x 0 ; μ θ ( x 1 , 1 ) , σ 1 2 ) p_\theta(x_0 \vert x_1) \sim \mathcal{N}(x_0;~\mu_\theta(x_1, 1), \sigma_1^2) p θ ( x 0 ∣ x 1 ) ∼ N ( x 0 ; μ θ ( x 1 , 1 ) , σ 1 2 )
확률값을 구하기 위해서는 이미지 상의 ‘모든 픽셀’에 대해 주어진 확률 분포를 적분해야만 구할 수 있다. 이때, 이미지가 [ − 1 , 1 ] [-1, 1] [ − 1 , 1 ] 로 정규화된 상태라고 생각한다면 가우시안 분포에서 − 1 -1 − 1 보다 작은 value는 모두 − 1 -1 − 1 로 확률을 매핑하고 1 1 1 보다 큰 value는 모두 1 1 1 로 확률을 매핑한다고 생각할 수 있다. 또한 원래 이미지 전체는 0 ∼ 255 0 \sim 255 0 ∼ 255 의 discrete RGB value로 표현되기 때문에 다음과 같이 변수를 1 / 255 1/255 1/255 로 끊어서 확률매핑이 가능하다.
예를 들어 x 1 x_1 x 1 에서 x 0 x_0 x 0 으로 매핑될 때, ground truth에 따르면 0 0 0 이라는 값이 나와야한다고 가정해보자.
그렇다면 p θ ( x 0 ∣ x 1 ) p_\theta(x_0 \vert x_1) p θ ( x 0 ∣ x 1 ) 는 해당 값 근방을 적분 구간으로 삼아 확률분포를 적분한 값으로 정의할 수 있는 것이다(아래 그림에서 노란색 부분의 넓이).
이를 모든 픽셀에 대해 곱하게 되면, 전체 픽셀 (1 ∼ D 1\sim D 1 ∼ D )에 대한 joint probability를 구할 수 있다.
p θ ( x 0 ∣ x 1 ) = ∏ i = 1 D ∫ δ − ( x 0 i ) δ + ( x 0 i ) N ( x ; μ θ i ( x 1 , 1 ) , σ 1 2 ) d x p_\theta(x_0 \vert x_1) = \prod_{i=1}^D \int_{\delta_{-}(x_0^i)}^{\delta_+(x_0^i)} \mathcal{N}(x; \mu_\theta^i (x_1, 1), \sigma_1^2) dx p θ ( x 0 ∣ x 1 ) = i = 1 ∏ D ∫ δ − ( x 0 i ) δ + ( x 0 i ) N ( x ; μ θ i ( x 1 , 1 ) , σ 1 2 ) d x
이때 적분 구간인 [ − δ − ( x 0 i ) , + δ + ( x 0 i ) ] [-\delta_-(x_0^i),~+\delta_+(x_0^i)] [ − δ − ( x 0 i ) , + δ + ( x 0 i )] 는 앞서 말했던 것과 같이 − 1 -1 − 1 보다 작은 값들과 1 1 1 보다 큰 값들을 각각 원래 이미지의 − 1 -1 − 1 과 1 1 1 에 매핑하고 나머지는 1 / 255 1/255 1/255 만큼 간격을 준다고 생각할 수 있다.
δ + ( x ) = { ∞ , if x = 1 x + 1 255 , if x < 1 , δ − ( x ) = { − ∞ , if x = − 1 x − 1 255 , if x > − 1 \delta_+(x) = \begin{cases}
\infty,&\text{if }x=1 \newline
x+\frac{1}{255},&\text{if }x<1
\end{cases},~~\delta_-(x) = \begin{cases}
-\infty,&\text{if }x=-1 \newline
x-\frac{1}{255},&\text{if }x>-1
\end{cases} δ + ( x ) = { ∞ , x + 255 1 , if x = 1 if x < 1 , δ − ( x ) = { − ∞ , x − 255 1 , if x = − 1 if x > − 1
중간과정 Optimization(L 1 : T − 1 L_{1:T-1} L 1 : T − 1 )
결국 위에서 쭉 말한 내용은 L T L_T L T 는 굳이 필요없다는 내용이었고 L 0 L_0 L 0 는 확률 분포의 적분을 통해 구할 수 있다는 말이었다. 그렇다면 실제로 더해진 노이즈를 예측해서 빼는 과정인 L 1 : T − 1 L_{1:T-1} L 1 : T − 1 를 KL divergence 식으로 구하면 어떻게 될까?
∑ t > 1 D K L ( q ( x t − 1 ∣ x t , x 0 ) ∣ ∣ p θ ( x t − 1 ∣ x t ) ) \sum_{t > 1} D_{KL} (q(x_{t-1} \vert x_t, x_0) \vert\vert p_\theta(x_{t-1} \vert x_t)) t > 1 ∑ D K L ( q ( x t − 1 ∣ x t , x 0 ) ∣∣ p θ ( x t − 1 ∣ x t ))
해당 식에서 q ( x t − 1 ∣ x t , x 0 ) q(x_{t-1} \vert x_t, x_0) q ( x t − 1 ∣ x t , x 0 ) 를 가우시안 분포로 나타내는 증명과정은 다음과 같다. 앞서 유도한 바와 같이 기존의 prior를 다음과 같이 고쳤었다.
q ( x t − 1 ∣ x t , x 0 ) = q ( x t ∣ x t − 1 , x 0 ) q ( x t − 1 ∣ x 0 ) q ( x t ∣ x 0 ) q(x_{t-1} \vert x_t,x_0) = q(x_t \vert x_{t-1},x_0) \frac{q(x_{t-1} \vert x_0)}{q(x_t \vert x_0)} q ( x t − 1 ∣ x t , x 0 ) = q ( x t ∣ x t − 1 , x 0 ) q ( x t ∣ x 0 ) q ( x t − 1 ∣ x 0 )
그리고 우리는 식을 구성하는 각각의 분포를 다음과 같은 tractable한 가우시안 형태로 알고 있다.
q ( x t ∣ x 0 ) ∼ N ( α ˉ t x 0 , ( 1 − α ˉ t ) I ) q ( x t ∣ x t − 1 ) ∼ N ( 1 − β t x t − 1 , β t I ) \begin{aligned}
q(x_t \vert x_0) \sim& \mathcal{N}(\sqrt{\bar{\alpha}_t}x_0, (1-\bar{\alpha}_t)I) \newline
q(x_t \vert x_{t-1}) \sim& \mathcal{N}(\sqrt{1-\beta_t}x_{t-1},\beta_tI)
\end{aligned} q ( x t ∣ x 0 ) ∼ q ( x t ∣ x t − 1 ) ∼ N ( α ˉ t x 0 , ( 1 − α ˉ t ) I ) N ( 1 − β t x t − 1 , β t I )
오호라 위의 식을 그대로 가우시안 식으로 바꿔쓰게 되면 가우시안 앞부분에 붙는 상수는 무시하고 exponential 안의 식만 따로 정리할 수 있다.
q ( x t − 1 ∣ x t , x 0 ) ∝ exp ( − 1 2 ( ( x t − α t x t − 1 ) 2 β t + ( x t − 1 − α ˉ t − 1 x 0 ) 2 1 − α ˉ t − 1 − ( x t − α ˉ t x 0 ) 2 1 − α ˉ t ) ) = exp ( − 1 2 ( x t 2 − 2 α t x t x t − 1 + α t x t − 1 2 β t + x t − 1 2 − 2 α ˉ t − 1 x 0 x t − 1 + α ˉ t − 1 x 0 2 1 − α ˉ t − 1 − x t 2 − 2 α ˉ t x 0 x t + α ˉ t x 0 2 1 − α ˉ t ) ) \begin{aligned}
q(x_{t-1} \vert x_t, x_0) \propto& \exp \left(-\frac{1}{2} \left( \frac{(x_t - \sqrt{\alpha_t}x_{t-1})^2}{\beta_t} + \frac{(x_{t-1} - \sqrt{\bar{\alpha}_{t-1}}x_0)^2}{1-\bar{\alpha}_{t-1}} - \frac{(x_t - \sqrt{\bar{\alpha}_t}x_0)^2}{1-\bar{\alpha}_t} \right) \right) \newline
=& \exp \left(-\frac{1}{2} \left( \frac{x_t^2 - 2\sqrt{\alpha_t}x_t x_{t-1}+\alpha_tx_{t-1}^2}{\beta_t} + \frac{x_{t-1}^2 - 2\sqrt{\bar{\alpha}_{t-1}}x_0x_{t-1} + \bar{\alpha}_{t-1}x_0^2}{1-\bar{\alpha}_{t-1}} - \frac{x_t^2 -2\sqrt{\bar{\alpha}_t}x_0x_t + \bar{\alpha}_tx_0^2}{1-\bar{\alpha}_t}\right) \right)
\end{aligned} q ( x t − 1 ∣ x t , x 0 ) ∝ = exp ( − 2 1 ( β t ( x t − α t x t − 1 ) 2 + 1 − α ˉ t − 1 ( x t − 1 − α ˉ t − 1 x 0 ) 2 − 1 − α ˉ t ( x t − α ˉ t x 0 ) 2 ) ) exp ( − 2 1 ( β t x t 2 − 2 α t x t x t − 1 + α t x t − 1 2 + 1 − α ˉ t − 1 x t − 1 2 − 2 α ˉ t − 1 x 0 x t − 1 + α ˉ t − 1 x 0 2 − 1 − α ˉ t x t 2 − 2 α ˉ t x 0 x t + α ˉ t x 0 2 ) )
조금 복잡해 보이긴 하는데 여기서 x t − 1 x_{t-1} x t − 1 부분에 대해서 quadratic(2차 함수) 형태가 되도록 공통 변수로 묶어주게 되면,
= exp ( − 1 2 ( ( α t β t + 1 1 − α ˉ t − 1 ) x t − 1 2 − ( 2 α t x t β t + 2 α ˉ t − 1 x 0 1 − α ˉ t − 1 ) x t − 1 + C ( x t , x 0 ) ) ) =\exp \left(-\frac{1}{2} \left( ( \frac{\alpha_t}{\beta_t} +\frac{1}{1-\bar{\alpha}_{t-1}} )x_{t-1}^2 - (2\frac{\sqrt{\alpha}_tx_t}{\beta_t} + 2\frac{\sqrt{\bar{\alpha}_{t-1}}x_0}{1-\bar{\alpha}_{t-1}})x_{t-1} + C(x_t, x_0) \right) \right) = exp ( − 2 1 ( ( β t α t + 1 − α ˉ t − 1 1 ) x t − 1 2 − ( 2 β t α t x t + 2 1 − α ˉ t − 1 α ˉ t − 1 x 0 ) x t − 1 + C ( x t , x 0 ) ) )
위와 같이 정리할 수 있다. 해당 식을 가우시안 분포라고 생각하게 되면 평균은 a x 2 + b x + c ax^2+bx+c a x 2 + b x + c 형태의 이차식에서 − b / 2 a -b/2a − b /2 a 와 같기 때문에
μ ~ ( x t , x 0 ) = α t x t β t + α ˉ t − 1 x 0 1 − α ˉ t − 1 α t β t + 1 1 − α ˉ t − 1 = ( α t x t β t + α ˉ t − 1 x 0 1 − α ˉ t − 1 ) × β t ( 1 − α ˉ t − 1 ) 1 − α ˉ t = α t ( 1 − α ˉ t − 1 ) 1 − α ˉ t x t + β t α ˉ t − 1 1 − α ˉ t x 0 \begin{aligned}
\tilde{\mu}(x_t,~x_0) =& \frac{\frac{\sqrt{\alpha}_tx_t}{\beta_t} + \frac{\sqrt{\bar{\alpha}_{t-1}}x_0}{1-\bar{\alpha}_{t-1}}}{\frac{\alpha_t}{\beta_t} +\frac{1}{1-\bar{\alpha}_{t-1}}} = \left( \frac{\sqrt{\alpha}_tx_t}{\beta_t} + \frac{\sqrt{\bar{\alpha}_{t-1}}x_0}{1-\bar{\alpha}_{t-1}}\right) \times \frac{\beta_t (1-\bar{\alpha}_{t-1})}{1-\bar{\alpha}_t} \newline
=& \frac{\sqrt{\alpha}_t(1-\bar{\alpha}_{t-1})}{1-\bar{\alpha}_t}x_t + \frac{\beta_t \sqrt{\bar{\alpha}_{t-1}}}{1-\bar{\alpha}_t}x_0
\end{aligned} μ ~ ( x t , x 0 ) = = β t α t + 1 − α ˉ t − 1 1 β t α t x t + 1 − α ˉ t − 1 α ˉ t − 1 x 0 = ( β t α t x t + 1 − α ˉ t − 1 α ˉ t − 1 x 0 ) × 1 − α ˉ t β t ( 1 − α ˉ t − 1 ) 1 − α ˉ t α t ( 1 − α ˉ t − 1 ) x t + 1 − α ˉ t β t α ˉ t − 1 x 0
위와 같이 나타낼 수 있고 분산은 1 / ( − 2 a ) 1/(-2a) 1/ ( − 2 a ) 과 같기 때문에
σ ~ t 2 = β t × ( 1 − α ˉ t − 1 1 − α ˉ t ) \tilde{\sigma}_t^2 = \beta_t \times \left( \frac{1-\bar{\alpha}_{t-1}}{1-\bar{\alpha}_t} \right) σ ~ t 2 = β t × ( 1 − α ˉ t 1 − α ˉ t − 1 )
이처럼 나타낼 수 있다. DDPM 논문에서는 해당 부분을 제대로 증명하지 않고 넘어가서 갑자기 붕 떴던 기억때문에 남겨놓는다.
이제 실제로 p θ ( x t − 1 ∣ x t ) p_\theta(x_{t-1} \vert x_t) p θ ( x t − 1 ∣ x t ) 와 q ( x t − 1 ∣ x t , x 0 ) q(x_{t-1} \vert x_t, x_0) q ( x t − 1 ∣ x t , x 0 ) 간의 KL divergence를 최소화하는 식을 살펴보면 N ( x t − 1 ; μ θ ( x t , t ) , Σ θ ( x t , t ) ) \mathcal{N}(x_{t-1}; \mu_\theta(x_t,t),~\Sigma_\theta(x_t, t)) N ( x t − 1 ; μ θ ( x t , t ) , Σ θ ( x t , t )) 는 해당 논문에서 미리 스케쥴링해서 학습할 파라미터로 설정하지 않았기 때문에 σ ~ t 2 = β t × ( 1 − α ˉ t − 1 1 − α ˉ t ) \tilde{\sigma}_t^2 = \beta_t \times \left( \frac{1-\bar{\alpha}_{t-1}}{1-\bar{\alpha}_t} \right) σ ~ t 2 = β t × ( 1 − α ˉ t 1 − α ˉ t − 1 ) 라고 할 수 있다. 실제 실험에서는 단순히 t t t 번째 step에서의 variance인 β t \beta_t β t 를 σ t 2 \sigma_t^2 σ t 2 로 사용해도 큰 차이가 없다고 언급하는데, 이는 사실상 variance 누적곱에 해당되는 α ˉ t = ∏ τ = 1 t ( 1 − β τ ) \bar{\alpha}_t = \prod_{\tau=1}^t(1-\beta_\tau) α ˉ t = ∏ τ = 1 t ( 1 − β τ ) 가 step 수가 크기 때문에 t t t 에 따라 큰 차이를 보이지 않기 때문이라고 생각했다(아니라면 말고). 암튼 결론은 μ θ ( x t , t ) \mu_\theta(x_t,~t) μ θ ( x t , t ) 만 예측하면 된다는 소리..
KL divergence 식을 가우시안에 적용하게 되면 exponential이 벗겨지면서 다음과 같이 간단하게 표현 가능하다. 아래 식에서 C C C 는 파라미터에 무관한 모든 term을 의미한다.
L t − 1 : = E q ( 1 2 σ t 2 ∣ ∣ μ ~ t ( x t , x 0 ) − μ θ ( x t , t ) ∣ ∣ 2 ) + C L_{t-1} := \mathbb{E}_q \left( \frac{1}{2\sigma_t^2} \vert\vert \tilde{\mu}_t(x_t, x_0) - \mu_\theta(x_t, t) \vert\vert^2 \right)+C L t − 1 := E q ( 2 σ t 2 1 ∣∣ μ ~ t ( x t , x 0 ) − μ θ ( x t , t ) ∣ ∣ 2 ) + C
근데 이 식을 그대로 쓰면 stochastic 부분땜에 미분이 불가능하다. 써 있는 x t x_t x t 를 죄다 reparameterization(x t = α ˉ t x 0 + 1 − α ˉ t ϵ x_t = \sqrt{\bar{\alpha}_t}x_0+\sqrt{1-\bar{\alpha}_{t}}\epsilon x t = α ˉ t x 0 + 1 − α ˉ t ϵ )으로 바꾸면 된다. 근데 이렇게 쓰면 너무 장황하니까 x t ( x 0 , ϵ ) x_t(x_0,\epsilon) x t ( x 0 , ϵ ) 으로 표현하고자 함. 참고로 μ ~ t \tilde{\mu}_t μ ~ t 에 있는 x 0 x_0 x 0 도 process 상에서 x t x_t x t 를 통해 예측되는 x 0 x_0 x 0 이기 때문에 reparameterization하여 x 0 = x t α ˉ t − 1 − α ˉ t α ˉ t ϵ x_0 = \frac{x_t}{\sqrt{\bar{\alpha}_t}}-\frac{\sqrt{1-\bar{\alpha}_t}}{\sqrt{\bar{\alpha}_t}}\epsilon x 0 = α ˉ t x t − α ˉ t 1 − α ˉ t ϵ 로 쓰면 된다.
L t − 1 − C = E q ( 1 2 σ t 2 ∣ ∣ μ ~ t ( x t ( x 0 , ϵ ) , x t ( x 0 , ϵ ) α ˉ t − 1 − α ˉ t α ˉ t ϵ ) − μ θ ( x t ( x 0 , ϵ ) , t ) ∣ ∣ 2 ) L_{t-1} -C = \mathbb{E}_q \left( \frac{1}{2\sigma_t^2} \left\vert\left\vert \tilde{\mu}_t\left(x_t(x_0,\epsilon), \frac{x_t(x_0,\epsilon)}{\sqrt{\bar{\alpha}_t}}-\frac{\sqrt{1-\bar{\alpha}_t}}{\sqrt{\bar{\alpha}_t}}\epsilon\right) - \mu_\theta(x_t(x_0,\epsilon), t) \right\vert\right\vert^2 \right) L t − 1 − C = E q ( 2 σ t 2 1 μ ~ t ( x t ( x 0 , ϵ ) , α ˉ t x t ( x 0 , ϵ ) − α ˉ t 1 − α ˉ t ϵ ) − μ θ ( x t ( x 0 , ϵ ) , t ) 2 )
앞서 구했던 식에 대입하면,
μ ~ ( x t , x 0 ) = α t ( 1 − α ˉ t − 1 ) 1 − α ˉ t x t + β t α ˉ t − 1 1 − α ˉ t x 0 = α t ( 1 − α ˉ t − 1 ) 1 − α ˉ t x t ( x 0 , ϵ ) + β t α ˉ t − 1 1 − α ˉ t ( x t ( x 0 , ϵ ) α ˉ t − 1 − α ˉ t α ˉ t ϵ ) = 1 α t x t ( x 0 , ϵ ) − β t ( 1 − α ˉ t ) α t ϵ = 1 α t ( x t ( x 0 , ϵ ) − β t 1 − α ˉ t ϵ ) \begin{aligned}
\tilde{\mu}(x_t,~x_0) =& \frac{\sqrt{\alpha}_t(1-\bar{\alpha}_{t-1})}{1-\bar{\alpha}_t}x_t + \frac{\beta_t \sqrt{\bar{\alpha}_{t-1}}}{1-\bar{\alpha}_t}x_0 \newline
=& \frac{\sqrt{\alpha}_t(1-\bar{\alpha}_{t-1})}{1-\bar{\alpha}_t}x_t(x_0,\epsilon) + \frac{\beta_t \sqrt{\bar{\alpha}_{t-1}}}{1-\bar{\alpha}_t}\left(\frac{x_t(x_0,\epsilon)}{\sqrt{\bar{\alpha}_t}}-\frac{\sqrt{1-\bar{\alpha}_t}}{\sqrt{\bar{\alpha}_t}}\epsilon\right) \newline
=& \frac{1}{\sqrt{\alpha_t}}x_t(x_0, \epsilon)-\frac{\beta_t}{\sqrt{(1-\bar{\alpha}_t)\alpha_t}}\epsilon \newline
=& \frac{1}{\sqrt{\alpha_t}}\left( x_t(x_0, \epsilon)-\frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}}\epsilon\right)
\end{aligned} μ ~ ( x t , x 0 ) = = = = 1 − α ˉ t α t ( 1 − α ˉ t − 1 ) x t + 1 − α ˉ t β t α ˉ t − 1 x 0 1 − α ˉ t α t ( 1 − α ˉ t − 1 ) x t ( x 0 , ϵ ) + 1 − α ˉ t β t α ˉ t − 1 ( α ˉ t x t ( x 0 , ϵ ) − α ˉ t 1 − α ˉ t ϵ ) α t 1 x t ( x 0 , ϵ ) − ( 1 − α ˉ t ) α t β t ϵ α t 1 ( x t ( x 0 , ϵ ) − 1 − α ˉ t β t ϵ )
위와 같다. 따라서 다시 식에 해당 값을 대입하게 되면,
E q ( 1 2 σ t 2 ∣ ∣ 1 α t ( x t ( x 0 , ϵ ) − β t 1 − α ˉ t ϵ ) − μ θ ( x t ( x 0 , ϵ ) , t ) ∣ ∣ 2 ) \mathbb{E}_q \left( \frac{1}{2\sigma_t^2} \left\vert\left\vert \frac{1}{\sqrt{\alpha_t}}\left( x_t(x_0, \epsilon)-\frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}}\epsilon\right) - \mu_\theta(x_t(x_0,\epsilon), t) \right\vert\right\vert^2 \right) E q ( 2 σ t 2 1 α t 1 ( x t ( x 0 , ϵ ) − 1 − α ˉ t β t ϵ ) − μ θ ( x t ( x 0 , ϵ ) , t ) 2 )
요렇게 됨. 그런데 여기서 네트워크는 굳이 x t ( x 0 , ϵ ) x_t(x_0, \epsilon) x t ( x 0 , ϵ ) 를 구할 필요가 없게 된다. 왜냐면 time step-t t t 에서는 이미 input으로 x t x_t x t 가 주어지고, x t − 1 x_{t-1} x t − 1 를 만들기 위한 분포 예측 과정이므로 forward 과정에서 stochastic하게 더해진 epsilon 부분만 예측하면 되는 것.
μ θ ( x t , t ) = 1 α t ( x t − β t 1 − α ˉ t ϵ θ ( 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) μ θ ( x t , t ) = α t 1 ( x t − 1 − α ˉ t β t ϵ θ ( x t , t ) )
따라서 x t x_t x t 에서 x t − 1 x_{t-1} x t − 1 를 샘플링할 때 계산하는 것은 x t − 1 x_{t-1} x t − 1 에 대해 예측된 평균인 1 α t ( x t − β t 1 − α ˉ t ϵ θ ( 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) α t 1 ( x t − 1 − α ˉ t β t ϵ θ ( x t , t ) ) 에 사전에 정의된 variance term인 σ t z \sigma_tz σ t z , z ∼ N ( 0 , I ) z\sim\mathcal{N}(0, I) z ∼ N ( 0 , I ) 을 더하면 된다.
x t − 1 = 1 α t ( x t − β t 1 − α ˉ t ϵ θ ( x t , t ) ) + σ t z 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_tz x t − 1 = α t 1 ( x t − 1 − α ˉ t β t ϵ θ ( x t , t ) ) + σ t z
그렇기 때문에 L 1 : T − 1 L_{1:T-1} L 1 : T − 1 에 대한 loss는 아주 간단하게 표현 가능하다(x t x_t x t 에 대한 부분은 제외하고 예측하면됨).
E x 0 , ϵ ( β t 2 2 σ t 2 α t ( 1 − α ˉ t ) ∥ ϵ − ϵ θ ( α ˉ t x 0 + 1 − α ˉ t ϵ , t ) ∥ 2 ) \mathbb{E}_{x_0,\epsilon} \left( \frac{\beta_t^2}{2\sigma_t^2 \alpha_t(1-\bar{\alpha}_t)} \parallel \epsilon - \epsilon_\theta(\sqrt{\bar{\alpha}_t}x_0 + \sqrt{1-\bar{\alpha}_t}\epsilon, t) \parallel^2\right) E x 0 , ϵ ( 2 σ t 2 α t ( 1 − α ˉ t ) β t 2 ∥ ϵ − ϵ θ ( α ˉ t x 0 + 1 − α ˉ t ϵ , t ) ∥ 2 )