Denoising Diffusion Probabilistic Models (DDPM) 논문 리뷰

DDPM 개요

출처: https://arxiv.org/pdf/2006.11239

 

Denoising Diffusion Probabilistic Models는 2020년 발표된 생성모델로 최근 이미지 생성분야에서 각광받고 있는 모델이다. DDPM을 직관적인 한줄로 이해하면 다음과 같다.

이미지에 점진적으로 노이즈를 추가해서 정규분포로 만든 뒤, reverser process로 다시 노이즈를 제거하면서 이미지를 복원하는 확률적 생성모델

 

물론 이렇게 써놓으면 다 이해한 것 같지만 실상은 그렇지 않다. 자세하지는 않지만 나름대로의 직관으로 논문을 이해해보도록 하겠다.

1. Introduction

• 최근 몇 년 사이에 GAN,VAE,AutoRegressive Model 등 다양한 딥러닝 생성 모델이 이미지, 오디오 등 여러 데이터에서 진짜처럼 보이는 샘플을 잘 만들어내고 있다.
• GAN이 대표적이었지만 다른 방식도 점점 좋아져서 거의 GAN 수준의 이미지를 생성하는 것이 가능해졌다.
• 이 논문에서는 Diffusion Model이라는 새로운 생성 방식을 다룬다. 이 Diffusion Model은 "노이즈를 더하는 과정"을 수식으로 정의하고 "노이즈를 빼는 과정"을 신경망이 학습해서 자동으로 수행하도록 만든다. 노이즈가 조금씩 추가되는 마르코프 연쇄 구조라서 매 단계가 전부 조건부 Gaussian만으로 간단히 표현된다. 이 덕분에 신경망을 간단히 훈련 시킬 수 있다.
• 또한 해당 논문에서는 Diffusion 모델의 수식과 샘플링 구조를 잘 고르면, Score Matching과 Langevin dynamics라는 기법들도 수학적으로 거의 같은 구조가 된다는 것을 증명했다. 이 방식으로 훈련하면 샘플의 퀄리티가 매우 향상된다.
• 그러나 모델이 실제로는 눈으로 볼 수 없는 미세한 부분까지 다 압축하려해서 효율이 떨어진다. 때문에 이미지 샘플의 퀄리티와 별개로 로그 가능도는 기존 likelihood 기반 모델보다 낮다.

2. Forward & Reverse Process

Reverse Process (p)

Diffusion model은 latent variable models(잠재 변수 모델)의 일종으로 다음과 같은 형태를 가진다:

reverse process

 

여기서 $x_1...,x_T$는 데이터 $x_0~q(x_0)$와 동일 한 차원을 갖는 잠재 변수들이다. $p_\theta(x_0:T)$는 Reverse Process라고 하며 다음과 같이 정의되는 Markov Chain이다. 이 Chain은 학습된 가우시안 transition을 따르며 $p(x_T) = N(x_T;0,I)$($xT$​는 평균 0, 분산 1의 정규분포에서 샘플링됨)에서 시작한다.

위 수식은 $xT$부터 $x_0$까지 차례로 노이즈를 제거하며 복원하는 마르코프 연쇄 과정의 정의를 보여준다. 각 단계 $t$에서 노이즈 이미지 $x_t$를 받아 $x_t-1$(덜 노이즈된 이미지)를 샘플링한다. 이때 평균 $\mu_\theta$와 분산 $\sigma_\theta$는 딥러닝 모델이 예측한다. 한마디로 Reverse Process는 정규분포처럼 보이는 노이즈에서 점점 선명한 이미지로 되돌리는 과정이며, 그 매 단계는 신경망이 학습한 가우시안 분포를 따른다고 말할 수 있다.

 

Forward Process (Q)

Forward Process는 Diffusion 확률 모델에서 데이터에 점진적으로 노이즈를 추가하는 마르코프 체인이다. Forward Process는 기호 Q로 표기한다. Forward Process $q(x_{1:T}|x_0)$는 데이터 $x_0$에 대한 일련의 잠재 변수 $x_1, \ldots, x_T$로 구성된다. 이 과정은 고정된(학습되지 않는) 마르코프 체인으로 정의된다.

