이 글에서는 2020 Neurips에 게재된 "Denoising Diffusion Probabilistic Models" (Shaofeng Zhang et al) 논문을 정리한다.
요즘들어 매우 자주 사용되고 언급되는 Diffusion model의 시대를 연 논문이라고 할 수 있다.
이 논문의 주요 Contribution은 Diffusion model의 Loss를 수정하여 더 정확한 노이즈 예측을 수행하고 이를 통해 고해상도 이미지 복원을 가능하게 하였다는 점이다.
이를 이해하기 위해 우선 Diffusion model부터 알아보자
Background
Diffusion model은 $p_{\theta}(x_0) := \int p_{\theta}(x_{0:T})dx_{1:T}$ 를 따르는 잠재 변수 모델이다. 여기서 $x_1, ... , x_T$는 각 time step별로 노이즈가 추가된 이미지이며 모두 차원이 같다.
$p_\theta(x_{0:T})$는 reverse process라고 부르며 이는 노이즈에서 점점 노이즈를 걷어내고 이미지를 생성하는 과정을 뜻한다(위의 그림에서 오른쪽 방향 화살표). reverse process는 다음과 같은 Marcov chain으로 정의된다.
※ Marcov chain에 관해서는 아래 글에서 참고할 수 있다. 베이지안 네트워크와 관련된 포스트인데, Marcov chain은 베이지안 네트워크에서 바로 이전 time step에만 영향을 받는 특수한 형태로 정의된다. 즉, DAG(방향성 그래프)의 형태가 위의 그림과 같은 케이스를 의미한다.
이번엔 반대로 이미지에서 노이즈를 조금씩 추가하여 완벽한 노이즈를 만드는 forward process를 살펴보자. 이 과정이 왜필요하냐면, forward process에서 얻은 값들을 토대로 reverse process를 학습시키기 때문이다.
Forward process는 $q(x_{1:T}|x_0)$으로 나타내며 마찬가지로 Marcov chain으로 정의된다. 이 과정엔 variance schedule $\beta_1, ... , \beta_T$에 따른 Gaussian noise를 추가하는 과정이 포함된다.
위의 식과 같이 이전 time step 이미지 분포의 분산에 $\beta_t$를 곱하고 평균에 $\root{1-\beta_t}$를 곱하여 gaussian noise를 추가한다.
여기서 $\beta$를 학습 가능한 파라미터로 둘 수도 있는데, 이 논문에서는 상수값을 사용한다.
훈련을 위한 loss는 다음과 같은 variational bound로부터 도출된다.
※ variational bound와 밑에 나올 reparameterization trick은 VAE논문을 보면 이해할 수 있다. 이것도 조만간 포스트 예정.
그리고 $\alpha_t := 1 - \beta_t$이고 $\alpha_t := \Pi_{s=1}^t \alpha_s$라고 하면, 다음과 같이 정의할 수 있다.
그리고 이를 통해서 위의 variational bound 식을 다음과 같이 정리할 수 있다. (증명이 논문 appendix에 나오지만 매우매우 복잡하다. 풀어보면 이렇게 나온다는 정도만 이해하자.)
여기서 3가지의 loss로 나뉘는 것을 볼 수 있는데 이 중 중요한 건 가운데의 $L_{t-1}$ term이다($L_T$와 $L_0$에는 학습가능한 parameter가 없기 때문에 미분하면 사라진다). 이 term의 각 값을 구할 수 있으면 학습을 진행할 수 있다.
$L_{t-1}$ term은 q와 p의 KL(Kullback-Leibler) distance로 정의된다. 그 중 q의 분포는 forward process에서 정의된 likelihood값을 기반으로 posterior를 구하면 위의 식과 같이 정의된다.(이것도 유도 증명이 논문 appendix에 나와있다)
※ Kullback-Leibler에 대해서는 아래 포스트에서 확인할 수 있다.
상대 엔트로피, 쿨백 라이블러 발산(Kullback-Leibler divergence) (tistory.com)
그럼 이제 KL의 왼쪽 분포(q)를 deterministic하게 구하는 데에 성공했으니 오른쪽 분포인 $p_\theta(x_{t-1}|x_t)$만 구하면 되는데, 이게 closed form으로 구할 수가 없다. 따라서 이 값을 근사해야 한다. 여기서 trainable parameter를 이용하여 근사하게 되는데, 이 부분을 포함하여 $L_{t-1}$을 깔끔하게 정의한 것이 이 논문의 핵심 contribution point라고 할 수 있다.
Denoising Diffusion Probabilistic Models
이 섹션에서는 각 loss들에 대해서 하나씩 살펴본다. 위에서도 얘기했지만 논문의 핵심은 $L_{T-1}$의 KL 오른쪽 분포인 p를 정의한 것이다. 그러나 loss전체의 이해를 위해서 $L_T$와 $L_0$도 살펴본다.
$L_T$
먼저 $L_T$를 살펴보자. $D_{KL}(q(x_T|x_0) || p(x_T)) 로 정의되는데 KL의 양쪽 분포 모두 학습이 필요하지 않은 정규분포로 파라미터를 설정했기 때문에 미분하면 사라지게 된다. 즉 무시해도 된다.
$L_{T-1}$
KL의 왼쪽 분포는 위의 background에서 소개했으므로 오른쪽 분포를 살펴보자. 해당 $p_\theta$를 정규분포로 나타내면 다음과 같다.
$p_\theta$를 보면 모델의 parameter인 $\theta$와 input $x_t$, $t$에 의해 결정되는 정규분포의 평균 $\mu_\theta$와 $\sum_\theta$로 정의된다. 여기서 분산에 해당되는 $\sum_\theta$는 $\sigma_t^2I$로 고정시켰다. 즉, 학습가능한 parameter없이 상수로 정의된다는 것이다. 그러면 우리는 정의해야할 수식이 더 적어졌다!
위의 수식에만 모델 parameter가 존재하게 되었다. 이제 이를 가지고 수식을 풀어보면 다음과 같이 유도할 수 있다.
분산을 상수로 고정하였기 때문에 KL을 풀면 다음과 같이 모델에 의해 forward process 의 posterior mean인 $\hat{\mu_t}$와 모델의 예측값인 $\mu_\theta(x_t,t)$의 차로 표현될 수 있다.
여기서 다음 식과 같은 reparameterization trick을 적용할 수 있다.
이를 이용하여 수식을 전개하면 다음과 같이 풀어낼 수 있다.
전개된 식을 보면 아직 모델 예측값 부분인 $\mu_\theta$부분은 전개가 되어있지 않은데 이 부분은 다음과 같이 전개할 수 있다.
이제 (10)과 (11)을 이용하여 다음과 같이 깔끔하게 정리할 수 있다.
깔끔하게 정리되었지만 $\alpha$와 $\beta$로 이루어진 수식이 구현상 복잡하게 느껴지므로 좀 더 simplify한 loss를 제안하였다.
위 식이 유도 과정의 마지막 단계이다. 즉, 이 논문은 위의 loss를 $L_{T-1}$으로 이용하여 정확하게 노이즈를 예측하고, 고품질 이미지(뿐만 아니라 대부분의 생성 task에서)를 생성할 수 있게 되었다.
$L_0$
$L_0$은 reverse process 관점에서 봤을 때 최종 이미지를 생성하는 마지막 디코더이다. 여기서는 이미지를 고해상도로 완벽하게 복원하는 것이 중요하다.
이미지의 각 픽셀은 {0, ... , 255}의 값을 가질 수 있다. 모델에서는 이 값을 [-1, 1] 사이의 값으로 스케일링 하여 사용한다. 따라서 $L_0$은 다음과 같이 쓸 수 있다.
$\mu$의 위첨자 i는 이미지의 특정 좌표를 나타낸다. 위의 식을 통해서 노이즈가 제거 된 정확한 픽셀 값을 추출하는 것을 볼 수 있다. $\delta$가 -1이하 1이상이면 무한대를 나타내므로 -1에서 1사이의 값이 출력되도록 학습된다.
Algorithm
정리한 loss를 토대로 다음과 같이 학습 알고리즘을 정리할 수 있다.
Experiments
다음은 실험 결과이다. IS와 FID에서 모두 좋은 성능을 보이는 것을 알 수 있다.
위의 그림은 실제로 DDPM을 이용해 생성한 이미지이다. 고품질의 고해상도 사진을 출력하는 것을 볼 수 있다.