[논문 리뷰] Denoising Diffusion Probabilistic Models (DDPM) + 구현 코드

0️⃣ Introduction

최근 다양한 딥 생성 모델(GAN, Autoregressive models, Flow-based models, 그리고 VAE) 등이 이미지와 오디오에서 뛰어난 샘플을 생성하고 있다.

이 논문에서는 Diffusion probabilistic models 분야의 진전을 이루는데, 이는 노이즈를 점진적으로 추가한 후 역으로 이를 제거해 원본 데이터를 복원하는 마르코프 체인 기반 모델을 의미한다.

참고로 여기서 Diffusion(확산)이란 이미지에 노이즈가 추가되는 과정을 확산이라고 한다.

기존에는 Diffusion model이 고품질 샘플을 생성할 수 있다는 명확한 증거가 없었지만, 이 논문은 그 가능성을 입증하며, 다른 생성 모델보다 더 나은 결과도 가능함을 보인다.

그 밖의 특징으로는

DDPM은 고품질의 샘플을 생성하지만 다른 likelihood 기반 모델보다 경쟁력 있는 log likelihood가 없다.

DDPM의 lossless codelength가 대부분 인지할 수 없는 이미지 세부 정보를 설명하는 데 사용되었다.

Diffusion model의 샘플링이 autoregressive model의 디코딩과 유사한 점진적 디코딩이라는 것을 보였다.

1️⃣ 논문 용어/수식 설명

1-1) Diffusion process(from data to noise)

논문 내에서 는 diffusion process이며 미세한 Gaussian noise를 추가하는 과정이다.

VAE에서 사용하는 reparameterization trick으로 학습할 수 있으나, 이 논문에서는 사전에 정의한 noise를 사용하기 때문에 학습할 필요가 없다.

추가로 는 점점 noisy해지는 과정 상에서 Markov chain이다.

즉, 만일 이 condition으로 주어진다면 이전의 에 independent하다.

또한 왼쪽 식에서 의 아래 첨자는 1에서 T까지의 process 과정을 표현하는데, 는 noise가 없는 데이터이고, 가 커짐에 따라 점점 noisy해져서 는 Gaussian noise가 된다.

오른쪽 식은 에 만큼의 노이즈가 추가되는 것을 의미한다.

1-2) Reverse process(from noise to data)

논문 내에서 는 reverse process이며 미세한 Gaussian noise를 걷어내는 과정이다.

아래 첨자로 표현된 는 neural network로 학습되는 확률 모델임을 의미한다.

Diffusion process와 마찬가지로 또한 Marcov chain이다.

그렇다면, 어떻게 process를 보고 가 denoising process가 되게 학습시킬 수 있을까?

2️⃣ DDPM loss

2-1) Variational Autoencoder(VAE)

VAE loss 수식을 유도하는 과정을 복습해보면 다음과 같다.

VAE loss는 latent를 로 표현하는데, 위 그림에서의 를 로 치환하면 좀 더 익숙한 형태로 볼 수 있다.

위 그림에서 확인할 수 있듯, VAE에서 수식을 전개해나갈 때 사용하는 트릭은

를 곱해줌으로써 intractable KL divergence를 없애고 ELBO(Evidence Lower Bound)를 남기기

ELBO 수식에서 bayesian 표현으로 갈라져나온 (위 그림의 파란박스)가 latent의 KL divergence (가 를 따르므로) + 엄밀히 말하면 해당 항은 KL divergence는 아니라고 한다. 다만, 해당 항에서 KL divergence를 손쉽게 유도할 수 있다.

ELBO 수식에서 latent KL divergence를 제외한 나머지가 reconstruction term으로 남아 두 loss term으로 표현됨

2-2) DDPM loss term

VAE의 수식 전개를 조금 비틀어서, DDPM의 loss 수식 전개 또한 살펴볼 수 있다.

자세히 보면 파란 박스만 바뀌었다.

결국 맨 아래 수식을 살펴보게 되면 와 의 KL divergence를 최소화시키는 term이 된다.

1의 마지막에서 해소하고 싶었던 “로부터 어떻게 로 하여금 denoising process를 학습시킬 수 있을까?”의 질문이 수식적으로 등장했다.

두 분포의 condition과 target variable이 서로 반대다. 이를 수식으로 풀어내는 과정이 다음과 같다.

2-3) “seems intractable” to tractable

(1) Bayes rule을 바탕으로 multi-variable에 대한 conditional, joint probability로 나눈다.

(2) VAE에서 사용한 트릭과 같이 를 곱해준다.