• $x_t$는 시간 스텝 $t$에서의 데이터(노이즈가 추가된 상태)이며 $x_{t-1}$는 이전 시간 스텝 $t-1$에서의 데이터이다.
• $\beta_t$는 시간 스텝 $t$에 따라 미리 정해진 작은 양의 분산 스케줄(variance schedule)인데 논문의 실험에서는$\beta_1=10^{-4}$부터 $\beta_T=0.02$까지 선형적으로 증가하는 스케줄을 사용했으며, $T=1000$으로 설정했다.
• $\mathbf{I}$는 단위 행렬을 의미하며 $\sqrt{1 - \beta_t}$항은 노이즈 추가 시 분산이 커지지 않도록 데이터 스케일을 다운하는 역할을 한다.

이러한 정의 덕분에 포워드 프로세스는 임의의 시간스텝 $t$에서 $x_0$로부터 직접 $x_t$를 샘플링하는 닫힌 형태(closed-form)를 가진다. 이 닫힌 형태 덕분에 학습 중의 임의 스텝 $t$에서 $x_t$를 효율적으로 샘플링할 수 있다. 

3. DDPM의 Loss 함수

위에서 forward, reverse 프로세스를 이해하고 오면 이 Loss 함수를 이해하는 것이 조금 수월해진다. DDPM에서 제시하는 로스 수식은 데이터의 negative log-likelihood에 대한 variational bound를 최적화 하는 것에 기반한다.

이 기본 Variational bound 수식의 각 항에 대한 설명을 해보자면

• 수식은 데이터 $x_0$에 대한 잠재 변수 $x_1, \ldots, x_T$의 전체 경로 $x_{0:T}$를 고려한다.
• $q(x_{1:T}|x_0)$ 앞서 설명한 forward process이고 $p_\theta(x_{0:T})$는 reverse process이며 학습 가능한 매개변수 $\theta$를 가진다.
• $p(x_T)$는 최종 노이즈 상태 $x_T$의 사전 분포이며 표준 정규 분포 $\mathcal{N}(x_T; 0, \mathbf{I})$로 설정된다.
• $q(x_t|x_{t-1})$은 포워드 프로세스의 스텝별 전환 확률이고 $p_\theta(x_{t-1}|x_t)$는 리버스 프로세스의 스텝별 전환 확률로 학습 가능한 가우시안 분포로 모델링된다. 

위의 기본 Variational bound는 직접 계산하기 어렵기 때문에 논문에서는 Variational bound를 재구성하여 분산을 줄인 형태를 사용한다.

이러한 KL divergence 항들로 분해 된다고 한다.

이 수식에 대한 각 항을 설명하자면

• $L_T$는 최종 스텝 $T$에서의 포워드 프로세스 분포 $q(x_T|x_0)$와 사전 분포 $p(x_T)$ 간의 KL Divergence이다. 논문에서는 포워드 프로세스 분산 $\beta_t$를 고정했기 때문에, $L_T$는 학습 중 상수가 되어 무시할 수 있었다.
• $L_{t-1}$는 시간 스텝 $t$에서 $t-1$로 가는 리버스 프로세스 전환 확률 $p_\theta(x_{t-1}|x_t)$이 포워드 프로세스의 사후 확률 $q(x_{t-1}|x_t, x_0)$를 얼마나 잘 모방하는지를 측정한다.
• $L_0$는 첫 번째 스텝 $x_1$에서 원본 데이터 $x_0$로 가는 리버스 프로세스 전환 확률의 네거티브 로그 우도이다. 

DDPM 논문에서는 $L_{t-1}$ 항에서 모델 매개변수 $\theta$를 노이즈 예측 모델 $\epsilon_\theta$로 재매개변수화 했다. 이는 평균 $\mu_\theta(x_t,t)$를 직접 예측하는 대신 데이터 $x_t$에 추가된 노이즈 $\epsilon$ 예측하도록 매개변수화하는 것이다. 결국에 노이즈가 없는 이미지를 생성하는 것은 노이즈가 있는 이미지에서 노이즈를 예측하고 예측된 노이즈를 제거하는 것과 같기 때문에 이 방식을 사용한다. 

