DDPM (NeurIPS 2020, Diffusion)
DDPM
Background
* diffusion process는 원래 종전의 [1]을 발전시킨 모델으로, variational inference를 이용하여 train된 parametrized Markov chain으로 image를 generation하는 과정이다.
→ variational inference는 posterior probability $p(x|z)$를 직접 구하는 대신 parametrized function $q(x)$를 대신 구하고, ELBO를 이용하여 둘의 차이를 minimize하는 것을 의미한다. 이때 ELBO를 maximize하면 둘의 KL divergence $D_{KL}(p(x|z) || q(x))$를 minimize하는 것과 같다.
물리학에서 말하는 diffusion처럼, diffusion process는 원래 structure에서 diffuse되어 noise를 만들고, 이를 역과정으로도 수행하여 noise로부터 image를 generate할 수 있게 된다. 각각의 markov chain process는 data에 작은 양의 Gaussian noise를 더한다. 이는 conditional Gaussian, 즉 neural network로 parametrize될 수 있다.
생성된 image를 $x_0$, signal이 제거된 image를 $x_T$라고 했을 때, diffusion model은 $p_θ(x_0):=\int p_θ(x_{0:T})dx_{1:T}$로 정의될 수 있다. 이 joint distribution $p_θ(x_{0:T}$는 reverse process라고 한다. 이때 각각의 step의 Guassian translation을 수행하는 Markov chain으로 정의된다:
$$ p_θ(x_{0:T}):=p(x_T)\prod^T_{t=1}p_θ(x_{t-1}|x_t),\qquad p_θ(x_{t-1}|x_t):=\mathcal{N} (x_{t-1};\mu_θ(x_t,t),Σ_θ(x_t,t))$$
→ 따라서 $x_0$는 $x_T$에서 여러 번의 Gaussian을 여러 번 곱해준 형태가 되는 것이다.
원래 data에 noise를 더하는 과정을 forward process 또는 diffusion process라고 하며 이는 원래 데이터 $x_0$에 variance schedule $β_t$에 따라서 Gaussian noise를 더해주는 Markov chain이다.
$$ q(x_{1:T}|x_0):=\prod^T_{t=1}q(x_t|x_{t-1}),\qquad q(x_t|x_{t-1}):=\mathcal {N} (x_t;\sqrt{1-β_t} x_{t-1},β_t I)$$
→ 이때 mean에 원래 signal에 $\sqrt{1-β_t}$를 곱해주는 것은 variance 값이 convex combination으로 구성되게 하기 위한 것이다.
Training은 ELBO를 minimize함으로써 수행될 수 있다.
$$\mathbb{E} [-\log p_θ(x_0)] ≤ \mathbb E_q \left [ -\log \frac{p_θ(x_{0:T}}{q)x_{1:T}|x_0} \right] \\ = \mathbb E_q \left[ -\log p(x_T)-\sum_{t≥1}\log\frac{p_θ(x_{t-1}|x_t}{q(x_t|x_{t-1})}\right ] =: \mathcal L$$
이때 각 variance는 reparametrization에 의해서 train될 수도 있고, hyperparameter로 설정될 수도 있다. 또한 이때 $x_t$의 sampling은 $x_0$을 mean으로 하는 Guassian에서의 sampling의 형태를 갖게 된다.
$$ q(x_t|x_0)=\mathcal{N}$$
이때 $α_t:=1-β_t, \qquad \bar α_t:=\prod^t_{s=1}α_s$
Forward Process & Reverse Process
- Forward process
forward process는 모든 $β_t$가 constant로 fix된 형태이다. 따라서 $q$는 learnable parameter가 없다.
- Backward process
backward process는 neural network로 $μ, Σ$를 학습하여 이루어진다. 처음에는 $x_0 \sim \mathcal {N}(0,I)$를 하게 되고,
위의 형태로 backward process를 정의한다. 이를 optimize하기 위해서 위의 ELBO term을 optimize해야 한다. 이때 $L_0$와 $L_T$는 optimize하지 않고 $L_{t-1}$만 optimize하게 된다. 이는 $L_T$에는 training이 가능한 parameter가 없고 $L_0$는 제외했을 때 performance가 더 좋았기 때문이다.
따라서 loss term은 다음과 같이 쓰여질 수 있다.
그러나 이를 더 간단히 하기 위해서 reparametrization trick을 사용해서 $ε~\mathcal {N}(0,I)$로 표현한다.
이를 추가 simplify하는 과정을 거쳐 다음과 같이 표현하게 된다.
이는 곧 $ε$을 predict하는 문제로 바뀔 수 있음을 의미한다. 물론 $x_0$를 predict할 수 있지만, 이 경우 performance가 떨어지는 문제가 있었다.
최종적으로는 다음 loss를 optimize하게 된다:
이때 t는 1과 T 사이에서 uniform하다.
이를 정리하면 다음 algorithm으로 표현할 수 있다:
Discussion
References
[1] Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning, pages 2256–2265, 2015.