(3) 2번에서 빨간색으로 표현된 부분을 intractable한 KL divergence Term으로 제거하고 ELBO를 남긴다.

(4) 3번과 같은 식을 표현만 달리하였다.

(5) p와 q가 markov chain임을 이용하여 Markov chain의 성질을 이용하였다.

(6) log의 성질로 곱으로 표현된 부분을 합으로 나타내었다.

(7) 6번의 수식과 같다.

(8) t=1일 때의 오른쪽 항을 분리했다.

(9) Markov chain의 성질로 나타낼 수 있다. 사실상 loss term을 만들어낼 때 제일 핵심이 되는 구간이라고 한다. 이 과정을 통해서 수식 내의 와 distribution이 같은 condition으로부터 같은 target distribution을 나타내게 되었다. 참고로 증명은 다음과 같다.

(10) log의 성질로 곱으로 표현된 부분을 합으로 나타내었다.

(11) 10번에서 log로 갈라져나온 term을 보면, log에서의 합은 곱으로 표현되므로 결국 다 상쇄되어 해당 term만 남는다.

(12) 11번에서 파란색 분모 부분을 앞으로 빼고 남은 term을 정리해서 남겨놓는다.

loss 수식 정리가 완료되었다. 총 3개의 term이 나왔는데, 각각 어떤 의미를 가질까?

2-4) Loss term 설명

가 만드는 noise 의 분포와 data 로부터 과정을 거쳐 만들어진 noise 의 분포간 KL divergence를 최소화한다.

이 loss는 process를 많이 진행할 시 분모의 가 와 independent한 Gaussian distribution이라고 가정할 수 있기 때문에 또한 그냥 Gaussian distribution에서 sampling하는 방식을 취하여 이 term에 대해서는 학습하지 않는다.

latent 상에서 로부터 의 분포를 예측하는 reverse process는 가 Gaussian process일 경우 아래의 수식으로 나타낼 수 있다.

이해는 제대로 못했는데, 평균은 노이즈를 고려한 와 의 가중합이고, 분산은 를 조정하여 시간 에서의 노이즈 수준을 계산하는 것이다.

VAE의 reconstruction loss와 같은 역할을 한다. latent 로부터 data인 를 추정하는 확률 모델의 parameter를 최적화한다.

3️⃣ Diffusion models and denoising autoencoders

3-1) Objective Function ()

길고 긴 과정을 거쳐 결국 위와 같은 식을 유도해냈다.

여기서 가장 중요한 term은 이다. 우리는 부터 시작하여 conditional하게 식을 전개하다 보면, tractable한 forward pass posterior 의 정규분포를 알 수 있는데, 이를 바탕으로 KL divergence를 계산하면 우리가 결과적으로 학습하고자 하는 를 학습시킬 수 있다.

3-2) Forward Diffuson Process ()

Objective Function의 첫번째 term으로 위에서 상술한 이유 때문에 는 prior 와 거의 유사하게 된다.

따라서 이 loss term은 항상 0에 가까운 상수이며, 학습 과정에서 무시된다.

3-3) Reverse Diffusion Process ()

Objective Function의 두번째 term으로 이를 계산하기 위해 첫 번째로 의 분포를 알아내고, 두 번째로 를 알아내기 위해 와 를 알아내야 한다.

이를 순서대로 하나씩 살펴보겠다.

이는 우리가 기존 diffusion model에서 알아봤던 것처럼 다음과 같은 형태로 정의할 수 있다.

여기서

이는 우리가 기존 diffusion model에서 알아봤던 것처럼 다음과 같은 형태로 정의할 수 있다.

다음으로, 의 표준편차는 라는 상수 행렬로 정의한다. 그러므로 이에 대해서는 학습이 필요하지 않다.

마지막으로 의 평균 는 다음과 같이 정의한다.

증명

Sampling

이렇게 위 과정을 통해 를 얻어내면, 이를 활용해서 우리는 로부터 을 샘플링 할 수 있다. 이 과정은 다음과 같은 알고리즘으로 표현할 수 있다.

위의 재료들(, , )을 잘 조합하여 (KL Divergence)를 계산하면 L_{t-1}을 아래처럼 쓸 수 있다.

아래 식을 보면 expectation의 기준이 에서 , 으로 달라지는데, forward process에 있어서 어떤 time에 대한 Gaussian Distribution이 time과 만 주어진다면 reparameterize를 통해 표현할 수 있기 때문에 이렇게 표현하는 것은 자연스럽다.