DDPM 논문 리뷰(Denoising Diffusion Probabilistic Models)

Diffusion Model의 초기 논문으로 유명한 DDPM에 대해 리뷰해보고자 한다. 수학도 수학인데, 논문의 디테일이 이해가 안되서 일주일동안 이것저것 찾아보면서 공부를 했다. 

<한 눈에 보는 정리>

Introduction

생성된 샘플을 먼저 보자. Diffusion model의 초기 모델임에도 불구하고 초고해상도의 이미지 역시 잘 생성하는 모습을 볼 수 있다. 

 

이 논문(이하 DDPM)은 ' Deep Unsupervised Learning using Nonequilibrium Thermodynamics'에 바탕을 둔 논문으로 마르코프 체인을 이용해 데이터에서 노이즈를 더해가는 과정의 역과정을 variational inference를 통해 추정하는 아이디어를 그대로 차용한다. DDPM은 Loss함수를 더 간단히 만듦과 동시에 Reparametrization trick을 이용해 residual estimation을 통해 성능을 항샹시켰으며, denoising score matching과 annealed Langevin dynamics와 결론적으로 같은 목적을 가지고 있음을 보여준다. 

 

Background

다음과 같이 forward process와 reverse process를 정의한다.

여기서 눈여겨 보아야 할 점은 forward process는 이미 고정되어 학습이 필요없다는 점인데, encoder를 학습해야 했던 VAEs나 forward process를 학습해야 했던 Normalizing Flow와의 차이점을 보여준다.

variational bound는 위와 같이 주어지며, forward process에서 다음과 같이 정의할 수 있다.

where

식 L은 아래와 같이 바꿀 수 있다.

L식에서 각각의 KL divergence term은 Gaussian이므로 analytic하게 계산 가능하다.

이제$L_T$, $L_{1:T-1}$, $L_0$ 순서로 알아보자. 

 

Forward process and $L_T$ 

DDPM에서 ${\beta}_t$는 ${10^{-4}}$부터 $0.02$까지 linear하게 증가하는 constant로 설정하였다. 또한 $q$ 역시 학습해야 할 파라미터를 포함하지 않으므로 $L_T$는 학습을 하는데 있어 고려하지 않아도 된다.

 

Reverse process adn $L_{1:T-1}$

다음과 같이 $p_{\theta}(x_{t-1}|x_t)$를 정의하자. 여기서 ${\Sigma}_{\theta}(x_t, t) = {\sigma}^{2}_{t}I$로 정의해 시간에만 의존하는 상수로 정의하자. 

실험적으로 $ {\sigma}^{2}_{t} = {\beta}_t$로 설정하거나, ${\sigma}^{2}_{t} = \tilde{{\beta}_t} = \frac{1-\bar{{\alpha}_{t-1}}}{1-\bar{{\alpha}_t}}{\beta}_t$로 놓을 수 있으며 첫번째 설정은 $x_0 \sim N(0, I)$일 때, 두번째 설정은 $x_0$가 하나의 점으로 고정되었을 때 optimal하다. 

 

즉, 이제 $p_{\theta}(x_{t-1}|x_t)$는 다음과 같이 정의된다.

이제는 ${\mu}_{\theta}$에 대해 생각해보자. $L_{t-1}$은 KL Divergence term이므로 정의에 대입하면 다음과 같이 식을 변형할 수 있다.

위 식에서 알 수 있듯이 ${\mu}_{\theta}$는 forward process의 posterior mean인 $\tilde{{\mu}_t}$을 추정하는 모델이다.

Posterior 

 

베이지안 추론에서의 Posterior)

베이지안 추론에서 우리가 어떤 사건이나 관측 데이터를 통해 얻은 정보와 기존에 가지고 있던 사전 지식(사전 확률, prior probability)을 결합하여 새로운 정보를 업데이트한다. 이때, 업데이트 정보를 사후 확률(posterior probaibility)라고 한다. 

 

