근래에( 이미 조금 되었지만 ) Generative Model로 연구가 무척 활발히 이루어지고 있는 Diffusion Model에 대해서 살펴보도록 하겠습니다. 본 아티클에서는 Diffusion에 대한 직관부터 DDPM 그리고 DDIM까지 다루고자 합니다. 그럼 시작하도록 하겠습니다. :)
1. Diffusion with Physical Intuition
1.1. Diffusion Destroys Structure
Data distribution -> Uniform Distribution
Diffusion은 우리가 가지고 있는 데이터 분포를 destroy하기 위해서 사용된다. 상단의 이미지에서 염료의 분포를 probablilty of denstiy funtion 즉 PDF로 생각해볼 수 있다. 즉, 염료의 density는 probablilty density를 나타낸다. 이때 ML의 관점에서 이야기해보자면 현재 담겨있는 물이 기존의 데이터의 분포다. 여기에 염료 즉 noise를 probabilistic하게 diffusion forward 과정을 거치게 되면 기존의 데이터 분포를 시간이 충분히 지났다는 가정하에 노이즈가 많이 더해져 uniform distribution이 된다. ( 염료가 확산되어 맨 마지막 보라색 액체처럼 보이는 형태를 지칭 ) 즉, 결국에 염료 분자들이 uniform하게 분포되는 PDF가 Uniform Distribution의 것과 같아진다는 것이다.
1.2. Recover Structure by Reversing Time
Diffusion Model의 코어는 여기에 있다. 앞서 확산이라는 현상의 프로세스와 결과를 살펴봤고, 확산이라는 현상을 확률적으로 모델링할 수 있음을 간단히 언급했다. Diffusion Model은 이러한 확산 현상의 역과정을 학습하고 이를 바탕으로 Reverse Process 즉 Noise로부터 Recover하는 것에 목적이 있다. 좀 더 자세히 살펴보자.
수많은 이미지들이 있고 이 이미지들의 분포인 data distribution이 있다고 생각해보자. 이러한 이미지들에 대해서 노이즈를 더해가면서 Diffusion Forward Process를 거쳐 ( 염료를 넣는 것을 노이즈를 주입하는 것으로 생각할 수 있게 되는 이유는 뒤에서 다시 설명하겠다. ) Uniform한 Distribution을 만들었다. 이 분포의 각 샘플들은 노이즈가 충분히 더해졌기에 그냥 일종의 노이즈와 같을 것이다. 이러한 과정을 역과정으로 즉 Uniform Noise Distribution으로부터 Data Distribution을 복원해내는 것이 Diffusion Model의 목표다.
그러나 사실 물리적으로 자연적으로 이러한 Diffusion의 역과정 즉 un-mix되는 과정은 잘 발생하지 않는다. 즉 macroscopic irreversiblity하다. 그런데 어떻게 이 역과정을 즉 복원을 가능하게 하는 것일까?
1.2.1. How Reverse Process Possible?
실제 확산 과정을 확대해보면 단서를 찾을 수 있다. 실제로 forward process와 reverse process에서 각각의 염료 분자의 움직임을 살펴보면 ( bright particle = 염료 분자 ) 이 분자들의 motion이 두 방향에서 모두 비슷해보인다. ( 거의 identical ) 이는 실제로 아인슈타인이 액체 속의 부유입자들의 확산에 대해서 접근한 방법을 통해 확인할 수 있는데, 확산은 실제로 이러한 부유입자들의 무작위 운동으로 인해 확산이 일어난다는 것이다. 이는 액체 분자들의 무작위 운동 때문에 일어나는 것으로 각각의 움직임은 결국 서로 영향을 주게 되고 모든 입자들의 지그재그 움직임을 만들어낸다. 즉 다시 말해 Brownian Motion과 Diffusion과의 관계에 대해 이야기하고 있다. 때문에 실제로 Forward 와 Reverse에서 Diffusion Kernel은 same functional form을 가진다.
더 나아가서 시간축을 기준으로 다음 time step에서의 분자의 위치는 매우 작은 가우시안으로부터 draw되는 것과 동일하게 생각 할 수 있게 된다. Forward와 Reverse 모두에 대해서 해당되는 말이다. 이것이 바로 Diffusion Process를 가능하게 해주는 원리이다. 왜냐하면 하나의 분자에 대해서 (= 하나의 데이터에 대해서 ) forward trajectroy를 seqeunce of small gaussian으로 모델링할 수 있기 때문이다. 그렇기에 실제로 Diffusion Model에서 forward과정에 각 이미지 샘플에 가우시안 노이즈를 추가해나가는 것이다. 한가지 잊지말아야할 점은 Forward 와 Reverse는 same functional form을 가진다는 것이다. 그렇기에 Reverse trajectory또한 seqeunce of small gaussian으로 모델링할 수 있게 된다. 그렇기에 Diffusion Model의 프로세스가 다음과 같이 디자인되는 것을 확인할 수 있다.
실제로 Diffusion Forward Process를 1d에 projection 시켜서 변화를 살펴보면 다음과 같다. 간단하게 점차적으로 각 data sample들이 Noise와 Mixing되면서 자연계의 분포로 변하는 것을 확인할 수 있다. 그러면 이렇게 structure는 다 지워지고 color noise들로만 남아있는 것을 확인할 수 있따.
동시에 Reverse Process ( DDPM 이후로는 Denoising Process라고도 함 ) 에서는 다음과 같이 역과정을 거치는 것을 확인할 수 있게 된다.
2. Mathematics Formation ( with. DDPM )
실제로 이러한 Diffusion의 물리적 성질이 Diffusion Model에 어떻게 녹아들어가있는지를 수학적으로 살펴보도록 하자. Diffusion Model은 Diffusion Process ( Forward ) 와 Reverse Process ( Denoising )으로 나뉜다. 앞서 언급했듯 Diffusion은 데이터의 분포를 structure를 무너뜨리는 과정이다. 앞서도 언급했듯 Diffusion Model의 목표는 바로 Reverse Process를 잘 학습해 ( 모델을 통해 ) Prior로부터 데이터의 분포를 복원하는 Generative Model을 만드는 것이다. 이를 위해서 Diffusion Model은 노이즈를 더해가며 데이터의 분포를 우리가 원하는 prior ( latent )로 무너뜨리고, 이를 다시 기존의 분포를 따르도록 학습시켜나가는 과정을 거친다. 이와 같은 각 Process는 점차적으로 time step에 따라서 진행되며 앞서 살펴봤던 Diffusion의 개념을 간단히 떠올려보더라도 Markovian Chain에 기반해서 이러한 process를 구성하는 것은 납득이 된다.
Markovian Chain
'특정 state의 정보는 바로 이전 상태에 의존적'이라는 markov property를 가지는 discrete probability process를 markoivan chain이라고 한다. 즉 time t+1에서의 probability 는 오로지 t의 probability에만 의존한다. 즉 condition이 오직 t
이때 이 각각의 step마다의 conditional distribution으로 정의되는 noising 혹은 denoising process는 위에서 물리적인 이론에 근거해 각각 가우시안임을 확인할 수 있다. 그런데, Reverse Process에 해당하는 noise를 add하는 것의 reverse condition distribution은 bayes' rule로 바꿔보면, q(x_t)와 q(x_t-1)을 알아야한다. 하지만, 일반적으로 stochastic optimization을 minibatch단위로 진행하는 학습구조에서 이를 수학적으로 구하기는 어렵다. 따라서, q(x_t-1|x_t)라는 reverse process의 한 step은 이를 approximation하는 p(x_t-1|x_t)를 머신러닝을 통해 학습하는 방식으로 진행된다.
따라서 전체적으로 Diffusion Model Overview는 다음과 같다. time step T를 hyperparameter로 설정하고 ( DDPM에서는 1000 사용 ) non-learnable한 diffusion process와 learnable한 reverse process를 통해 기존의 데이터 분포를 따르도록하는 generative model를 만들어내는 것이다.
2.1. Forward Process
그렇다면 각 process는 어떤 방식으로 모델링되어 있고, 어떻게 동작하는지를 조금 더 수학적으로 서술해보도록 하겠다. 앞서 충분히 언급했듯, forward process는 각 이미지(데이터 1개)에 작은 가우시안 노이즈를 점진적으로 더해가는 과정이다. 이때 이 가우시안 노이즈는 brownian motion에서도 언급했듯 바로 이전의 상태에만 영향을 받는다는 가정하에 진행되고 따라서 markov chain으로 구성된다. 그렇다면 이 markov chain은 결국 이전 상태가 주어졌을 때 가우시안 노이즈가 주입된 상태를 나타내는 조건부 확률분포를 의미할 것이고 이는 다시 말해 주입되는 가우시안 노이즈를 의미할 것이다. 따라서 q(x_t|x_t-1)이라는 조건부 확률분포는 Gaussian Distribution으로 정의된다.
Why Conditional Gaussian Distribution == Add Gaussian Nosie?
왜 이러한 Gaussian Noise 주입 과정이 Conditional Gaussian Distribution에서 샘플링하는 방식으로 구성되는지를 확인해보도록 하겠다. 아래의 과정을 살펴보면 되는데, 우선 다음과 같은 두 statements는 동일하다. 즉 이전 이미지를 mean으로 가지고 variance가 1인 Gaussian 분포에서 샘플링하는 것은 이전 이미지에다가 노이즈를 주입시키는 것과 동일하다. 왜 그럴지 가우시안의 PDF를 풀어서 식을 전개해보게 되면, 아래와 같이 어렵지 않게 p(xt)를 conditional gaussian이 아닌 shifted variable에 대해서 적용된 gaussian으로 바꿀 수 있게 되고, 이는 결국 Convolution Operation과 동일해지고, 결과적으로 노이즈를 주입하는 것과 동일하게 표현된다. 하단의 표현이 정말 잘 정리된 표현 같다.In other words, we have show that asserting the distribution of a timestep conditioned on the previous one via the mean of a Gaussian distribution is equivalent to asserting that the distribution of a given timestep is that of the previous one with the addition of Gaussian noise.
이러한 conditional gaussian distribution은 하단의 식과 같이 바로 이전의 이전의 이미지와 specific variance schedule ( β )에 의해 결정된다. 정리하자면, markov chain에 의해서 바로 이전 이미지에 가우시안 노이즈를 주입해준다는 것은 이전 이미지가 주어졌을 때 다음 이미지의 probability density를 정의하는 conditional gaussian distribution으로 나타낼 수 있다는 것이다. ( Noise Equation )
β β가 정확히 의미하는 것은 무엇일까? image는 픽셀들로 이루어져 있다. 이러한 픽셀들에 gaussian noise를 더해주게 될텐데, larger β가 의미하는 거은 pixel 분포에 더 넓게 즉 greater variance로 노이즈를 추가해준다는 것을 의미한다. 또한, mean의 관점에서는 normal curver가 origin 혹은 center로 움직이도록 한다는 것이다. 그러므로 이러한 이미지는 original에서 멀어짐에따라 점점 더 corrupt된다는 것을 의미하고 따라서 β는 structure destroy정도 혹은 가우시안 노이즈 주입 정도 등으로 이해할 수 있게 된다.
이때 markov chain에 의해서 정의되기에 임의의 time step t까지의 noise process를 다음과 같이 각 조건부 확률분포들의 product로 표현할 수 있게 된다.
2.1.1. β ( variation schedule ) Scheduling
그렇다면 처음에 swiss roll과 같은 형태였던 데이터의 분포는 structure는 점점 무너질 것이고, data sample 하나에 대해서 살펴보자면 다음과 같이 점점 noise와 mixing되어가며 noise map이 되어가는 것처럼 보일 것이다. 이때 앞서 Diffusion에 대해서 물리적인 직관을 살펴보면서 언급했던 중요한 부분 중에 하나는 '작은' 가우시안 노이즈가 더해져야한다는 것이었다. 이를 위해서 deterministic한 forward process에서 더해줄 가우시안 노이즈의 정도를 정해줄 수 있는데 이는 하이퍼파라미터인 β에 의해서 조정된다. DDPM에서는 0.0001 ~ 0.02까지 점진적으로 베타를 늘려가면서 스케줄링해주게 된다. 더 나아가서 이러한 스케줄링 방법에 대해서 실험을 통한 결과를 제시하고 있고 이는 다음과 같다. 하지만, 근래에는 cosine annealing과 같은 스케줄링이 가장 일반적으로 좋은 퍼포먼스를 내는 방법이라고 알려져 있다.
Data Distribution의 Diffusion Process에 따른 Structure Destroy / Data Sample에 Diffusion Process에 따른 nosie mixing
2.1.2. Reparameterization Trick
물론 이러한 gaussian noise를 주입해나가는 과정은 다음과 같은 conditional gaussian distribution에서 노이즈를 샘플링해서 더해나가는 과정을 의미하기에 graident가 흐르기 위해서는 이러한 샘플링 프로세스를 VAE에서 사용했던 것처럼 Reparameterization Trick을 사용해야한다. 그렇게 각 time step에서 노이즈가 샘플링되어 더해지는 과정을 reparameterization trick을 통해 closed form으로 사용해야한다. 그렇게 되면 다음과 같이 써지는 것을 확인할 수 있다.
2.2. Reverse Process
Diffusion Model의 Training은 이 Reverse Process를 통해서 이루어진다. 직관적으로 말하자면, Diffusion Process에서 주입한 노이즈들과 Denosing Process에서 denoising하게 될 노이즈가 유사해지도록 학습을 진행하게 된다.(다만, 역방향의 markov chain이라는 점) 이때 앞서 언급했듯 p(x_t-1|x_t)는 다음과 같이 표현할 수 있고 이러한 denoising의 역할을 담당하는 conditional distribution을 noise와 same functional form으로 생각할 수 있기에 이러한 denoising을 학습하고자 하는 것이다. 또한denoising과정은 conditional gaussian distribution에서 샘플링된 것을 사용하는 구조이기에 결국 이러한 distribution의 파라미터인 mean&variance를 diffusion process에서와 유사하도록 estimation하는 것이다. 하단의 코드 예시에서 결국 Denoising을 담당하는 모델이 output으로 mean과 variance를 estimation하고 있다는 것을 통해 이를 다시 한번 확인할 수 있다. ( 명확히 인지해야한다. 학습대상이라고 적혀있는 것은 mean과 variance라는 파라미터를 학습한다는 것이 아니라 네트워크가 있고 해당 네트워크 파라미터를 학습해나가면서 denoising에 더 적절한 mean과 variance를 model이 estimation해내도록 학습해나간다는 것이다. )
Data Distribution의 Reverse Process에 따른 Recover Structure / Data Sample에 Reverse Process에 따른 data generation
2.2.1 Objective Function ( Optimization Problem )
이러한 Reverse Process가 Training되기 위해서 Diffusion Model은 기존의 VAE와 같은 생성모델처럼 training data의 likelihood를 maximize하는 방향으로 reverse markov transition을 학습한다. 정리하자면 p(x0) 즉 training data의 likelihood를 maximize하도록하는 것이 diffusion model의 학습목표이다. 그러나 VAE에서와 마찬가지로, diffusion model에서는 reverse process의 q를 명확하게 구할 수 없기에 이를 근사하는 p를 사용했다. 따라서 다시 VAE와 유사하게 이러한 likelihood(정확히는 negative log likelihood) 대신에 ELBO를 통해 upper-bound를 minimizing하도록 학습을 진행한다.
VAE에서 MLE가 왜 intractable한지
2.2.1.1. ELBO
따라서 결과적으로 풀고자하는 것은 하단의 L_vlb라는 upper bound를 minimize하는 optimization problem이 된다.
2.2.1.2. markov assumption & log transformation
이러한 variational bound를 다시 한번 markov assumption에 의해서 xT에 대한 것과 나머지를 분리해낼 수 있고, log transformation을 통해서 summation으로 표현할 수 있다.
2.2.1.3. Bayes' rule at second term
이때 Expectation안의 3가지 term 중에 가운데 term의 q(x_t|x_t-1)은 bayes' rule에 의해서 p와 유사하게 condition의 순서를 바꿀 수 있고 이 역시 log로 수식이 구성되어 있으므로 다시 한번 나눌 수 있게 된다.
여기서 다시 한번 두번째 term을 정리하면 t=1 ~ T-1 까지의 결과들이 서로 다 지워지게 되어 다음과 같이 깔끔한 결과를 얻을 수 있게 된다. 이를 다시 기존의 equation에 대입하게 되면 지워질 것들은 지워지고 맨 아래와 같이 정리된다.
중간 과정에서 도출된 optimization하기를 원하는 식
2.2.1.4. Hint : KL-Divergence
이렇게 정리된 수식의 첫번째 term과 두번째 term을 자세히 들여다보면 KL-Divergence와 유사한 것을 확인할 수 있다. 첫번째 term은 q(x_T|x_0)와 p(x_T)간의 KL-Divergence, 두번째 term은 q(x_t-1|x_t,x_0)와 p(x_t-1|x_t)간의 KL-Divergence들의 합으로 정리할 수 있게 된다.
KL-Divergence
최종적으로 Variational Bound에 대해서 정리된 식으로 이를 minimize하는 것이 목표
그렇게 되면 결국 L_vlb는 다음과 같이 KL divergence들로 표현이 되며 L0 / L1~L_T-1 / L_T 와 같이 세개의 범주로 구분할 수 있게 되고 위에서 유도했던 결과를 결합해서 표현하자면 다음과 같아지고 이것의 우리의 objective function이다.
이때 L_T 는 forward process에서 정해지는 time-dependet constants인 베타값들에 의해 deterministic하게 되기에 현재 학습하고자하는 learnable parameter의 관점에서는 constant로 생각할 수 있다.
2.2.1.5. Second term KL-Divergence propotional to mean distance
반면에 L_1:T-1 에서 reverse makov transitions는 gaussian이라는 가정이 있었기에 다음과 같이 표현할 수 있고, 이때 이 mean과 variance를 parametrize하고 학습을 진행하는 것이 diffusion model의 목표이다.따라서, 이러한 mean과 variance를 time step t와 x_t가 주어졌을 때 estimate할 수 있는 neural network f를 두고 이 mean과 variance를 예측하도록 한다. ( 위에서 code example을 통해 한번 살펴봤던 그 부분이다. ) 어쨋든 이러한 mean과 variance를 estimate하는 model parameter를 learn하고자하고 이때의 L_t-1이라는 term은 다음과 같이 KL-Divergence에 의해서 정리되는데 두 distribution모두 gaussian이기에 reverse process에서의 mean을 x0와 xt가 condition일때의 reverse markovian에 의해서 구한 것과 network에 의해서 estimation 된 것 사이의 L2-norm에 비례하게 되고 이를 활용하게 된다.
이것이 가능한 이유는 variance term을 time schedule varaiance를 이용해서 상수화하는 것이 가능하다고 ddpm에서 입증했기 때문이다.
이를 위해서는 각각의 mean에 대해서 이해해야하고 다시 이를 위해서는 선행적으로 q(x_t-1|x_t,x_0)에 대해서 먼저 수학적으로 풀어볼 필요가 있다. 즉, reverse process의 conditional probabilites인 q(x_t-1|x_t)를 estimate해야하는데, 이는 쉽게 구할 수 없기에 앞서 이를 Neural Network를 통해서 p(x_t-1|x_t)라는 approximation model을 활용해서 이를 학습하고자 한다. 그런데 이러한q(x_t-1|x_t)는 intractable할지라도 여기에 x0 즉 데이터를 condition으로 주게되면 tractable하다는 점에 주목할 필요가 있다. 하단의 수식 전개 과정을 잘 살펴보게 되면, 결국 우리가 원하는 것은 이러한 reverse process의 conditional gaussian을 아는 것이고 이는 mean과 variance에 의해서 parameterized된다. x0를 추가적으로 condition으로 부여한 conditional probability를 bayes' rule에 의해서 전개하면 다음과 같이 표현되고, 이는 Gaussian이라는 가정이 있으니 mean과 variance를 각각 파란색과 빨간색 부분으로 해석할 수 있다. 그렇게 식을 아래와 같이 정리할 수 있다.
이때 앞서 froward process에서 q(xt|xt-1)의 conditional probability를 확장해서 임의의 xt에 대해서 x0가 condition으로 주어졌을 때의 결과가 product에 의해서 정의되는 것을 확인했다. 그때 베타를 다음과 같이 알파를 활용해서 치환하고 식을 전개하면 다른 variance를 가진 2개의 가우시안의 merge는 다음과 같이 standard deviation이 더해진 형태로 전개된다는 것을 고려할 때 어렵지 않게 적을 수 있게 된다.
2.2.1.6. Finally estimation goal is epsilon! so reverse process is called denoising process.
이에 따르면 x0를 다음과 같이 xt와 noise를 활용해서 표현할 수 있게 된다. ( 앞서 gaussian distribution에서 샘플링하는 것과 노이즈를 주입하는 것이 같다에 대해 적었던 내용 참고. ) 이를 다시 뮤에 대해 전개된 수식에 대입해주면 하단의 하늘색처럼 결과가 정리되고, 이때 알파는 하이퍼파라미터인 베타들로 이루어져있으니 deterministic하게 되고, 결국 풀고자하는 것은 입실론이라는 노이즈 component를 예측하는 것으로 바뀌게 된다. ( forward process에서 넣어주는 noise가 reverse process에서 denoising을 통해서 이 입실론에 해당하는 값을 예측하도록 하게 된다. )
그렇기에 DDPM에서 학습해야하는 것은 결국 diffusion process와 denoising process의 노이즈를 다음과 같이 매칭시켜주는 것이 목표가 된다.
이를 정리하자면 결국 다음과 같이 입실론에 해당하는 noise에 대해서 다음과 같이 loss function을 쓸 수 있게 된다.
정리하자면 Diffusion의 Loss에 대한 수식은 아래와 같이 표현이 되고, 이는 결국 Noise를 주입했을 때의 noise와 denoising하기 위해서 denoise해주는 noise간의 차이를 줄여주도록 학습이 되고, 이렇게 함으로써 기존의 데이터 분포를 복원하는 recover가 학습된다.
2.2.1.7. Full Loss with Constant or Small term.
추가적으로 실제 학습 페이즈에서는 앞서 언급했듯 simplify된 loss를 사용하게 되는데 이때 Regularization Term의 경우 굳이 존재하지 않더라도 Diffusion과정 자체가 gaussian의 가정에 있으며 T가 무한대로 갈 때 istropic gaussian이 된다는 가정을 두고 있기에 실제 학습시에 사용하지 않게 된다. 즉, 굳이 가우시안을 따르도록 regularization을 해줄 필요가 없기에 이는 삭제 된다. 다음으로 Reconstruction Term의 경우 실제로 Diffusion이라는 물리적 현상 자체가 인접한 small time step에 대해서 small gaussian noise를 주입해주는 과정으로 진행되기에 x0와 x1에서의 차이는 거의 없다. 실제로 DDPM에서는 이에 대해서 ablation study또한 진행했고, reconstruction term은 너무 작기에 학습에 유의미한 영향을 끼치지 않기에 제거했다. 그래서 결과적으로는 저렇게 단순한 loss를 optimization하는 문제로 바뀌게 된다.
여기까지가 일반적인 Diffusion Model에 대해서 이해해야할 부분들과 DDPM에 대한 내용이다. 최근에는 이러한 markov chain에 의한 과정이 번거롭고 inference time을 너무 늘리는 등의 단점이 존재하기에 이를 implicit하게 해결하고자 했던 DDIM부터 raw data level에서의 diffusion process를 적용하는 것은 data에 대해서 (입자에 대해서) 가우시안 노이즈를 주입해야했던 ( 브라운 운동의 가정 ) 문제에서 벗어나 데이터를 dimension reduction시킨 즉 latent level에서 (어쨋든 데이터에 대해서 의미론적으로 같은 상태 ) diffusion process 및 reverse process를 적용하는 latent diffusion ( -> stable diffusion )까지 다양한 연구들이 활발히 진행되고 있다. 추후에 해당 내용들에 대해서 천천히 적어보도록 하겠다.
