Understanding Diffusion Models

Paper Review 9 : 생성모델의 수식적 이해

|Intro

이번에 리뷰할 논문은 'Understanding Diffusion Models: A Unified Perspective'입니다. 구체적으로 이미지 생성과 관련된 기본적인 사항에 대해서 수식적으로 잘 정리된 논문이어서 가져왔습니다. 이번 포스팅에서는 자체적으로 코드를 작성하는 것보다는 수식 전개를 살펴보고, 그 개념을 이해하는 것에 초점을 맞춰 진행해보고자 합니다.

또한 해당 논문의 수식과 관련된 코드를 찾아, 분석한 사항에 맞는지 확인하고 잘못된 점이 있다면 수정한 뒤 결과를 살펴보도록 하겠습니다. 본 논문은 이미지 생성 모델(Diffusion Models)의 백그라운드가 되는 여러 알고리즘 및 수식을 통합적으로 잘 정리한 메타분석 논문으로 아래와 같이 생성 모델 논문에서 기본 전제 조건으로 깔고 가는 개념에 대해 수식적으로 해설해 여러 개념들을 잘 합쳐놓은 논문이라 할 수 있습니다. 

생성모델의 핵심 개념(출처: youtube[2])

결국 생성모델은 위 그림과 같이 '정답분포 p(x) 및 그로부터 샘플링한 결과를 얻을 수 없는 상태에서, AI모델(neural net)을 이용해 실제 p분포에 근사시킬 수 있는 q분포를 이용하는 것'이며, 구체적으로 어떻게 p분포와 근사한 q분포를 생성할 것인가에 초점을 맞추고 있습니다.

무한 원숭이 정리를 표현한 그림(by chatGPT)

제가 이 부분을 읽을 때마다 드는 생각 중 하나는 바로 '무한 원숭이 정리'입니다. 무한 원숭이 정리란 키보드를 칠 수 있는 무한한 원숭이가 있다고 가정했을 때 이들이 임의로 한 글자씩 입력하면 당연히 그 출력 결과의 대부분은 제대로 된 단어나 구문조차 형성하지 못하겠지만, 그 중 일부는 셰익스피어의 햄릿과 같은 결과물을 재현내낼 수도 있다는 것입니다. 간단히 표현하자면 우리가 완전히 랜덤한 확률로 키보드의 특정 버튼을 클릭하는 프로그램을 만들고, 이를 무한한 시간동안 실행시킨다면 그 중에는 우리가 잘 아는 유명한 드라마나 영화의 대본이 나올 수도 있다는 것이죠.

하지만 이러한 가정은 수학적으로는 모르나, 실질적으로는 의미가 없습니다. '무한'이라는 개념을 현실로 가져와 사용할 수가 없기 때문이죠. 생성 모델은 이러한 '무한'의 개념을 '유한'하게 만드는 방법을 다루고 있습니다. 완벽하게 랜덤한 것이 아닌 앞뒤의 글자를 확인하거나, 이전 페이지에 등장한 인물의 이름이나 특징을 기억하거나 함으로써 말이죠. 본 논문에서는 이를 가능하게 하는 수식적 근거를 제시합니다. 하나씩 읽어보며 진행해보도록 하겠습니다.

|Background : ELBO, VAE, HVAE

생성 모델을 이해하는 첫 단계는 바로 생성 모델의 핵심 개념을 이해하는 것입니다. Diffusion 기반의 생성 모델을 포함한 모든 종류의 생성 모델은 입력을 통해 출력이 결정되며, 이를 위해 입력한 값을 낮은 차원의 Latent vector로 embedding했다가 다시 원상복구시키는 과정을 학습함으로써 적절한 Latent Vector를 입력하는 것만으로 우리가 원하는 결과물을 생성할 수 있게 됩니다. 

이러한 과정은 기존의 VAE(Variational AutoEncoder)와 동일한 방식입니다. 모델의 구조 자체는 encoder와 decoder로 구성된 AE(AutoEncoder)와 동일하지만, Encoder를 통해 압축된 벡터들의 분포까지도 고려한다는 점에서 VAE는 차별성을 가지며 생성 모델로서의 핵심 개념을 포함하고 있다고 할 수 있습니다.

물론 세상에 존재하는 모든 종류와 형상의 데이터를 확보할 수 없고, 확보한 일부의 데이터만으로 전체 데이터의 특징을 반영하는 Encoding 규칙을 구현하는 것은 무척 어려운 일입니다. 본 포스팅에서는 그러면 어떻게 관측한 데이터(학습가능한 데이터)로부터 전체 데이터의 분포를 근사할 수 있는지에 대해 알아볼 것입니다.

|ELBO

기본적으로 생성모델은 실제 분포 p(x)와 근사한 q(x)를 표현하는 잠재 구조(latent structure)를 구현하고, 해당 잠재구조를 통해 특정 입력이 들어왔을 때 의도한 형태로 재생성해내는 것을 목표로 하고 있습니다. 재생성하는 과정은 우리가 Decoder라고 명명하는 인공지능 모델에 할당한다면, 우리에게 남은 것은 근사한 분포 q(x)를 만들어내는 방법에 대한 고민이며, 그 첫 시작은 ELBO(Evidence Lower Bound)입니다.

Likelihood p(x)에 대한 공식

위 그림에서와 같이 우리는 관측된 데이터와 잠재 분포를 Joint distribution을 통해 상상할 수 있고(근사시킬 수 있고), 관찰된 데이터의 가능도 p(x)를 최대화함으로서 생성모델의 구현이 가능하도록 할 수 있습니다. 이때 관찰된 데이터로부터 joint distribution p(x, z)를 통해 가능도 p(x)를 만들어내는 방법은 두 가지가 있으며, 위 그림의 두 가지 수식에 해당합니다.

만약 p(x)를 이렇게 한 번에 구해서 최대화할 수 있다면 괜찮겠지만 현실적 한계로 인해 위 공식들 중 어느 하나를 이용해 p(x)를 구현하는 것은 한계가 있습니다. 구체적으로 Equation (1)을 사용하려면 복잡하기 그지없는 모델 내부의 내재된 변수들(Latent Variables) 간 가능한 모든 조합에 대해 적분해야 하나 현실적으로 불가능하며, Equation (2)를 사용하려면 Ground Truth Latent Encoder인 P(z|x)가 필요하지만 이를 사전에 알 수 있는 방법이 없기 때문입니다. 이로 인해 위 두 가지 공식을 이용해 관찰된 데이터(Evidence)로부터 실제 데이터를 근사해낼 수 있는 공식인 ELBO(Evidence Lower Bound)를 구현하며, 도출하는 과정은 아래와 같습니다.

ELBO 도출 과정 ①

위 그림은 Equation (1)에서부터 ELBO를 도출하고, 왜 ELBO를 최대화하는 것이 실제 p(x)와 유사해질 수 있는지를 표현하는 과정입니다. 위 그림에서 Jensen 부등식이란 다음과 같습니다.

Jensen 부등식에 대한 설명

즉, 위 도출 과정에서 log 함수가 볼록함수 f가 되어 적용된 결과로 볼 수 있습니다. 다시 정리해서 보자면 ELBO의 수식은 다음과 같습니다.

ELBO 수식

하지만 이렇게 도출한 수식은 ELBO에 대한 수식적 이해 정도면 몰라도 직관적으로 왜 ELBO를 최대화해야하고, 정말 ELBO를 최대화하면 실제와 같아지는지에 대한 납득이 어렵습니다. 다시말해 관찰된 데이터와 ELBO 사이의 관계가 명확하게 와닿지 않습니다. 때문에 이번엔 Equation (2)까지 추가로 적용해 조금 다른 방법으로 공식을 유도하면 아래와 같습니다.

ELBO 도출과정 ②

위 수식 전개 과정을 보면 알 수 있듯이 Equation (2)를 적용하고, 이후 KL Divergence의 정의를 이용함으로써 보다 직관적으로 ELBO 값이 왜 항상 본래의 값 이하로 나타나는지 알 수 있습니다. 참고로 KL Divergence는 두 확률분포 간의 차이를 의미하며 Cross Entropy와 Entropy의 결합으로 이뤄집니다. 구체적으로 KL Divergence = Cross Entropy - Entropy와 같은 식이 성립하며, 이를 통해 p분포와 q분포 간의 거리를 측정하게 됩니다[3].

Entropy와 Cross Entropy의 수식[3]