조건부 확률로서의 Posterior) 

posterior는 일반적으로 어떤 관측갑싱 주어졌을 때, 특정 변수의 조건부 확률 분포를 가리킬 때 사용한다. 예를 들어 측된 데이터 $y$가 주어졌을 때, 어떤 파라미터 $\theta$의 posterior는 $p(\theta|y)$로 표현된다.

 

Diffusion model에서의 posterior)

Diffusion model의 forward prcess에서 $x_t$가 주어졌을 때, $x_{t-1}$의 확률 분포를 posterior알고 한다. 

여기서 Reparametrization trick을 통해 식을 아래와 같이 수정할 수 있다.

위 식에서 ${\mu}_{\theta}$가 $\tilde{{\mu}}_{t}$를 예측해야 하기 때문에 아래와 같이 ${\epsilon}_{\theta}$를 function approximator로 선택할 수 있다.

즉 $x_{t-1}$을 $p_{\theta}(x_{t-1}|x_t)$로부터 샘플링한다는 것은 다음을 계산한다는 것과 동일하다.

이 과정은 Langevin dynamics와 닮았다.

또한 우리의 $L_{t-1}$은 다음과 같이 쓸 수 있다.

이 과정은 denoising score matching과 닮았다. 

다음은 학습과 샘플링 알고리즘이다.

Data scaling, reverse proces decoder, and $L_0$

reverse process의 마지막 항은 독립적인 discrete decoder로 위와 같이 정의한다. decoder를 다음과 같이 정의함으로써 데이터에 노이즈를 추가하거나 scaling 연산의 Jacobian을 계산할 필요가 없어진다.

 

위 사실을 종합해  계수항 없이 간단하게 training objective를 나타낼 수 있다.

 

$L_{simple}$은 작은 $t$에 해당하는 손실 항목의 가중치를 낮추어 더 어려운 노이즈 제거 작업에 집중할 수 있도록 한다. 작은 $t$는 노이즈가 매우 적은 상태를 의미하는데, 이 상태에서의 노이즈 제거 작업은 어렵지 않으나, $t$가 작을수록 계수가 커 큰 가중치를 부여한다. $L_{simple}$에서는 가중치를 없애므로써 더 큰 $t$에 대한 노이즈 제거 작업에 집중하도록 한다.

 

Experiments

T = 1000, groupd normalization + U-Net, parameters are shared across time, self-attention at 16 * 16 feature map resolution

Sample quality & Reverse process parametrization and training objective ablation

• unconditional CIFAR10 기준으로 IS는 9.46, FID는 3.17를 기록했다.
• $\mu$를 예측하는 기본 모델보다 $\epsilon$을 이용해 예측하는 모델이 훨씬 안정된 성능을 보였음을 표 2를 통해 알 수 있다.

* 이 논문에서 Negative Log-likelihood(NLL)과 Lossless Codelength는 같은 의미로 사용되고 있다.

NLL) 모델이 데이터를 얼마나 잘 맞추는지를 측정하는 척도로 $-\log{p(x)}$로 측정한다.

Lossless Codelength) 데이터를 압축하고 다시 원래대로 복원했을 때, 정보 손실이 없는 상태를 유지하는 압축 길이. 정보이론에 따르면 데이터 $x$의 최적의 Lossless Codelength는 NLL로 계산할 수 있다. 

 

Progressive coding

테이블 1을 통해 DDPM이 codelength 기준 EBM이나 annealed importance sampling보다는 좋지만 다른 모델보다는 경쟁력이 없다는 사실을 발견했디. 그럼에도 DDPM이 높은 품질을 보이는 것을 감안한다면 DDPM은 losser compressor로서의 역할을 할 수 있는 inductive bias를 가지고 있다고 할 수 있다.

 