• $E_{x_0,\epsilon}[\cdot]$은 학습 데이터셋에서 샘플링된 원본 데이터 $x_0$와 표준 정규 분포 $N(0, I)$에서 샘플링된 노이즈 $\epsilon$에 대해 계산된다. 시간 스텝 $t$도 1부터 $T$ 사이에서 균일하게 샘플링하여 기댓값에 포함한다.
• $\frac{\beta_t^2}{2\sigma_t^2\alpha_t(1-\bar{\alpha}_t)}$는 시간 스텝 $t$에 의존하는 가중치 항이다. 이는 원본 Variational bound $L$의 각 시간 스텝별 항에 대한 가중치 역할을 한다.
• $\beta_t$는 Forward Process의 분산 스케줄이며 $\sigma_t^2$는 Reverse Process의 고정된 분산이다.
• $\Vert \cdot \Vert^2$는 $\epsilon$과 $\epsilon_\theta(\cdot, t)$ 사이의 유클리드 거리 제곱, 즉 평균 제곱 오차(MSE)를 나타낸다
• $\epsilon$은 시간 스텝 $t$에서 원본 데이터 $x_0$에 실제로 추가된 노이즈이다. 이 값이 모델 $\epsilon_\theta$가 예측해야할 target 노이즈이다.
• $\sqrt{\bar{\alpha_t}}x_0 + \sqrt{1-\bar{\alpha_t}\epsilon,t}$은 시간 스텝 $t$에서 노이즈가 추가된 데이터 $x_t$를 나타낸다. 이 $x_t$가 신경망 모델 $\epsilon\theta$의 입력으로 사용된다.
• $\epsilon_\theta(\cdot,t)$는 학습가능한 신경망 모델 \theta의 출력이다.

정리하면, 시간 스텝 $t$와 해당 시간 스텝의 노이즈 데이터 $x_t$가 주어졌을 때, 신경망 모델 $\epsilon_\theta$가 $\x_t$에 실제로 추가되었던 노이즈 $\epsilon$을 학습하는 것을 목표로 하며, 이 오차에 시간 의존적인 가중치를 곱하여 최소화한다. 이 형태는 denoising score matching과 유사하다고 한다. 논문에서는 Variational bound를 직접 최적화 하는 대신 샘플 품질을 향상시키고 구현을 단순화하기위해 단순화된 목적 함수 $L_{simple}$을 사용했다.

$L_{simple}$은 이 수식 (12)에서 시간 단계 $t$에 의존하는 가중치 항 $\frac{\beta_t^2}{2\sigma_t^2\alpha_t(1-\bar{\alpha}_t)}$을 제거한 형태이다. 즉, 변이 하한의 각 시간 단계 항에 대한 가중치를 제거하고 모든 시간 단계의 노이즈 예측 오차에 동일한 중요도를 부여한 것이다.

위의 내용들을 종합하자면 DDPM은 Variational Inference를 사용하여 Negative log likelihood에 대한 Variational bound를 최적화한다. 이는 모델의 매개변수 $\theta$를 조정하여 학습된 Reverse Process $p_\theta$가 실제 데이터 분포 $p(x_0)$과 유사하도록 만드는 것을 목표로 한다. 결론적으로, DDPM의 핵심 학습 목표는 노이즈 제거 스코어 매칭과 유사한 형태의 단순화된 목적 함수 $L_{simple}$을 통해, 주어진 노이즈 레벨 $t$에서 원본 데이터에 추가된 노이즈 $\epsilon$를 예측하는 함수 $\epsilon_\theta$를 학습하는 것이라고 할 수 있다.

 

여기까지가 DDPM의 Forward, Reverse Process, Loss 함수에 관한 내용이다. 이해하는데에 에너지를 쏟아서 구현은 추후에 해보도록하겠다.