Entropy와 Cross Entropy의 수식은 위 그림과 같습니다. 표기 방식이 달라 모호해보이지만, ground truth인 p분포끼리의 확률분포, 즉 H(p, p)가 곧 Entropy인 셈입니다. 따라서 Cross Entropy와 Entropy는 다음과 같은 식이 성립하며, 두 값의 차이는 p와 q 분포의 차이에 의해서 발생하는 것으로 해석할 수 있게 됩니다. 때문에 두 확률분포가 동일하다면 0이 되고 차이가 커질수록 값이 커지게 됩니다. 즉 최소값이 0인 것이죠.

Entropy와 Cross Entropy의 관계[3]

이러한 점을 고려하면 'log P분포는 ELBO+KL Divergence 값이며, 이는 항상 ELBO 이상이다'라고 말할 수 있습니다. 또한 데이터의 likelihood가 상수라는 것을 고려한다면, ELBO를 최대화되는 경우는 KD Divergence가 최소가 되는 경우와 일맥상통하는 의미가 됩니다. 때문에 우리는 최적화 목표를 ELBO를 최대화시키는 것 혹은 KL Divergence를 최소화시키는 것이라고 정리할 수 있으며, 목적함수가 되는 것입니다.

|Variational AutoEncoder

VAE의 공식에서는 ELBO를 바로 최대화시킵니다. VAE는 자기 자신을 압축했다가 다시 재구성하는 AutoEncoder 계열의 알고리즘이며, ELBO와는 아래와 같이 연결고리가 있습니다.

VAE에 적용을 위한 ELBO

이때 위 수식 전개의 마지막에서 reconstruction term과 prior matching term은 각각 decoder와 encoder를 최적화하는 것을 의미합니다. 구체적으로 reconstruction term은 encoder를 통해 압축된 latent vector z가 decoder를 통해 우리가 학습한 데이터 z로 재구성하는 것을 의미합니다. prior matching term은 우리가 encoder로 압축한 (learned) variational distribution(qΦ(z|x))가 latent variable distribution(p(z))와 얼마나 유사한지를 측정하게 되며, 이를 통해 모델이 잠재 변수에 대한 불확실성과 다양성을 유지하도록 유도해 Dirac delta function으로 붕괴되지 않도록 합니다.

참고로 위 수식에서 θ와 ϕ는 각각 Decoder와 Encoder의 파라미터를 의미하며, pθ(x∣z)는 잠재변수 z로부터 원본 데이터 x를 재구성하는 조건부 확률분포를, qϕ(z∣x)는 주어진 입력 x에 대해 가능한 잠재 변수 z의 분포를 나타내며, p(x)는 잠재변수에 대한 사전분포로 우리가 잠재공간에 대해 사전에 가진 믿음을 의미한다 할 수 있습니다. 

또한 우리가 정의한 목적함수는 몬테 카를로 추정에 의해 근사될 수 있습니다. 몬테 카를로 추정은 잠재변수 z의 샘플들을 qϕ(z∣x)로부터 추출해 기대값을 추정하는 방법입니다. 이렇게 근사된 식은 다음과 같습니다.

몬테 카를로 추정에 의해 근사된 목적함수

하지만 이때 몬테카를로 추정으로 샘플링한 z(l)은 확률적으로 생성되기 때문에 미분 불가능하며, 이는 머신러닝의 목적함수로서의 기능을 하지 못하는 문제를 발생시킵니다. 이러한 문제를 해결하기 위해 VAE에서는 미분이 가능한 특정 분포(ex. 다변량 가우시안)대로 모델링하도록 설계를 진행하고 재매개변수화 기법(Reparameterazation Trick)을 적용합니다.

구체적으로 VAE에 ELBO를 적용함에 있어 qϕ(z∣x)는 대각 공분산을 가진 다변량 가우시안을 모델링하는 것으로 선택되며, 이는 입력 x에 대해  평균 μϕ(x)와 분산 σϕ^2(x)를 가진 다변량 가우시안 분포에 단위행렬 I가 함께 고려된 형태입니다. 또한 사전 확률 p(x)는 표준 다변량 가우시안 분포로 선택됩니다. 이를 수식으로 나타내면 다음과 같습니다.

VAE 인코더 변수

이를 통해 미분불가능한 z를 직접 Sampling하는 것이 아닌, 입력 x와 보조잡음변수(auxiliary noise variable) ϵ를 이용한 결정론적 함수(deterministic function)를 통해 간접적으로 Sampling하게 됩니다. 구체적으로 노이즈 ϵ를 기준 분포로부터 샘플링하고, 노이즈 ϵ를 변환해 z 샘플을 생성함으로써 미분가능하도록 유도합니다. 이를 수식으로 나타내면 아래와 같습니다.

z 샘플링 수식

위 수식의 x는 reparameterization trick을 이용해 연산된 값으로 수식은 아래와 같습니다.

reparameterization trick 수식

정리하면 우리의 목적은 데이터를 encoder로 압축했다가 decoder 복구시키는 AutoEncoder 방법론을 적용하며, 이때 인코딩된 z의 분포까지도 고려하는 Variational AutoEncoder를 살펴보았습니다. 다만 이 경우 인코딩된 z와 z의 원본 x의 분포 p(x|z)는 이미 알고 있어 문제가 없지만, 세상에 존재하는 모든 입력가능한 x와 그 x의 인코딩된 z의 분포 q(z|x)를 직접적으로 알 수 없기 때문에 ELBO를 이용해 근사할 수 있는 식을 도출했고, 머신러닝의 목적함수로 사용할 수 있게(미분가능하도록) Gaussian 분포를 모델링하도록 설계하고 Reparameterization Trick을 적용하였습니다.

이를 간단히 개념도로 정리하면 아래와 같습니다.

이렇게 구현한 VAE를 여러 단계에 걸쳐 반복하고 이러한 반복을 최소화해 단번에 계산할 수 있는 방법 등이 적용되면 본 논문의 최종 목적인 Diffusion에 대한 구현이 완료됩니다. 하지만 그 전에 여기까지 진행한 VAE에 대한 코드를 통해 실제 수식이 코드로 구현되는 과정을 살펴보도록 하겠습니다. 해당 코드는 직접 구현하지는 않았고, 사전에 잘 구현된 코드를 가져와 사용하였습니다. 출처는 여기[4]입니다.

|VAE Code

해당 코드는 MNIST를 압축했다가 복구하는 VAE를 모델링해 학습하는 코드입니다. 우선 VAE를 이루는 구성요소를 살펴보면 다음과 같습니다.

VAE Code[4]

위 코드를 살펴보면 Encoder를 이용해 바로 z로 압축해 출력하는게 아닌 압축한 값의 평균(mean)과 로그분산(log var)을 출력하고, 이를 임의의 노이즈 ϵ을 이용해 reparameterization trick에 적용함으로써 간접적으로 z를 샘플링한 뒤, 이렇게 샘플링한 z를 Decoder에 통과시켜 원본으로 재구성하는 과정입니다.

이러한 연산을 위해 Encoder는 원본 데이터를 압축하고, 압축한 데이터의 평균과 분산을 출력해야 하며, Decode는 Reparameteriation Trick으로 샘플링된 z를 전달받아 원본으로 재구성해야 합니다. 이에 대한 코드는 아래와 같습니다.

encoder and decoder of VAE[4]

이렇게 재생성된 x_hat은 우리가 목적함수로 지정한 ELBOW 수식을 이용한 손실함수를 구현함으로써 모델의 학습이 가능해지도록 합니다. 이에 대한 수식을 다시 살펴보면 아래와 같았습니다.

VAE에 적용되는 ELBO 수식

위 수식에서 볼 수 있듯 우리의 목적함수(손실함수) ELBO는 reconstruction term과 prior matching term으로 구성돼 있으며, 각각 디코더의 출력값과 인코더의 출력값을 이용해 계산할 수 있습니다. 구체적으로 위 값은 디코더 출력과 원본 x의 확률분포에 대한 Cross Entropy에서 인코더의 출력인 z의 평균과 로그분산을 이용한 KL Divergence값을 뺀 것입니다. 이를 코드로 구현하면 아래와 같습니다.

VAE의 손실함수 정의

이렇게 구현된 모델을 간단하게 30epoch 동안 학습시킨 뒤 입력 x와 재구성된 x_hat을 출력해 비교해보면 아래와 같습니다.

input x와 reconstruct x

|Hierarchical Variational Autoencoder

HVAE(Hierarchical Variational Autoencoder)는 여러 개의 VAE를 결합시킨 형태입니다. 정확히는 Markov Chain이 적용된, MHVAE(Markovian HVAE)입니다. 이는 여러 개의 Encoder와 Decoder로 구성돼있고, 각각의 Encoder와 Decoder는 이전 단계에 존재하는 모든 Latent Vector를 고려하는 것이 아닌 바로 직전의 Latent Vector만을 고려하는 식으로 진행되는 것입니다. 이를 그림으로 표현하면 아래와 같습니다.

