Denoising Diffusion Probabilistic Models (DDPM)
Diffusion model의 아이디어는 잉크가 물속에서 퍼져나가는 과정 (forward process)과 이를 역으로 되돌리는 과정 (reverse process)으로 이루어져 있다. 이미 퍼져나간 잉크를 다시 한데 모으는 것은 매우 어려울 것이다. 그러나 시간을 잘게 쪼개면 순간의 입자 운동은 예측 가능할 것이고, 거꾸로 모으는 것도 가능할 수 있다.
Objective
여러 likelihood-based generative model이 그렇듯이 p(x)를 잘 모델링하여 high-quality fake sample을 생성하는 것이 diffusion model의 objective이다. VAE와 유사하게 latent model이지만, 시간에 따른 여러 trajectory path가 있고 latent space의 크기가 data space가 같도록 유지되는 것이 차이점이다.
Forward Process
어떤 data x(0)에 T step 동안 noise를 섞는 과정이다. DDPM 논문에서는 Markov chain function q를 Gaussian form으로 가정했으나 DPM 논문에서는 Bernoulli form으로 가정하기도 했다 (그러나 아직 이를 활용한 연구를 보지는 못했다).
q(x(t)|x(t-1)) function을 위와 같이 설정한 이유는 1) time step이 진행됨에 따라 변하는 x의 scale을 보존하기 위함이고, 2) time step이 충분히 클 때 x를 isotropic Gaussian distribution으로 수렴하게 하기 위해서다. 또, 3) nice property라고 불리는 성질을 이용하면 x(0)에서 x(t)를 바로 추론할 수 있다.
DDPM은 Gaussian distribution에서의 sampling을 위해 reparameterization trick을 주로 사용한다.
Reverse Process
위 forward process의 반대 방향으로 진행되고 learnable parameter theta가 존재한다 (forward process에는 learnable parameter가 없음).
어떤 function의 parameter theta를 학습시키기 위해서는 reverse process의 ground truth가 필요하다. 이를 q(x(t-1)|x(t))라고 하자. 문제를 쉽게 풀기 위해 그대로 사용하지는 않고 x(0)에 대한 condition을 넣는다.
Ground Truth of Reverse Process
q(x(t-1)|x(t))를 q(x(t-1)|x(t), x(0))으로 바꾸어 풀면 다음과 같은 형태의 mu와 sigma를 얻을 수 있다 (nice property와 Bayes rule를 이용하여 전개).
Loss Function
log(p(x))에 대해 variational bound를 구해보면 (식 전개 원리는 VAE와 유사) 3가지 term으로 나누어지는데,
L(T)는 constant이고 L(t-1)은 reverse process가 가능하도록 p_theta를 학습시키는 과정이다. L(0)의 경우 x(0)의 data distribution에 따라 별도의 decoder를 설계할 수 있다. L(t-1)은 두 Gaussian distribution의 KL divergence를 계산하는 것이기 때문에 variance를 고려하지 않는다면 두 mu function의 mean squared error를 구하는 것과 같아진다. p_theta의 mu function을 적절한 형태로 가정한다면 L(t-1)은 time step t에서의 noise predictor가 된다.
Training phase에서는 위에서 언급한 대로 time step t에서의 noise predictor를 학습시키는 과정이다. epsilon_theta 모델의 input은 nice property로 계산된 x(t)와 time step t이고 output은 Gaussian noise (x(t)와 data shape이 같음)가 된다. U-Net 계열의 모델이 사용되고, time step t는 Transformer에서의 position가 유사하게 sinusoidal 형태로 embedding되어 CNN feature에 더해진다.
Sampling phase에서는 Gaussian noise로 출발해 autoregressive model에서의 sampling처럼 현재 x(t)을 다음 time step으로 넘겨준다. 좋은 성능을 기대하려면 충분히 큰 T 설정이 중요할 것 같은데, 1) x(T)가 isotropic Gaussian distribution임이 보장되어야 하고 2) 각 time step의 noise prediction task가 너무 어렵지 않아야 하기 때문이다. 이는 반대로 sampling 시간을 증가시키기 때문에 {1, ..., T}에서 일부만 사용하는 방법들이 고안되고 있다.
내 생각
다른 계열의 generative model에 비해 학습에 필요한 제약 조건이 적고 안정적인 수렴이 가능할 것으로 생각된다. 같은 계열인 VAE와 비교했을 때 diffusion model이 왜 더 나은 성능을 보이는지 감은 있지만 논리적으로 설명하기는 아직 어렵다 (더 tight한 bound?). 또, information bottleneck 구조가 아니기 때문에 사용할 수 있는 latent feature가 명확하지 않은 점도 있다. 모든 독립적인 time step에 대한 shared noise predictor를 좀 더 효율적으로 학습시킬 수 있는 방법도 생각해보면 재밌을 듯.
References
[1] J. S. Dickstein et al., Deep Unsupervised Learning using Nonequilibrium Thermodynamics, ICML 2015
[2] J. Ho et al., Denoising Diffusion Probabilistic Models, NeurIPS 2020
[3] https://www.lgresearch.ai/kor/blog/view/?seq=190&page=1&pageSize=12
[4] https://lilianweng.github.io/posts/2021-07-11-diffusion-models/
[5] https://huggingface.co/blog/annotated-diffusion