rate은 $L_1 + ... + L_T$, distortion은 $L_0$로 정의해보자. CIFAR10에서 rate은 1.78bits/dim, distortion은 1.97bits/dim이며, 이는 0에서 255사이의 스케일에서는 0.95의 RMSE에 해당한다. 이는 모델이 생성한 샘플에서 나타나는 distortion은 사람이 인지하기 어려운 작은 차이지만, 모델이 이를 압축된 표현으로 설명하기 위해 노력한다는 의미로 볼 수 있다.

 

* inductive bias : 모델이 학습 과정에서 특정한 유형의 데이터를 잘 처리하도록 설계된 성질로 Diffusion Model에서는 loss compress에서 탁월한 성능을 보이는 inductive bias가 있다고 봄

* rate : 데이터를 표현하는데 필요한 비트수로 낮은 압축률은 데이터를 덜 압축한 상태를 의미하며, 높은 압축률은 데이터를 많이 압축한 상태를 의미한다.

* distortion : 압축 후 원본 데이터와의 차이로 distortion이 보통 데이터의 품질 저하로 나타난다.

 

progressive lossy compression

Progressive lossy compression을 도입하여 rate-distortion을 살펴보자. 알고리즘 3과 4는 minimum random coding과 같은 절차로 $x \sim q(x)$인 샘플을 한다. 여기서 $p$와 $q$는 어떤 분포라도 될 수 있으며, receiver는 사전에 $p$만 알 수 있다. $x_0 \sim q_0(x_0)$에 적용하면 알고리즘 3과 4에 의해 $x_T, ..., x_0$의 샘플이 순차적으로 전송되며, 총 기대 코드 길이는 식 (5)와 동일하다.

* Progressive lossy compression : 데이터를 점진적으로 압축하면서, 필요한 비트 수와 압축 과정에서 발생하는 왜곡(정보 손실)의 관계를 탐구하는 방법

* Minimum Random coding : 데이터 압축에서 최소한의 비트를 사용하여 정보를 전송하는 방법으로 평균적으로 KL Divergence만큼의 비트를 사용해 데이터를 전달할 수 있다.

수신자는 $t$ 시점의 $x_t$를 사용해 $x_0$를 점진적으로 추론할 수 있다.

또한 위 그림은 CIFAR-10 테스트 세트의 rate-distortion 곡선을 나타내는데, 각 시간 $t$에서 distortion은 $\sqrt{\frac{||x_0-\hat{x_0}||^2}{D}}$로 계산되며, rate은 현재까지 수신된 비트의 누적 수로 계산된다. distortion은 rate-distortion 곡선의 낮은 압축률 구간(데이터를 적게 압축)에서 급격히 감소하는데, 이는 비트의 대부분이 실제로 인지할 수 없는 왜곡을 설명하는데 할당된다는 것을 의미한다.

 

* rate-distortion curve : rate와 distortion의 관계를 나타내는 그래프로 일반적으로 데이터를 더 많이 압축하면(압축률이 낮아지만) 왜곡이 커지는 경향이 있다. 즉, rate과 distortion은 반비례 관계다.

 

progressive generation

또한 random 비트로부터 점진적으로 생성하는 과정도 진행해보았는데, 이는 알고리즘 2를 사용하여 역방향 과정에서 샘플링하면서 역방향 과정의 결과인 $\hat{x_0}$를 예측하는 과정이다.아래의 그림을보면 큰 이미지의 특징들이 먼저 나타나고 세부 사항들이 나중에 나타나는 것을 볼 수 있다. 

 

아래 그림은 stochastic prediction $x_0 \sim p_{\theta}(x_0|x_t)$ with $x_t$ frozen for various t를 보여준다. 각각의 $x_t$에 대해 생성되는 조건부 이미지를 구분한 것으로 만약 $t$가 작다면 세부사항을 제외한 모든 것이 유지되지만 $t$가 크면 큰 이미지의 특징들만이 유지된다. 