MHVAE 구조[1]

이러한 구조는 마치 여러 개의 VAE가 쌓여있는듯한 구조이며, 때문에 재귀적(Recursive) HAVE라고 부르기도 합니다. 이러한 MHVAE의 결합분포(Joint Distribution)와 사후분포(posterior)는 수식적으로 다음과 같이 전개될 수 있습니다.

MHVAE의 결합분포와 사후분포

이 수식은 우리가 이전에 살펴보았던 ELBO 수식 전개 과정에 적용될 수 있는데요, 우선 전개된 ELBO 수식은 다음과 같았습니다.

Simple ELBO extend

이는 이전에 ELBO를 전개하며 젠센 부등식을 이용해 볼록함수인 log를 고려해 전개된 수식인데요, 위 수식의 결합분포와 사후분포를 MHVAE 버전으로 치환하면 다음과 같이 작성할 수 있습니다.

MHVAE식으로 전개된 ELBO 수식

이렇게 적용된 수식은 이후에 세 가지 제한(조건)이 추가됨으로써 Variational Diffusion Model로 구현될 수 있게 됩니다.

|Variational Diffusion Models

우리는 이전가지 생성 모델의 근간이 되는 ELBO와 VAE, 그리고 MHVAE를 적용하는 과정을 통해 수집가능한 데이터로부터 생성 모델을 구현할 수 있는 가능성을 확인할 수 있었습니다. 그렇다면 다음 단계는 바로 이렇게 구현한 수식을 기반으로 추가적인 구체화 및 목적함수화를 진행하는 것입니다.

|3 Key Restrictions(Assumptions)

Variational Diffusion Model(VDM)은 우리가 위에서 살펴본 MHVAE에 세 가지 제약(restrictions)가 추가됨으로써 구현되며 이는 각각 다음과 같습니다.

---

1. 

The Latent Dimension is exactly equal to the data dimension.

잠재차원은 원본 데이터 차원과 동일하다(=Encoder를 통해 차원 압축하지 않는다).

2. 

The structure of the latent encoder at each timestep is not learned; it is pre-defined as a linear Gaussian model. 

각 타임스텝의 인코더는 학습되는 것이 아닌 사전에 정의된 선형 가우시안 모델을 사용한다.

In other words, it is a Gaussian distribution centered around the output of the preivous timestep.

이는 다시 말해, 이전 타임스텝의 출력을 중심으로 한 가우시안 분포가 생성하는 것이다.

3. 

The Gaussian parameters of the latent encoders vary over time in such a way that the distribution of the latent at final timestep T is a standard Gaussian.

잠재 인코더의 가우시안 파라미터는 최종 타임스텝 T에서의 잠재분포가 표준 가우시안이 되도록 시간에 따라 변화한다(=어떤 입력 x여도 최종 타임스텝 T에서의 zT는 항상 가우시안 분포가 되도록 인코더가 구현된다).

---

이러한 제한 사항에 따라 MHVAE까지 전개된 수식이 다시 업데이트되게 되며, 하나씩 살펴보도록 하겠습니다.

|1st Restriction : 데이터 압축 X

첫 번째 가정에 의해 입력 데이터 x의 차원과 잠재 벡터 z의 차원은 동일하게 됩니다. 이를 통해 이전에는 timestep 1에서는 x로, 2부터는 z1, z2 ... zT와 같이 표현했던 사후확률(Posterior)을 일괄적으로 x0, x1, ... xT와 같이 표현할 수 있으며 수식적으로는 다음과 같이 정리됩니다.

VDM의 사후확률

|2nd Restriction : Gaussian Linear Encoder - not learned

두 번째 가정은 Encoder를 학습시키는 것이 아닌 Gaussian Linear Model로 정한다는 것입니다. 이를 통해 사전에 평균과 분산을 지정하거나 파라미터로 지정해 학습시킬 수 있습니다. 구체적으로 가우시안 인코더를 평균 µt(xt) = √αtxt−1와 분산 Σt(xt) = (1 − αt)I로 파라미터화하며, 각 단계의 공분산이 비슷한 수준(scale)이 되도록 공분산 계수 αt를 적용하고, 이를 통해 각 인코딩 과정에서 분산이 유지(variance preserving)됩니다. 이 과정에서 αt가 고정된 값이 아닌 계층적 깊이 t에 따라 변할 수 있는(학습가능한) 파라미터 계수로서, 모델의 유연성을 증가시킬 수 있습니다. VDM의 인코딩 과정을 수식으로 나타내면 다음과 같습니다.

encoding process

이와 같이 VDM의 Encoder는 각 타임스텝별로 평균과 분산이 정의되는 Gaussian으로 모델링되기 때문에 더이상 Φ로 파라미터화되지 않습니다(위 첫 번째 제약조건 수식에서 Φ가 사라진 이유입니다). 이로 인해 VDM에서는 Decoder의 파라미터인 θ에 초점을 맞추게 됩니다. 

|3rd Restriction : p(xT) = Standard Gaussian

마지막 가정은 최종 타임스텝 T에서의 x의 분포는 표준 가우시안 분포가 된다는 것입니다. 이를 통해 우리는 p(xT)가 평균이 0이고, 분산이 I인 정규분포라는 것을 알 수 있고, 첫 번째 제약조건을 통해 모든 잠재변수들을 x를 기준으로 표현할 수 있었습니다. 때문에 이를 정리해 HVAE에서 정의한 결합확률(Joint Distribution) 수식을 추가적으로 다음과 같이 전개할 수 있습니다.

VDM의 결합확률과 p(xT)의 분포

|VDM Optimization

VDM Structure

MHVAE에 추가된 3가지 제약조건을 고려해 p(x)를 근사할 수 있도록 식을 전개해, ELBO 극대화가 곧 p(x)에 대한 근사로 이어지도록  VDM 최적화가 완료되면, VDM의 샘플링 과정은 p(xT)로부터 가우시안 노이즈를 샘플링하고 T단계에 걸쳐 잡음 제거 과정 pθ(xt-1|xt)을 반복적으로 실행해 novel x0를 생성하는 것만큼 간단해집니다. 이 과정을 수식적으로 하나씩 살펴보겠습니다.

|ELBO for VDM

VAE까지 진행된 ELBO 식 전개

첫 번째로 이전에 VAE 과정까지 살펴본 ELBO의 수식입니다. 정확히는 log p(x)를 전개하고, 마지막에 젠슨 부등식을 이용해 볼록함수 log를 안으로 집어넣은 후의 수식입니다. 여기서부터 위의 3가지 제약조건에 따라 전개된 수식을 적용해 다음과 같이 전개할 수 있습니다.

제약조건이 적용된 수식

이 수식에서 연속적 곱셈을 의미하는 파이(Π)의 첫 번째 값들(pθ(x0|x1)과 q(x1|x0))을 하나씩 빼면 다음과 같이 전개할 수 있습니다.

첫 번째 값을 앞으로 뺀 전개

참고로 이때 q(xt|xt-1)는 상태 간의 전이 확률(Transition Probability)을 나타내며, pθ(xt-1|xt)는 파라미터 θ에 의존해 주어진 후속 상태(xt)에 대한 이전 상태(xt-1)의 우도를 나타내는 우도(Likelihood)를 의미합니다. 특히 q(xt|xt-1)는 x0에서 시작돼 이전 이미지에 Gaussian Noise가 더해져가는 과정을 의미하는데요, 이는 최초 입력인 x0가 xt를 생성, 예측하는데 중요한 영향을 끼친다는 것을 의미합니다. 때문에 이를 조건에 포함시켜 표시하면 다음과 같습니다.

x0를 조건에 포함

엄밀히 말하면 우리가 설계한 모델은 마르코프 성질을 따르기 때문에 x0가 굳이 필요하지 않고, 그 이전의 상태만으로 충분하며, 추가적으로 사전에 인코더를 학습가능한 것이 아닌 가우시안 선형 모델로 설계했기 때문에 q(xt|xt-1, x0)는 q(xt|xt-1)와 동일합니다. 이러한 전개는 다다음의 베이즈 규칙을 적용해 수식을 간단히 하기 위함입니다. 베이즈 규칙을 적용하기 전에 시각적인 직관성을 높이기 위해 수식을 로그의 성질과 파이(Π)의 성질을 이용해 분리하면 다음과 같습니다.

수식 분리

이렇게 분리한 수식에 베이즈 규칙을 적용할 수 있습니다. 베이즈 규칙(Bayes' Rule)은 조건부 확률을 계산하는데 사용되는 규칙으로, 추가 정보가 주어졌을 때 어떤 가설의 확률을 업데이트하는 방법에 대한 것입니다. 과거 통계를 공부할 때, 2차 세계 대전에서 추락한 전투기 조종사의 위치를 특정하기 위해서 사용됐던 통계기법이기도 하며 이에 대한 기본적인 수식은 다음과 같습니다.

베이즈 규칙

이를 우리가 전개한 수식에 적용하면 아래와 같습니다. 

베이즈 규칙

위 수식에서 q(xt|xt-1, x0)는 이전 상태(xt-1)과 초기 상태(x0)가 주어졌을 때 현재 상태(xt)의 조건부 확률이고, q(xt-1|xt, x0)는 현재 상태(xt)와 초기 상태(x0)가 주어졌을 때 이전 상태(xt-1)의 조건부 확률이며, q(xt|x0)와 q(xt-1|x0)는 각각 초기 상태(x0)가 주어졌을 때 현재 상태(xt)와 이전 상태(xt-1)의 조건부 확률을 의미합니다. 

이전에 설명했다시피 위 전개를 위해 x0가 조건으로 추가된 결과입니다. 또한 위 규칙이 우리의 수식에 적용되면 다음과 같아지게 됩니다.

베이즈 규칙이 적용된 수식

위의 두 번째 수식 중 붉은 테두리로 감싼 부분만을 따로 빼서 다음과 같이 전개 및 상게할 수 있습니다.

이렇게 전개한 식을 따로 빼서 전체 식을 다시 작성하면 다음과 같습니다.

위 그림에서 볼 수 있듯 총 세 개의 수식으로 나뉘고 첫 번째 수식과 두 번째 수식의 공통된 부분을 다시 제거할 수 있습니다. 공통된 부분을 제거한 뒤 첫 번째 수식의 p(xt)를 두 번째 수식으로 이동하고, 기대값의 성질을 이용해 세 수식을 각각 분리하면 다음과 같이 나타낼 수 있습니다.

기대값의 성질을 이용한 수식 분리

이후 위 수식에 마르코프 성질을 고려한 변형이 적용되게 됩니다. 첫 번째 수식은 x0가 주어졌을 때, x1에 대한 로그우도(Log Likelihood)의 기대값(Expectation)을 나타내며, 이 과정에서 전체 상태 시퀀스 x1:T를 고려하지만 마르코프 성질을 고려하면 x1에만 초점을 맞추게 됩니다(직전, 직후의 값에만 의존하도록 조건부 독립이 가정되었기 때문에).

두 번째 수식의 경우 마지막 상태 xT와 그 이전 이전 상태 xT-1 사이의 전이확률(q(xT|x0))과 사전확률(p(xT))의 로그비율에 대한 기대값(Expectation)을 계산하며, 첫 번째 수식과 마찬가지로 마르코프 성질을 고려해 전체 시퀀스가 아닌 xT와 xT-1에 대해 초점을 맞추게 됩니다.

마지막 세 번째 수식의 경우, 각 시점 t에서의 xt에 대해 xt+1에 조건부인 로그우도와 xt-1에 조건부인 로그 전이확률의 비율에 대한 기대값을 합산(Σ)합니다. 이 또한 마찬가지로 전체 시퀀스를 모두 고려하는 것이 아닌, xt를 결정짓는(의존적인) xt-1과 xt+1에 대해서만 초점을 맞추게 됩니다.

정리하면 마르코프 성질을 적용함으로써 복잡한 확률 모델을 간소화하는 것으로, 특정 조건 하에서의 기대값 계산에 초점을 맞추는 것입니다. 이러한 내용을 적용하면 위 수식은 다음과 같이 전개될 수 있습니다.

마르코프 성질을 적용한 간소화 수식

위의 두 번째 수식과 세 번째 수식에 역수를 취한 뒤(+가 -가 됩니다), KL Divergence를 이용해 표현하면 다음과 같습니다.

KL Divergence를 적용한 수식 전개

위 수식에서 첫 번째 term은 reconstruction term으로 해석되며, VAE에서와 동일하게 근사될 수 있고, 몬테카를로 추정에 의해 최적화될 수 있습니다. 두 번째 term은 x0에서 최종적으로 생성된 xT가 얼마나 가우시안 분포와 일치하는지를 측정하는 것으로, 우리가 논의하고 있는 VDM에서는 이를 학습하는 것이 아닌 Gaussian Linear Model을 적용함으로써 항상 xT는 표준 가우시안 분포를 따르도록 설계되었기 때문에 항상 0이 되게 됩니다.

마지막 term은 denoising matching term으로 해석되며, 디코더를 통해 xt-1이 xt로 디코딩될 때 실제 xt와 얼마나 다른지를 측정하게 됩니다. 결과적으로 완벽하게 Denoising을 했다면 이 또한 0이 되게 되며, 이는 곧 reconstruction term을 극대화함으로써 p(x)에 근사하는 결과를 얻을 수 있게 되는 것입니다.

|Sampling

ELBO를 통해 목적함수를 정의할 수 있게 되었다면, 그 다음 단계는 바로 어떻게 z를 샘플링할 것인가에 대한 논의로 넘어가게 됩니다. VAE에서 논의했던 사항은 직접적으로 확률적 샘플링을 적용할 경우, 이는 미분불가능하기 때문에 재매개변수화 기법(Reparameterization Trick)을 적용할 수 있도록 다변량 가우시안 분포로 모델링하도록 설계했었습니다. 이를 통해 간접적으로 인코딩된 x의 평균과 분포, 그리고 파라미터 노이즈를 이용해 z를 샘플링함으로써 미분가능하도록 하였습니다. 당시 정의한 수식은 아래와 같았습니다.

VAE Sampling 수식

우리는 VAE 이후, 해당 VAE를 여러 단계에 걸친 Hierarchical VAE로 확장했으며, 여기에 Markov 성질을 추가한 Markov HVAE를 가정한 후, 세 가지 제약조건을 통해 지금의 Variational Diffusion Model을 구현하였습니다.

이에 따라 샘플링 방법에도 변화가 발생하는데요, VDM의 각 과정은 데이터를 압축하지 않으며 공분산이 비슷한 수준으로 유지되도록 Scale하며 이에 따라 인코더의 전이확률은 다음과 같은 분포를 따르게 됩니다.

encoder transition

이는 인코딩 과정에서 xt의 평균은 root(αt)xt-1이며, 분산은 (1 − αt)I임을 의미합니다. 이를 zt, 정확히는 xt는 다음과 같이 샘플링된다고 수식적으로 표현할 수 있습니다.

VDM Sampling

이때 위 수식에서 xt-1을 전개하면 다음과 같습니다.

xt-1 전개

이때 노이즈(ϵ)는 가우시안 정규분포를 따르는 임의의 변수이기 때문에 두 변수의 합에 대해서 평균은 두 변수 평균의 합과 같고, 두 변수의 합에 대한 분산 또한 각 변수의 분산의 합과 같습니다. 이를 수식적으로 나타내면 아래와 같습니다.

Sum of normally distributed random variables[5]

이러한 성질을 우리가 전개한 수식의 노이즈에 적용하면 다음과 같습니다.

noise 분포 정리

이에 따라 전개된 수식은 다음과 같습니다.

t-2 시점의 샘플로 설명된 xt

위 과정을 통해 xt를 t-1 시점의 샘플링으로 나타냈고, 이를 다시 t-2 시점의 샘플링으로 나타낼 수도 있었습니다. 이러한 과정을 재귀적으로 반복하면 다음과 같이 전개, 정리될 수 있습니다. 이는 xt가 우리가 가정한 제약조건 등으로 인해 x0로부터 샘플링된 수식임을 표현하는 것이기도 합니다.

xt 샘플링 수식

|실제 전이확률(q(xt-1|xt))의 형태 추정

이러한 샘플링 수식을 통해 우리는 q(xt|x0)의 가우시안 형태를 유도할 수 있고, 이를 응용해 q(xt-1|x0)를 묘사하는 가우시안 파라미터화를 유도할 수도 있스빈다. 이처럼 q(xt|x0)와 q(xt-1|x0)를 알고 있다면 우리가 학습시킬 xt-1을 xt로 전이시키는 확률분포인 q(xt-1|xt, x0)를 구할 수도 있는데요, 그 전에 베이즈 정리에 의해 q(xt-1|xt, x0)를 전개하면 다음과 같습니다.

베이즈 정리에 의한 전개

이 수식에 대해 우리가 구한 샘플링 수식을 적용해 분포를 나타내면 다음과 같이 확장될 수 있습니다.

샘플링 수식을 적용한 전개

이렇게 전개한 뒤에는 전체적인 분포의 형상이 어떤지 확률밀도함수(probability density function, PDF)를 활용해 q(xt-1|xt, x0)의 분포를 확인하고자 하는데요, 이를 위해 가우스 분포의 확률밀도함수를 알아야 합니다. 정규분포(가우스 분포)에서의 확률밀도함수는 다음과 같이 전개됩니다.

PDF of Standard normal distribution[6]

이를 이용해 우리는 q(xt-1|xt, x0) 분포가 어떤 형상을 취하는지 확인할 수 있는데요, 확률분포의 확률밀도함수를 표현할 때 정규화 상수나 분모를 명시하지 않음으로써 전체 분포의 형태만을 강조할 수 있습니다. 이러한 수식 전개를 위한 예시를 간단히 작성하면 다음과 같습니다.

확률밀도함수를 이용한 형태 추정

이를 수식에 적용하면 다음과 같이 확장할 수 있습니다.

확률분포 추정을 위한 수식 전개

또한 위 수식의 마지막의 C(xt, x0)는 xt와 x0로만 구성된 상수항을 의미하며, 이를 제거하고 다시 다음과 같이 근사될 수 있습니다.

q(xt-1|xt, x0)의 분포

이는 q(xt-1|xt, x0)가 정규분포를 따르며, 평균(μq(xt,x0))은 xt와 x0의 함수이고, 분산(Σq(t))은 α 계수들의 함수로 나타낼 수 있음을 시사합니다. 이때 α계수는 각 timestep 별로 고정되어 사용되며, 모델의 시간에 따른 변화를 제어하게 됨으로써 모델의 불확실성이나 변동성을 나타내는 부분으로 해석될 수 있습니다. 

참고로 위 식의 마지막 부분에서 근사하는 과정에서는 이전에 근사되며 제거된 xt, x0, α로만 구성된 C가 내재적으로 반환(implicitly returned)되며, (x- μ)의 제곱식의 형상을 띔으로써 xt와 x0의 평균(μq(xt,x0))을 구현합니다. 이에 대한 식을 표현하면 아래와 같습니다.

q(xt-1|xt, x0)의 분포 근사 유도

위 근사 과정을 통해서 우리가 원하는 것은 우리가 전개/구현한 근사 잡음제거 전이확률(pθ(xt-1|xt))이 실제 잡음제거 전이확률(q(xt-1|xt, x0))을 가능한 일치시키도록 하는 것입니다. 이를 위해 두 전이 단계 모두 가우시안 분포를 사용해 모델링함으로써 통계적 속성을 명확히 하였습니다

또한 모든 α값이 timestep별로 하이퍼 파라미터와 같은 형태로 사전에 정해졌기 때문에, 이를 통해 정의되는 분산이 전이 단계의 과정에서 안정화되고 일관성을 유지하는데 도움을 주게 됩니다. 다만 근사 잡음제거 전이확률(pθ(xt-1|xt))은 x0에 기반하지 않기 때문에, 평균(μθ(xt,t))를 xt의 함수로 파라미터화해야 한다고 합니다.

이러한 개념을 그림으로 살펴보면 아래와 같습니다.

근사적인 노이즈 제거 전이 단계(pθ(xt-1|xt))를 실제 노이즈 제거 전이 단계(q(xt-1|xt))에 근사시키는 그림

즉, 두 개의 분포( pθ(xt-1|xt) 와 q(xt-1|xt, x0) )는 가우시안 모델링으로 설계되며, 이 둘을 근사시킴으로써 우리의 목적을 달성할 수 있습니다. 이때 두 개의 분포를 근사시키는 과정은 KL Divergence Term으로 설명될 수 있으며, 두 분포의 분산은 아래와 같이 x와 무관하게 timestep t에 따라 구분되며 사전에 정의되거나 계산될 수 있습니다.

denoising step의 분산

이처럼 두 분포의 분산이 동일하기 때문에 KL Divergence를 최적화한다는 것은 timestep별 두 분포의 평균이 동일할수록 두 분포가 일치함을 의미하게 됩니다. 이를 고려해 KL Divergence Term을 전개하면 다음과 같습니다.

두 분포의 분산이 동일한 경우의 KL Divergence Term

위 수식의 전개를 통해 최종적으로 전개된 결과는 두 분포의 평균에 대한 유클리드 노름(2-노름)이 최소화되는 지점입니다. 구체적으로 두 분포의 분산이 동일할 경우, 두 분포의 차이를 설명하는 KL Divergence를 최소화하는 것은 결국 두 분포의 평균 차이를 최소화하는 지점(argmin)이라는 의미입니다. 

이때 μθ와 μq는 다음과 이전에 전개한 바와 같이 다음과 같이 정의될 수 있습니다.

위 수식에서 xˆθ(xt, t)는 노이즈 이미지 xt와 time index t로부터 x0를 예측하려는 신경망에 의해 파라미터화되는 값입니다. 이를 다시 위 KL Divergence Term을 전개하는 식에 대입하면 아래와 같습니다.

KL Divergence Term 전개

위 수식을 최적화하기 위해선 신경망을 통해 최적화된(예측된) xˆθ(xt, t)가 원본 이미지(x0)와 일치화하도록 하는 것입니다. 이는 VDM을 최적화하는 것이 임의의 노이즈가 추가된 xt에서 원본 이미지 x0를 예측하기 위한 신경망을 학습시키는 것으로 귀결됨을 의미합니다. 

이를 이전에 전개한 ELBO 수식에 대입하면 위 그림의 denoising matching term에 해당하게 되며, 이는 곧 각 타임스탭(2~T)에 대한 기대치를 최소화함으로써 각 기대값의 합을 최소화하는 식으로 근사될 수 있을 것입니다. 

|Learning Diffusion Noise Parameters

이전까지의 과정을 통해 우리는 각 타임스탭별로 예측한 분포와 실제 분포와의 평균 차이를 최소화하는 것이 VDM 최적화로 이어지는 것을 확인할 수 있었습니다. 하지만 모든 타임스텝별로 이러한 과정을 반복하는 것은 특정 xt를 구하기 위해 이전 과정을 누적적으로 반복해야 하기 때문에 무척이나 비효율적입니다. 때문에 특정 시점의 xt를 단번에 구하기 위해 수식을 추가적으로 전개해 조금 더 간소화해보도록 하겠습니다. 

수식 전개 전에 간소화를 위한 SNR(신호 대 잡음 비율, Signal-to-Noise Ratio)를 알아야 하는데요, 일반적으로 SNR의 수식은 μ^2 / σ^2으로 정의됩니다. 이때 높은 SNR은 더 많은 신호가 있음을 의미하고, 낮은 SNR은 더 많은 노이즈가 있음을 의미합니다. 이를 가우시안으로 모델링되는(가우시안 분포로 정의되는) 우리의 수식에 전개할 개할 경우 다음과 같습니다.

가우시안 분포에서의 SNR 정의

이를 적용하기 위해 우리가 이전 단계에서 정의했던 각 타임스탭별 denoising matching term을 SNR을 기준으로 표현할 수 있도록 다음과 같이 전개할 수 있습니다.

denoising matching term에 대한 추가적인 식 전개

이를 통해 우리는 각 타임스텝별 Denoising Matching Term이 결국 (SNRt-1 - SNRt)[t에서의 평균차이]를 2로 나눈 값이 되는 것을 확인할 수 있습니다. 이를 수식으로 정리하면 다음과 같습니다.

SNR을 이용한 수식전개

이는 확산모델에서는 타임스텝 t가 증가함에 따라 SNR이 단조롭게(monotonically) 증가해야 함을 의미합니다. 구체적으로 교란된/노이즈 낀 입력(perturbed input) xt가 시간이 지남에 따라 노이즈가 많아지다가 t=T에서 표준 가우시안과 동일해진다는 개념을 공식화한 것입니다. 이러한 단조성을 고려한 파라미터화를 적용하면 각 타임스탭 t에서의 SNR을 다음과 같이 정의될 수 있습니다.

타임스탭 t에서의 SNR 수식

이를 이용해 각 타임스텝 t에 대한 α값으로 정의된 가우시안 분포의 평균과 분산을 sigmoid 함수를 통해 표현하면 다음과 같습니다.

타임스텝별 파라미터로 정의된 평균과 분산의 수식

이렇게 정의된 항들은 다양한 연산에 필수적입니다. 예를 들어 최적화의 과정에서, 재매개변수화 기법(reparameterization trick)을 적용함으로써 xt를 생성하기 위해 그 이전의 누적적인 연산과정을 중첩한 결과가 아닌, 입력 x0로부터 임의의 노이즈가 생성된 xt를 직접적으로 생성하는 과정을 위해 사용됩니다.

|Three Equivalent Interpretations

우리는 직전까지의 단계를 통해 xt에서 x0를 예측하도록 모델링함으로써 VDM을 최적화할 수 있음을 수식적으로 확인하였습니다. 이를 더 확장해 추가적인 해석을 통한 다양한 시각에서의 VDM 최적화 방안을 살펴보도록 하겠습니다. 

|2nd Interpretation : learn to predict source noise ϵ0

첫 번째 확장, 즉 VDM 최적화 수식 전개의 두 번째 해석은 바로 원본 노이즈에 대한 것입니다. 결론부터 말하면 xt에서 x0를 예측하도록 학습시키는 것과 마찬가지로 정규 분포를 따르는 원본 노이즈 ϵ0를 예측함으로써 VDM을 최적화시킬 수 있으며, 실험적으로는 이편이 더 성능이 좋았다는 것입니다. 이에 대해서 하나씩 살펴보도록 하겠습니다.

우리는 이전에 q(xt|x0)의 확률분포를 전개하며, xt의 분포를 x0와 ϵ0에 대한 수식으로 정리한 바 있습니다. 이를 다시 x0에 대해서 전개해보면 다음과 같습니다.

x0를 기준으로 한 전개

위 x0를 대입해 잡음제거 전이확률분포의 평균 μq(xt, x0)를 전개해보겠습니다.

x0 대입 후 전개

이렇게 전개된 사항들에 대해 수식 분리, 공통된 값 곱하기 등을 적용해 추가적으로 수식을 전개하면 다음과 같습니다.

μq(xt, x0) 전개 결과

위 수식에 따라 근사된 전이확률분포의 평균 μθ(xt, t)은 다음과 같습니다.

근사된 전이확률분포의 평균 μθ(xt, t)

우리는 이 두 확률분포의 평균에 대한 유클리드 노름을 최소화시킴으로써 VDM을 최적화시킬 수 있음을 알고 있습니다. 따라서 이를 다시 해당 식에 대입해 최적화 문제를 풀면 다음과 같아지게 됩니다.

VDM 최적화 식 확장

위 식의 마지막에서 강조된 ϵ^θ(xt,t)는 x0에서 xt를 결정하는 원본 노이즈(source noise) ϵ0∼N(ϵ;0,I)를 예측하도록 학습하는 신경망을 의미합니다. 즉, 이는 원본 노이즈 ϵ0를 예측하도록 모델링하고, 학습됨으로써 VDM을 최적화할 수 있다는 것을 의미합니다.

|3rd Interpretation : learn to predict score of the image(▽logp(xt))

마지막 세 번째 해석은 Tweedie's Formula를 이용해 각 확률분포를 Score 함수를 적용해 표현합니다. Tweedie's Formula를 적용함으로써 지수 가족 분포의 평균을 추정할 때, 샘플의 최대 가능도 추정치(경험적 평균)와 추정치의 스코어를 포함하는 수정 항을 더함으로써 샘플 편향을 완화할 수 있다고 합니다. 예를 들어, 관찰된 샘플이 모두 기본 분포의 한쪽 끝에 위치한다면(지나치게 편향되었다면) 음의 스코어가 커져 샘플의 naive maximum likelihood 추정치를 진짜 평균쪽으로 수정하게 됩니다.

tweedie's formula

가우시안 변수 z ∼ N (z; µz, Σz)에 대한 tweedie's formula 공식은 위와 같습니다. 이를 왼쪽 그림과 같은 분포를 가진 샘플 xt의 진정한 사후확률분포 q(xt|x0)를 예측하기 위해 대입해 전개하면 아래와 같습니다.

Tweedie's Formula가 적용된 전이확률분포 수식

Tweedie's Formula에 의하면 xt가 생성되는 진짜 평균을 위한 최적 지점은 µxt = √ α¯tx0 일 때로 정의되며, 이를 위 수식에 대입해 x0를 기준으로 아래와 같이 전개할 수 있습니다.

Tweedie's Formula에 의한 최적 x0

이렇게 전개된 x0는 이전 확장 해석에서 전개된 노이즈 기준 전개와도 유사한데요, 이를 이용해 수식을 전개함으로써 위 수식의 스코어 함수 ∇xtlogp(xt)에 대해 살펴보도록 하겠습니다. 이때 스코어 함수 ∇xtlogp(xt)는 xt의 변화에 따른 로그확률분포(logp(xt))의 변화를 의미하며, 이전의 수식을 대입한 전개는 아래와 같습니다.

노이즈를 통해 설명된 스코어 함수 ∇xtlogp(xt)

위 수식을 통해 확인할 수 있는 것은 tweedie's formula를 통해 정의된 스코어 함수 ∇xtlogp(xt)를 모델링하는 것이 결국 원본 노이즈의 반대를 모델링하는 것과 동일하다는 것입니다(그 외에는 αt로 이뤄진 상수항의 차이로 인한 스케일링 인자 차이만 있음). 이는 데이터 공간에서 로그 확률을 최대화하는 방향으로 움직이는 스코어 함수는 노이즈 생성의 반대 방향으로 움직이게 되며, 간단히 말해 '잡음 제거'가 이뤄지는 것으로 설명할 수 있습니다.

이렇게 x0와 스코어 함수에 대한 수식이 정리됐다면 2번째 확장 해석에서 했던 방식과 동일하게 각 잡음제거의 전이확률 평균을 확장해 전개하면 됩니다. 아래와 같습니다.

실제 전이확률에 대한 확장 전개

그리고 이를 다시 근사 잡음제거 전이확률 평균에 적용하면 다음과 같아지게 됩니다.

근사된 잡음제거 전이확률 평균

위 수식에서 sθ(xt,t)는 임의의 노이즈 레벨 t에 대해 데이터 공간에서 xt의 스코어 함수 ∇xtlogp(xt)를 예측하도록 학습하는 신경망입니다. 이때 스코어 함수 ∇xtlogp(xt)는 변수 xt에 대한 logp(xt)의 변화율을 의미하는 것으로 , xt가 주어진 데이터 공간 내에서 어떻게 이동할지를 나타내는 그래디언트를 의미하게 됩니다.

이를 최적화 문제에 대입해 적용하면 아래와 같습니다.

tweedie's formula를 이용한 최적화 식 전개

우리는 위에서 스코어 함수는 노이즈 제거의 과정으로 움직이도록 하는 함수임을 확인했으며, 이를 음수로 적용하는 것은 결국 노이즈를 생성하는 것으로 해석할 수 있습니다. 즉 위 수식 또한 이전에 전개한 수식들과 마찬가지로 타임스텝 t에서 단조롭게 생성된 임의의 노이즈를 예측하도록 모델링함으로써 VDM을 최적화할 수 있음을 의미합니다.

지금까지의 과정을 통해 우리는 VDM을 최적화하기 위한 세 가지 목표를 도출했으며, 이는 각각 ①원본 이미지 x0를 예측하는 신경망, ②원본 노이즈 ϵ0를 예측하는 신경망, 그리고 ③임의의 잡음 수준 ∇xtlogp(xt)에서의 이미지의 스코어를 예측하는 신경망을 학습함으로써 달성될 수 있음을 확인했습니다. 

이때 마지막 확장 전개에 대해서는 여전히 모호하며 직관적으로 이해가 어려운 부분이 존재하는데요, 이에 대해서 VDM이 아닌 다른 모델을 이용해 아래에서 진행토록 해보겠습니다.

|Score-based Generative Models

VDM(Variational Diffusion Model)을 최적화하는 세 가지 방법 중 마지막 에너지 함수에 기반한 설명은 다소 모호할 수 있습니다. 구체적으로 스코어 함수를 예측함으로써 VDM이 최적화될 수 있음은 확인했지만, 이게 다른 방법들과의 차이는 무엇이고 그 자체적으로 어떤 의미가 있는지 알 수는 없었습니다. 때문에 이를 보다 직관적으로 설명할 수 있는 다른 모델을 통해서 추가적으로 설명이 가능한데요, 그것은 바로 Score-Based Generative Model(여러 종류가 있습니다)입니다.

Score-based Generative Model에서 사용되는 임의의 유연한 확률분포(arbitrarily flexible probability distributions)는 다음과 같이 정의될 수 있습니다.

arbitrarily flexible probability distributions

위 수식에서 Zθ는 정규화 상수로, 모든 가능한 x값들에 대한 확률값들의 총합 혹은 적분할 경우의 값이 1이 되도록 합니다. fθ(x)는 에너지 함수이고, e^(-fθ(x))는 x의 함수로서 분포의 형태를 결정짓습니다. 참고로 정규화 상수 Zθ는 분포를 정규화하기 위한 목적으로만 사용되며, 분포의 형태나 x의 특정 값에 따라 변하지 않습니다.

위 수식의 양변에 ∇xtlog를 씌워주면 우리에게 익숙한 에너지 함수의 형태가 되는데요, 이를 적용해 수식을 전개해보도록 하겠습니다.

수식 전개 결과

우리는 이전에 확인한 바와 같이 에너지 함수가 sθ(xt,t)와 근사하도록 유도됨을 확인했으며, 이로 인해 위와 같이 수식이 전개될 수 있습니다. 이러한 스코어 모델은 다음 수식과 같이 Ground Truth 스코어 함수와의 Fisher Divergence를 최소화함으로써 최적화될 수 있습니다. 

Fisher Divergence

이때 Fisher Divergence의 수식은 주어진 확률 분포 p(x)에 대한 스코어 함수와 모델이 추정한 스코어 함수 sθ(x) 사이의 차이를 측정하기 위한 것으로 일반적으로 위 그림과 같습니다. 다만 이러한 개념을 유지하되, 수식을 변형해 우리에게 익숙한 형태로 작성하면 다음과 같습니다.

Fisher Divergence 변형

이러한 수식을 좀 더 깊게 이해하려면 스코어 함수가 구체적으로 무엇을 의미하는지를 이해해야 하는데요, 본 논문에서는 스코어 함수를 '모든 x에 대해, x의 로그 가능도(log likelihood)에 그래디언트를 취함으로써 본질적으로 데이터 공간에서 어떤 방향으로 움직여야 가능도(likelihood)를 더 증가시킬 수 있는지를 설명하는 것'으로 설명하고 있습니다. 쉽게 설명하자면 스코어 함수는 데이터 x가 존재하는 전체 공간에 대한 벡터 필드를 정의하고, 데이터 분포에서 가장 높은 확률 밀도를 가지는 지점들을 향하게 되는 것입니다. 이러한 과정을 시각적으로 표현하면 다음과 같습니다.

Using Langevin dynamics to sample from a mixture of two Gaussians[7]

위 그림은 두 개의 가우시안 분포들로부터 샘플링한 x들이 가장 높은 확률밀도를 가지는 지점들(modes)로 향하는 과정을 시각화한 것입니다. 위 그림과 같이 참 데이터 분포의 스코어 함수를 학습함으로써, 임의의 공간, 임의의 지점에서 시작해 스코어가 모드에 도달할 때까지 반복함으로써 우리가 원하는 샘플들을 생성하는 것이 가능해집니다. 이러한 과정을 Langevin Dynamics라 하며, 수식적으로는 다음과 같습니다.

Langevin Dynamics Equation

위 수식의 시작점이 되는 x0는 균일 분포와 같은 사전 분포에서 무작위로 샘플링되며, 추가적인 노이즈 항의 역할을 하는 ϵ ~ N(ϵ; 0, I)은 생성된 샘플들이 하나의 모드로 축소되지 않고 다양성을 부여하는 역할입니다. 구체적으로 학습된 스코어 함수는 결정론적인 함수이기 때문에, 노이즈 항을 포함해 샘플링하는 것은 생성 과정에서 확률성을 추가해 결정론적인 궤적(늘 같은 궤적)을 피할 수 있도록 합니다. 이러한 방법론은 우리가 이미지와 같은 복잡한 분포를 모델링하는 과정에서 실제 스코어 함수를 알 수 없음에도, 스코어 매칭을 통해 이러한 제한을 극복하고 확률적 경사하강법을 통해 최적화될 수 있음을 제시합니다.

Visualization of three random sampling trajectories generated with Langevin dynamics

노이즈를 사용하는 이유를 조금 더 구체적으로 설명하자면, 노이즈가 포함되지 않은 순수한 스코어 매칭(vanila score matching)은 크게 세 가지 한계를 지니고 있습니다. 첫 번째는 데이터가 저차원 다양체(manifold)에 국한되며 스코어 함수가 잘 정의되지 않는 것이고, 두 번째는 스코어 매칭이 저밀도 지역에서 정확하지 않아 드물게 나타나는 데이터 포인트에 대해 학습하기 어려우며, 마지막으로 실제 스코어 함수를 사용하더라도 랑주뱅 동역학 샘플링이 데이터 공간 전체에서 잘 섞이지 않을 수 있다는 것입니다.

하지만 여기에 가우시안 노이즈가 추가됨으로써 데이터 샘플이 전체 공간에 걸쳐 분포할 수 있게 되어 첫 번째 문제를 해결할 수 있고, 가우시안 노이즈를 크게 추가함으로써 각 모드가 더 넓은 영역을 커버하게 되어 저밀도 지역(드물게 나타나는 데이터 포인트)에서도 학습 신호를 증폭시켜 두 번째 문제를 해결할 수 있으며, 마지막으로 증가하는 분산을가진 여러 단계의 가우시안 노이즈를 추가함으로써 실제 데이터 분포의 혼합 게수를 반영하는 중간 분포를 생성해 세 번째 문제를 해결할 수 있게 되는 것입니다.

이처럼 분포를 스코어 함수를 통해 표현하고, 랑주벵 동역학(Langevin dynamics)과 같은 마르코프 체인 몬테 카를로 기법을 통해 스코어 함수를 사용함으로써 샘플링을 하는 것을 통틀어 점수 기반 생성 모델링(Score-Based Generative Modeling)이라고 합니다. 우리는 T시점까지의 noise levels의 양의 시퀀스를 선택하고, 이를 점진적으로 교란된 데이터 분포의 시퀀스를 아래와 같이 공식적으로 정의할 수 있습니다.

t시점의 pertubed data x의 분포 수식

동시에 신경망 sθ(x, t)는 모든 noise levels에 대한 스코어 함수를 학습하기 위해 스코어 매칭을 사용해 학습됩니다. 이를 수식으로 나타내면 아래와 같습니다.

위 수식에서 λ(t)는 노이즈 레벨 t에 조건을 부과하는 양의 가중치 함수입니다. 즉 위 수식의 목적은 타임스텝 t에서 점진적으로 교란된 데이터 분포를 예측하는 신경망을 구현하는 것이며, 이 과정에서 노이즈가 포함돼 위에서 설명한 세 가지 한계를 극복할 수 있게 되는 것입니다. 이러한 최적화 수식은 우리가 이전에 Variational Diffusion Model을 학습하기 위해 유도한 목표와 거의 정확하게 일치하며, 샘플링 과정 또한 거의 같습니다. 

샘플링 과정에 대해 구체적으로 설명하자면 저자들이 제안하는 annealed Langevin dynamics sampling은 t=T, T-1, ..., 2, 1에 대해서 순차적으로 Langevin dynamics를 이용해 생성됩니다. 이러한 샘플링의 과정에서 초기화(최초의 샘플링)는 균일 분포와 같은 사전 분포에서 선택되며, 각 후속 샘플링 단계는 이전 시뮬레이션의 최종 샘플에서 시작됩니다. 또한 Noise Level이 시간단계 t에 걸쳐 꾸준히 감소하고, 시간이 지남에 따라 step size를 줄이기 때문에 샘플들은 결국 진정한 모드로 수렴하게 됩니다(조금  어렵게 표현했지만 결국 MHVAE에서와 동일한 논리적 프로세스를 내재한다는 의미입니다).

다만 이때 생길 수 있는 의문은 어떻게 확산 모델을 무한한 수의 타임스텝으로 자연스럽게 일반화할 수 있느냐에 관한 것입니다. 좀 더 쉽게 표현하자면 확산 모델의 개념을 연속적인 차원으로 확장하는 방법이 없을까에 대한 의문입니다. 기존의 확산 모델은 시간을 이산적인 단계로 나눠 각 단계마다 이미지에 점진적인 노이즈를 추가하고, 이를 역전시켜 원본 이미지를 재구성하지만, 무한한 타임스텝으로의 확장은 이 과정을 '연속적인 확률 과정'으로 모델링하게 함으로써 모든 타임스텝에 대해 노이즈를 추가하고 제거하는 과정을 보다 정교하게 표현할 수 있도록 합니다. 즉, 어떻게 하면 이산적인 확률과정을 연속적인 확률과정으로 바꿔서 성능 향상을 유도할 수 있을까에 대한 의문인 셈입니다.

Solving a reverse SDE yields a score-based generative model[7]

이에 대한 해결책이 바로 Stochastic Differential Equation(SDE)이며, 이를 사용함으로써 연속적인 변화 과정을 수학적으로 정의할 수 있고, 다양한 노이즈 레벨에서 이미지가 어떻게 변화할지를 설명합니다. 따라서 이 과정을 역으로 진행하는 것은 곧 무작위 노이즈로부터 점차적으로 원본 데이터의 구조를 복원하는 과정이며, 이를 위해서는 각 noise level에 대한 스코어 함수를 정확히 추정해야 합니다. 

추가적으로 SDE의 다양한 매개변수화를 통해 다양한 유형의 데이터와 노이즈 패턴에 적합한 맞춤형 확산 과정을 설계할 수 있으며, 이는 노이즈 추가 및 제거의 구체적인 방식을 조정할 수 있게 하여, 특정 데이터 유형의 복잡성과 특성을 더 잘 반영하는 생성 모델을 개발할 수 있게 합니다. 요약하면, 이러한 접근 방식은 생성 모델의 범용성과 유연성을 크게 향상시킬 수 시킬 수 있는 가능성을 제시하며, 이러한 일련의 과정을 시각화하면 아래와 같습니다.

SDE를 통한 노이즈 생성(left)과 제거(right) 과정[7]

개인적으로 이 부분을 직관적으로 이해하는 과정이 조금 험난했는데요(본 논문에서 스코어 베이스 모델에 대해선 간단히 짚고 넘어가기도 합니다), 만약 스코어 베이스의 분산 모델링에 대해 좀 더 자세한 설명이나 코드 등이 궁금하다면 제가 참조한 포스팅[7]을 읽어보시는 걸 추천드립니다. 실제 본 논문에서 참조한 스코어 함수 및 SDE와 관련된 논문을 작성하기도 하신 Yang Song님의 블로그 포스팅입니다.

|Guidance

드디어 본 논문의 마지막 챕터입니다. 여태까지 우리가 진행한 사항을 간단히 요약해보면 우리는 이미지에 노이즈를 적용한 뒤, 이를 재구성하는 과정이 수식적으로 가능한 것인가를 증명했습니다. 그 결과로써 (가우시안 분포 등의 가정이 필요하긴 했지만) 우리는 크게 세 가지 방법을 통해 원본 이미지를 재생성하는 것이 가능하다는 것을 확인할 수 있었습니다. 세 가지 방법은 결국 우리가 신경망을 통해 무엇을 예측하도록 학습하는지를 정의하는 것으로, 첫 번째는 원본 이미지 x0를, 두 번째는 원본 노이즈 ϵ0를 예측하는 것이고, 마지막은 바로 스코어 함수 ∇xtlogp(xt)를 예측하는 것이었습니다. 이러한 방법들 중 하나를 통해 우리는 원하는 이미지를 재생성할 수 있게 되며, 현실적인 한계로 한정된 데이터로도 복잡한 이미지를 표현할 수 있는 가능성을 얻게 되었습니다. 

그리고 마지막 챕터인 Guidance에서는 이러한 목적함수에 조건 정보 y를 추가적으로 고려하도록 함으로써 우리가 원하는 형태로 유도할 수 있는 방법을 추가하고자 합니다. 예를 들어 가장 마지막에 정의한 스코어 함수를 이용한 수식 표현에 조건 정보 y를 고려해 표현하면 아래와 같습니다.

조건 정보 y가 고려된 스코어 함수

이처럼 생성 과정에 컨디션 정보 y를 고려토록 하는 것을 Guidance라고 하며, 위 수식을 조금 더 가다듬은 Classifier Guidance와 Classifier Free Guidance가 일반적으로 사용되고 있습니다.

|Classifier Guidance

대표적인 Guidance의 방법 첫 번째인 Classifier Guidance의 수식은 아래와 같습니다. 조건부 확률의 수식을 베이즈 규칙에 의해 전개하고, xt에 대해 그레디언트를 하는 과정에서 0이 되는 ▽logp(y)를 제거하면 최종 수식으로 전개됩니다.

Classifier Guidance 수식

위 수식에서 첫 번째 unconditional score가 일반적으로 우리가 guidance를 고려하지 않은 스코어로 Diffusion Model로 학습되며, 그 뒤에 붙는 adversarial gradient가 조건 정보 y로 인해 추가된 것으로 Classifier를 통해 학습됩니다. 위 수식을 목적함수로 사용해 최적화할 경우, 항상 조건 정보 y를 고려하는 식으로 학습되기 때문에 좀 더 세밀한 제어를 위해 하이퍼 파라미터 γ를 추가해 다음과 같이 수정합니다.

수정된 Classifier Guidance 수식

위 수식을 통해 γ가 0일 때는 기존과 동일하게 guidance가 존재하지 않을 때의 스코어를 학습하게 되고, γ가 커질수록 조건정보를 더 고려한 샘플 분포를 학습하게 됩니다. 하지만 이렇게 설계된 Classifier Guidance 모델은 임의의 노이즈가 존재하는 입력을 처리할 수 있어야 하기 때문에 반드시 Diffusion Model과 함께 학습되어야 하는 제한이 존재하는데요, 이를 해결하기 위해 다음의 Guidance가 등장하게 됩니다.

|Classifier Free Guidance

Classifier Free Guidance(CFG)는 별도의 Classifier를 학습시키지 않고 이를 하나의 모델에 녹여내는 방식으로 전개합니다. 이를 위해 하이퍼 파라미터 γ를 함께 고려하도록 설계되며 그 수식은 아래와 같습니다.

CFG 수식

위 수식은 하이퍼 파라미터 γ가 추가되기 전 수식을 대입하고, 이를 전개해 모든 항을 하이퍼 파라미터 γ에 대한 수식으로 풀어낸 것입니다. 이를 통해 별도의 분류기 모델없이 단일 모델만으로 조건 정보를 고려한 이미시 생성이 가능해집니다. 위 수식에서 하이퍼파라미터 γ가 0일 때는 이전의 Classifer Guidance와 동일하게 guidance가 존재하지 않을 때의 분포만 학습하게 되고, 1일 때는 조건 정보에만 집중해서 학습하게 되지만, 추가적으로 γ가 1보다 클 경우에는 조건부 스코어에 가중치를 늘릴뿐 아니라 무조건부 스코어에서 멀어지는 방향으로 학습하게 됩니다. 이는 샘플의 다양성을 감소시킬 수 있지만, 대신 조건 정보와 일치하도록 샘플을 생성하는 효과가 있습니다.

동일한 하이퍼파라미터를 통해 단일 모델로 구현가능한 위 수식을 최적화한 모델(CFG)는 모델 학습에 들어가는 비용이 적을뿐 아니라, 단일 확산 모델을 학습시키는 것에서 큰 변형을 가하지 않으면서도 조건부 생성이라는 목적을 달성할 수 있기 때문에 '우아하다'고 합니다.

여태까지의 포스팅을 통해 Diffusion Model의 이론적 기반이 되는 수식과 용어에 대해 살펴봤습니다. 제가 진행하는 프로젝트 중에서 이미지 생성 모델을 이용하는 것이 많아, 좀 더 상세한 리뷰를 위해 이처럼 정리해봤는데요. 본 논문에서 제시한 수식이나 개념들은 향후 다른 Diffusion 기반 생성 모델 논문에서 기본으로 알고 있다 가정하고 넘어가는 경우가 많기에 가장 먼저 리뷰해야하는 논문이라고 생각해 포스팅 해보았습니다. 

구체적으로 이후에 리뷰할 논문들은 이러한 이론적 베이스를 기반으로 비교적 쉽게 이해할 수 있는 DDPM 혹은 DDIM에 대해서 리뷰해볼 것인데요, 해당 논문들은 일반적으로 Sampler로 분류되는 계열의 알고리즘이며 본 논문 포스팅에서 제시된 신호대비잡음 비율(Signal to Noise Ratio, SNR)부터 분산 모델이 최적화될 수 있는 다양한 방법에 이르기까지의 수식과 가정을 알고 있어야 쉽게 이해할 수 있습니다.

|Reference

[1] Calvin Luo. Understanding Diffusion Models: A Unified Perspective. 

https://arxiv.org/abs/2208.11970

[2] ImcommIT. 생성이랑 확률이 뭔 상관? 

https://www.youtube.com/watch?v=ENMtsWy52WA&list=LL&index=3

[3] 서울대학교 컴퓨터공학부 Club Gauss. 정보이론과 머신러닝-엔트로피.

https://cse.snu.ac.kr/sites/default/files/node--notice/ 머신러닝-엔트로피.pdf

[4] Jackson-Kang. Pytorch-VAE-tutorial. 

https://github.com/Jackson-Kang/Pytorch-VAE-tutorial/blob/master/01_Variational_AutoEncoder.ipynb

[5] Wikepedia. Sum of normally distributed random variables. 

https://en.wikipedia.org/wiki/Sum_of_normally_distributed_random_variables      

[6] Wikepedia. Probability density function. 

https://en.wikipedia.org/wiki/Probability_density_function  

[7] Yang Song. Generative Modeling by Estimating Gradients of the Data Distribution. 

https://yang-song.net/blog/2021/score/