* conceptual compression : 이미지나 데이터를 복원하는 과정에서 핵심적인 구조나 큰 특징들을 먼저 복원하고, 세부 사항은 나중에 복원하는 방식의 압축

 

 

Connection to autoregressive decoding

식 (5)는 아래와 같이 쓸 수 있다.

이제 diffusion process의 길이 $T$를 데이터 차원의 수로 설정하고,  forward process에서 $q(x_t|x_0)$가 첫 $t$좌표가 마스킹된 상태로 모든 확률 질량을 $x_0$에 할당하게 한다.(즉, 첫 t개의 좌표를 마스킹하고 나머지 좌표를 $x_0$와 동일하게 유지한다.) $p(x_T)$는 모든 질량을 빈 이미지에 할당하며(Noise 상태), 편의상 $p_{\theta}(x_{t-1}|x_t)$를 완전히 표현력 있는 조건부 분포로 가정한다.(즉, 주어진 $x_t$를 통해 $x_{t-1}$를 예측한다.)

 

이로 인해 $D_{KL}(q(x_T)||p(x_T)) = 0$이 되고, $D_{KL}(q(x_{t-1}|x_t) || p_{\theta}(x_{t-1}|x_t))$를 최소화하면 $p_{\theta}$는 $t+1, ..., T$를 변경하지 않고 복사하며, 주어진 $t+1, ..., T$에 대해 $t$번째 좌표를 예측하도록 훈련된다. 따라서 이 특정한 diffusion process를 통해 $p_{\theta}$를 훈련하는 것은 autoregressive model을 훈련하는 것과 동일하다. 

 

따라서 우리는 Gaussian diffusion model을 일종의 autoregressive model with generalized bit ordering that cannot be expressed by reordering data coordinates로 볼 수 있다. 선행 연구들은 이러한 재배열(reordering)이 inductive bias를 가져온다고 밝혔는데, Gaussian diffusion 역시 유사한 목적으로 사용될 수 있다고 추측된다. Gaussian Noise는 Masking Noise보다 이미지에 더 자연스럽게 추가될 수 있기 때문이다. 또한 $T$는 꼭 데이터 차원과 같을 필요는 없고, 더 짧게 설정할 수 있다.

 

* Generalizaed Bit Ordering : 데이터의 각 좌표를 단순히 재배열하는 것 이상의 복잡한 순서를 의미한다. 즉, 데이터 생성 과정에서 특정 순서에 따라 비트를 배치하고, 이를 기반으로 모델이 학습하는 방식이다. 

 

Interpolation

DDPM에서는 latent space에서 $q$를 stochastic endoer로 이용하여 원본 이미지 $x_0, x'_0 \sim q(x_0)$를 보간할 수 있다. $x_t, x'_t \sim q(x_t|x_0)$를 사용하여 선형 보간된 벡터 $\bar{x_t}$를 reverse process로 decoding하여 변환한다.($\bar{x_0} ~ p(x_0 | \bar{x_t})$) 결론적으로 원본 이미지를 선형 보간한 손상된 버전에서 artifact를 제거하기 위해 reverse process를 사용한 과정이 아래 그림에서 묘사되고 있다. 아래 그림에서 노이즈는 고정했으며, 아래 그림의 오른쪽 그림은 원본 이미지의 보간 및 복원의 결과를 보여준다. 안경 같은 속성은 보간되지 않는다는 점에 주목하자.

 

Reference

1. Unet의 이해 https://velog.io/@lighthouse97/UNet%EC%9D%98-%EC%9D%B4%ED%95%B4

UNet의 이해

2. DDPM 논문 코드 구현하기 https://velog.io/@yeomjinseop/DDPM-%EA%B5%AC%ED%98%84%ED%95%98%EA%B8%B03

[Annotated Diffusion] DDPM-(3): 코드 세부 구현하기

3. transformer https://wikidocs.net/31379

16-01 트랜스포머(Transformer)