Diffusion Sampler(DDPM&DDIM)

Paper Review 10, 11

|Intro

이번에 리뷰할 논문은 Diffusion Model의 가장 기본적인 Sampler인 DDPM[2]과 이를 개선한 DDIM[3]입니다. 이번 포스팅은 두 개의 논문을 자세히 해설한다기 보다는 이전에 포스팅한 'Understanding Diffusion Model[1]'에서 설명한 내용을 기반으로 DDPM과 DDIM 논문을 수식 기반으로 읽어보고, 이를 코드로 구현해보고자 합니다. (이번 포스팅에선 이전에 리뷰한 Understanding Diffusion Models[1]에서 설명한 수식이나 개념 설명을 최소화했으니, 모르는 부분이 있다면 해당 논문 및 제 이전 포스팅을 참조하시기 바랍니다)

Diffusion Model을 구성하는 요소는 많습니다. 단순히 Backbone 역할을 하는 모델 외에도, 이를 Denoising하는 방법 그 자체에 대한 부분을 다루는 Sampler나, 이러한 Sampler를 어떤 규칙에 의해 적용할지 결정하는 Scheduler들이 그것들 중 하나입니다. 

또한 이러한 구성요소들 하나하나에도 수많은 논문들이 존재하며, 아직까지 무엇이 우위다라고 보기 명확하게 어려운 상황인 경우가 많습니다. 하지만 그러한 것들의 시작점은 명확합니다. 바로 Sampler, 그 중에서도 DDPM입니다. 구체적으로 Diffusion 모델을 최초로 유효하게 만든 DDPM에서 시작해, 해당 알고리즘의 어느 부분을 개선했다라는 포인트에서 여러 Sampler들과 Scheduler들이 파생된 셈입니다.

때문에 본 포스팅에서는 이러한 DDPM과 이를 개선한 DDIM을 함께 살펴보며 각각의 개념과 성능, 그리고 어떻게 개선되어 갔는지를 살펴볼 것입니다. 구체적으로 DDPM을 먼저 살펴보고, DDIM 논문을 이후에 살펴보면서 어떤 식으로 논문이 전개되고, 코드로 구현했을 땐 어떤 결과가 있었는지를 살펴볼 것입니다.

참고로 이번에 리뷰할 DDPM 논문의 제목은 'Denoising Diffusion Probabilistic Models'[2]이고, DDIM 논문은 'DENOISING DIFFUSION IMPLICIT MODELS'[3]입니다.

|Background

본 포스팅의 과정에서는 Diffusion Model의 가정 중 하나인 입력과 출력의 크기가 동일한 AutoEncoder 구조를 가정하는 VDM(Variational Diffusion Model)을 기본으로 전제하며, 이를 위해 VAE가 아닌 UNET이라는 모델을 사용합니다. UNET은 AE 계열의 모델과 유사하게 데이터를 압축했다가 다시 원래 크기로 되돌리며 원하는 형태의 결과물(Semantic Segmentation, Object detection 등)로 재생성해내는 알고리즘입니다. 이에 대해선 후에 코드를 살펴보며 다시 간단히 리뷰하도록 하겠습니다.

|Paper #1 - DDPM

본 논문에서 제시한 개념도는 아래와 같습니다. 아래 그림을 기준으로 노이즈 이미지에서 원본 이미지를 재생성하는 방향을 정방향(Forward), 원본 이미지에서 노이즈를 만들어 가우시안 노이즈 이미지로 만드는 방향을 역방향(Reverse)으로 정의합니다.

DDPM 개념도

DDPM을 본격적으로 들어가기 전에 DDPM[2]에서 UDM[1](Understanding Diffusion Models 논문을 이제부턴 'UDM'으로 줄여서 칭하겠습니다)과 다르게 정의한 용어나 별도로 정의한 내용에 대해서 짚고 넘어가도록 하겠습니다.

|용어 정의 + 방식 차이

DDPM에서도 우리가 UDM 논문에서 살펴봤듯이 근사된 사후확률 분포와 전이확률 분포를 이용해 최적화함으로써 우리가 원하던 목표인, 제한된 시간 내에 이미지를 생성해내는 과정을 달성하고자 합니다. 다만 DDPM에서는 원본 이미지 x0를 표준 가우시안 노이즈가 생성된 xT로 보내는 분포( pθ(x0:T) )를 reverse process라는 용어로 정의하며, 수식은 다음과 같습니다.

reverse process 수식

이때 (UDM에서 설명한 VDM과 동일한 가정을 따르기에) xT는 표준 가우시안 분포, 즉 p(xT ) = N (xT ; 0, I)입니다. 그리고 근사된 사후확률(q(x1:T |x0))은 forward process 혹은 diffusion process이라는 용어로 정의하며, 수식은 다음과 같습니다.

forward process 수식

위 수식의 분포를 살펴보면 베타 스케줄 βt에 따라 데이터에 점차적으로 가우시안 노이즈를 추가하는 마르코프 체인으로 고정됨을 알 수 있습니다. 

이는 우리가 UDM에서 정의했던방식과는 다소 다릅니다. 우리는 당시 α를 평균 중심의 해석을 진행했으나, 이를 분산을 기준으로 해석할 수도 있습니다. 구체적으로 이전에는 α를 이용해 q(xt|x0) = N (xt; √ α¯tx0,(1 − α¯t)I)와 같은 식으로 정의했지만, 분산을 기준으로 식을 전개하고 효율적인 학습을 위해 DDPM 논문에서는 αt와 α¯t를 다음과 같이 정의합니다.

이후 이들을 이용해 다음과 같이 정의할 수 있습니다.

위 수식을 통해 알 수 있듯 forward process가 표준 가우시안 분포를 따르게 됨에 따라 폐쇄형(closed form)으로 계산이 가능해지며, 이를 통해 임의의 timestep t에서의 xt를 바로 샘플링하는 것이 가능해지게 됩니다.

이러한 DDPM의 정의에 따라 용어를 정의한 뒤에는 모델 학습을 위해 UDM과 동일하게 ELBO에 대한 로그 가능도를 최적화함으로써 수행됩니다(음의 로그가능도로 진행하는 차이가 있습니다). 그 수식은 아래와 같습니다. 비교를 위해 이전에 리뷰한 UDM에서의 공식과 본 DDPM 논문에서 공식을 함께 보여드립니다.

UDM 논문[1]의 ELBO적용 - 수식(47)

ELBO 수식 적용 - 수식(3)

위 수식에서 pθ(x0:T)는 시간 T까지의 전체 경로 x0:T에 대한 결합확률로, 마르코프 속성을 이용해 분해하고 수식을 정리하면 다음과 같이 표현할 수 있으며, 이를  'L'로 정의합니다. 

L에 대한 정의

이후에는 VDM의 세 가지 제약조건과 베이즈 규칙을 적용해 아래와 같이 reconstruction term과 prior matching term, 그리고 denoising matching term으로 구분하였습니다.

UDM 논문[1]의 전개된 수식

이러한 내용을 DDPM 논문에서는 아래와 같이 표현했습니다.

DDPM 최적화 수식[2]

위 수식을 기준으로 다시 표현해보면, Lt가 prior matching term으로 0에 수렴할 것이고, 결국 denoising matching term인 Lt-1과 reconstruction term인 L0가 최소가 되는 지점을 찾음으로써 최적화가 가능해질 것입니다. 조금 더 설명하자면 위 수식은 x0를 조건부로 고려할 때 추적가능한 순방향 사후확률(forward process posteriors)과 역방향을 구현하는 신경의 pθ(xt−1∣xt)를 비교해 최적화하는 것입니다. 

이때 pθ(xt−1∣xt)는 xt가 주어졌을 때, xt-1의 분포를 정규분포로 모델링하는 가우시안 분포(N)이며, 순방향 사후확률 또한 원본 이미지에 노이즈를 더하는 것으로 가우시안 분포(N)로 모델링됩니다. 즉, 아래가 성립합니다.

이로 인해 KL Divergence가 두 가우시안 분포를 직접적으로 비교하는 것이 가능해지며, 기존의 몬테카를로 추정 대신 두 분포가 가우시안 분포를 가지고 있으므로 폐쇄형 표현식을 이용한 라오-블랙웰 방식(Rao-Blackwellized fashion)으로 계산됨으로써 더 적은 분산, 더 높은 정확도, 더 빠른 연산이라는 장점을 취할 수 있게 되는 것입니다.

|Diffusion models and denoising autoencoders

|Forward Process & LT

DDPM의 Forward Process Variance인 βt는 재매개변수화를 통해 학습될 수도 있고, 상수값으로 고정될 수도 있습니다. 또한 DDPM에서는 근사 사후확률 q(노이즈를 추가하는 과정)에 학습가능한 파라미터가 없기 때문에, prior matching term인 LT는 학습 과정에서 상수이며, 이는 무시될 수 있음을 의미합니다.

이를 우리가 실제 사용할 sprite dataset을 통해 예시를 들면 아래와 같습니다. 참고로 sprite dataset은 우리가 비디오 게임 등에서 사용되는 게임 캐릭터 및 아이템 이미지입니다.

forward process 시각화

|Backward Process & L1:T-1

DDPM의 Reverse Process에서 도출하는 수식은 우리가 이전 포스팅에서 VDM을 최적화할 수 있는 세 가지 방법을 살펴보았고, 그 중 두 번째 최적화 방법에서는 다음과 같이 최적화 문제를 풀었었습니다.

2nd Interpretation of VDM Optimization[1]

그리고 DDPM 논문에서 제시한 간소화된 최적화 수식은 다음과 같습니다.

VDM Opitmization of DDPM[2]

위 두 수식을 살펴보면 결국 β를 통해 표현을 간결하게 한 것임을 확인할 수 있습니다(정확히는 이전에 읽은 논문 자체가 Diffusion 관련 사항들을 잘 정리해놓은 것이기 때문에 DDPM 논문에서 나왔던 내용을 잘 정리한 것에 가깝지 않나 싶습니다. 실제로 x0를 예측하는 것보다 원본 노이즈 ϵ0를 예측하는 편이 성능이 더 좋았다는 내용도 동일하게 나옵니다).

|Data scaling, reverse process decoder and L0

이미지 데이터는 일반적으로 3차원(RGB)의 0~255 사이의 정수 값을 가진 데이터이며 수식을 적용하는 과정에서 이를 고려해야 합니다. 첫 번째로 데이터 스케일링 과정에서 0~255 사이의 정수를 [-1, 1] 사이의 값들로 스케일링함으로써 신경망의 reverse process가 표준정규 사후분포 p(xT)로부터 일관되게 스케일링된 입력에서 작동하도록 보장하게 됩니다.

두 번째로 reverse process에서 이산적인 log likelihood를 얻기 위해, reverse process의 마지막 수식을 우리가 가정한 가우시안 분포 N (x0; µθ (x1, 1), σ1^2 I)에서 유도된 independent discrete decoder로 설정하며, 이에 대한 수식은 아래와 같습니다.

reverse process의 마지막 term

위 수식에서 D는 이미지의 총 픽셀개수를 의미하고, i는 각각의 픽셀을 의미하며, 마지막 샘플링 µθ (x1, 1)를 무소음(noiseless)으로 정의합니다. 이는 후술할 샘플링 알고리즘에서 마지막 단계인 x0를 샘플링할 때, 모델의 다양성을 위해 이전 단계에서는 추가되던 노이즈를 추가하지 않는다는 것을 의미하며, 그 이유는 x0의 i번째 픽셀값이 x1으로부터 에측된 평균의 i번째 픽셀값 µθ (x1, 1)과 동일하면 로그 가능도가 최대가 되기 때문입니다.

이러한 처리를 통해 DDPM은 이산적인 데이터를 다루면서도 잡음을 추가하거나 스케일링 작업의 Jacobian을 log likelihood에 포함시키는 과정(데이터를 변환하는 과정에서 발생하는 부피 변화를 고려하도록 하는 것) 없이도 원본 특성을 잘 압축하고 복원할 수 있게 되는 것입니다. 

|Simplified training objective

이전까지의 과정을 통해 우리는 L0를 LT로 보냈다가 다시 L0로 복원하는 과정에 대한 손실함수를 정의했습니다. 이러한 손실함수들은 모두 미분가능하도록 유도했기 때문에 그대로 사용할 수도 있지만, 다음과 같이 간략화시켜 사용함으로써 보다 효율적으로 연산이 가능해집니다.

간소화된 손실함수 식

위 수식은 이전 포스팅의 유도 과정을 통해 도출된 친숙한 형태의 손실함수이며, 구체적으로 본 포스팅에서 이전에 정의한 손실함수들에서 노름 앞에 곱해지는 상수를 제거한 것과 동일한 수식입니다. 위 수식에서 t는 1과 T 사이의 timestep을 표현하는 것으로 t=1일 때는 L0가 되고, t>1일 때는 수식 12와 같아지게 됩니다.

이에 대한 자세한 설명은 제 이전 포스팅 혹은 여기[5]를 참조하시는 걸 추천드립니다.

|Experiments

이렇게 손실함수를 정의했다면 이번에는 실제 데이터와 코드를 통해 구현해볼 것입니다. 우선 코드로 구현하기 전에 본 논문에서 제시한 알고리즘을 다시 살펴보도록 하겠습니다.

|Algorithm 1 - Training

Training algorithm

여태까지 정리한 내용을 바탕으로 위 알고리즘을 해석하면 다음과 같습니다.

---

1. 아래 과정을 반복

2. 임의의 학습 데이터 x0를 선택

3. 임의의 시간 단계를 선택, t는 임의의 시간 단계에 대해 uniform함

4. 표준 정규분포로부터 노이즈 ϵ을 샘플링

5. 위 단계에서 선택한 x0, t, ϵ으로부터 xt를 만들고, 모델에 입력해 출력으로 ϵθ를 생성.

   이를 4단계에서 샘플링한 노이즈 ϵ와 비교해 그래디언트를 구하고, 경사하강법 수행

6. 수렴하면 위 과정을 종료

|Algorithm 2 - Sampling

Sampling Algorithm

다음으로는 샘플링 알고리즘입니다. 이전에 이미 수차례 언급한 바와 같이 여기서는 xt를 통해 x0를 예측하기보다 ϵθ를 통해 ϵ을 예측하는 것에 초점을 맞춥니다. 위 알고리즘을 해서하면 다음과 같습니다.

---

1. xT를 표준 정규분포로부터 샘플링

2. T에서부터 1에 이르는 시간단계를 t에 대입하며 for loop

3. t>1일 때, 추가할 노이즈 ｚ를 표준 정규분포에서 샘플링.

    t=1일 때는 추가되지 않으며, 이는 L0는 noiselessly 해야하기 때문)

   (이때 ｚ는 모델의 예측결과에 더해짐으로써 확률적 다양성을 확보하기 위해 별도로 추가된 노이즈입니다.)

4. xt-1은 실제 xt와 xt를 입력으로 받아 모델이 출력하는 값과의 분포 차이, 그리고 t>1일때 추가되는 노이즈로 구성됩니다.

5. for문이 끝날 때까지 반복합니다.

6. for문의 결과로 연산된 x0를 반환합니다.

|Code - Preview

Diffusion 구현은 이전의 논문에서 봤다시피 3가지 제약조건이 추가된 VDM으로 구현됩니다. 이는 출력으로 나오는 결과의 형상이 입력의 형상과 동일해야 하며, 이를 충족하는 모델 및 학습을 위한 손실을 계산하는 클래스를 간단히 구현하면 아래와 같습니다.

Simple DDPM

이에 대해 조금 더 자세히 살펴보도록 하겠습니다.

|Code - Modeling

우선 우리가 앞서 살펴본 3가지 제약조건이 추가된 VDM의 구조를 충족하되 조금 더 성능이 좋은 모델인 Unet으로 업데이트해보도록 하겠습니다. Unet의 구조에 대해 간단히 살펴보면 아래와 같습니다.

unet structure

Unet의 구조는 공간차원을 줄이되 채널은 늘리는 과정을 반복하다가 다시 채널을 줄이고 공간차원을 늘려가며 본래의 형상으로 복원함으로써 모델에게 요구하는 결과물(ex. Semantic Segmentation)을 생성해내는 방법입니다. 일반적으로 의미론적 분할 등의 방법에 많이 사용되나 이러한 구조를 backbone으로 하여 이미지를 해석하고, 그 뒤의 추가적인 구조를 통해 다른 task에도 적용되는 구조입니다.

최근엔 Unet 구조보다 더 좋은 모델들을 이용해 학습시킨 모델들이 하나둘씩 등장하며 더 좋은 성능을 보인다는 포스팅을 종종 볼 수 있는데요, 다만 가장 많이 사용되고 연구된 stable diffusion이 unet 구조를 채택해 open source로 수많은 개발자 등이 사용 및 연구되고 있어 Diffusion 계열에서 unet이 가장 많이 활용되는 모델이 아닐까 싶습니다.(아직까진요)

Unet 구조를 구현한 코드를 살펴보면 다음과 같습니다. 이미지 내 글씨가 너무 작아 확인이 어렵다면 제 github에 올라온 jupyter notebook을 직접 보며 확인하는 것을 추천드립니다. 해당 코드는 여기를 눌러 확인할 수 있습니다.

unet modeling

참고로 DDPM 논문에서 위 모델의 역할에 대해선 아래 그림의 하이라이트된 부분을 의미합니다.

모델의 수식[2]

|Code - Training

모델의 입력[2]

모델을 구현했다면 그 다음으로는 모델에 입력으로 넣을 부분을 구현하는 코드가 필요합니다. 위 그림의 하이라이트된 부분이며, 이를 코드로 구현하면 다음과 같습니다.

model input code

위 코드를 통해 timestep t에 따라 선형적으로 감소하는 noise schedule을 ab_t로 구현하고, 이를 이미지 x와 곱하며 이미지에 노이즈를 더하며 여기에 실험적으로 성능이 높아진다 알려진 임의의 노이즈를 더하게 됩니다.

|Code - Sampling

이러한 코드를 통해 모델을 학습시킨 뒤에는 샘플링 코드를 통해 결과를 확인할 수 있는데요.

Sampling 알고리즘[2]

위의 샘플링 알고리즘을 코드로 구현하면 아래와 같습니다.

Sampling code

여기까지의 코드로 diffusion model을 학습시킨 결과는 아래와 같습니다. 참고로 본 학습을 위해서 sprite dataset을 사용했는데요, 아래에서 확인할 수 있듯 16x16의 작은 사이즈의 비디오 게임에서 사용되는 캐릭터와 아이템 등에 대한 이미지 데이터셋입니다. 해당 데이터셋은 사이즈가 작아 npy 형태로 제 github에도 올려두었으며, 여기에서 확인할 수 있습니다.

먼저 50에포크 학습시킨 결과입니다.

50에포크 학습 결과

대략적인 형상은 보이지만 몇몇 이미지에서 대충 뭉개진 이미지가 많이 보입니다. 다음으로 200에포크 학습 결과입니다.

200에포크 학습 결과

200에포크 학습에선 보다 분명하게 이미지를 생성해내는 것을 확인할 수 있습니다 (물론 영웅, 아이템은 잘 생성하지만, 몬스터쪽은 어려워하는 것 같긴 하네요). 마지막으로 DDPM의 Sampling 과정에서 우리는 다양성을 확보하기 위해 임의의 노이즈를 추가하였었는데요, 수식에선 다음에 해당합니다.

다양성을 위해 임의의 노이즈를 추가하는 수식[2]

이 부분이 정말로 긍정적인 영향을 미치는 것이 맞을까요? 이 부분을 없애면 어떻게 될지 확인해보면 좋을 것 같습니다. 이를 위해 코드를 아래와 같이 수정하고 샘플링한 결과를 살펴보도록 하겠습니다.

다양성을 위한 노이즈 제거

이전과 동일하게 200에포크를 학습시킨 모델로 32개의 임의의 샘플을 뽑아 생성하면 아래와 같습니다.

임의의 노이즈가 제거된 결과

임의의 노이즈가 제거됐을뿐인데 대부분의 이미지가 기존과 달리 다양성 없이 거의 유사한 형태의 결과물들만을 생성해내는 것을 확인할 수 있습니다. (참고로 이러한 임의의 노이즈를 추가하는 것은 이후에 보다 빠른 생성속도를 위해, 혹은 텍스트나 이미지가 추가적인 컨디션으로 제시되는 과정에서 제거되기도 합니다)

|Paper #2 - DDIM

DDIM[3]은 DDPM의 샘플링 과정을 개선함으로써 훨씬 빠른 생성이 가능하도록 합니다. 구체적으로 기존에 확률론적으로 진행되던 샘플링 과정을 결정론적으로 바꿈으로써 더 적은 timestep만으로 DDPM 이상의 성능을 낼 수 있게 됩니다. 즉 사용하는 모델은 동일하지만, 이를 적용하는 방식에 따라 속도 및 성능 차이를 낸다는  것인데요, 이를 그림으로 비교하면 아래와 같습니다.

기존 diffusion process

기존 DDPM에 의한 diffusion process를 그림으로 나타내면 위와 같았습니다. 이를 DDIM에 의한 프로세스로 변경하면 아래와 같이 일부 단계를 건너뛰며 진행할 수 있게 되는 식입니다.

DDIM Process

이러한 알고리즘을 적용한 샘플링 속도 차이를 직접 실험해 비교한 결과는 아래와 같았습니다.

sampling 속도 비교

위 그림과 같이 3차원의 32x32 이미지를 대상으로 실험했을 경우, DDIM은 DDPM에 비해 약 30배 전후의 속도 차이를 보여줍니다. 이는 결정론적 알고리즘을 통해 DDPM이 정상적인 성능을 내기 위해 약 1,000 timestep 이상이 요구되던 반면, DDIM의 경우 아래 그림과 같이 50 timestep 정도면 충분하기 때문입니다. (물론 더 긴 timestep을 적용할 경우 DDIM의 성능 - 이미지 퀄리티 -도 어느정도 더 올라가긴 합니다.)

|How to Speed up?

timestep별 DDIM의 재구성오차[3]

그렇다면 어떻게 DDIM에서는 더 적은 수의 Timestep으로 같거나 더 좋은 수준의 성능을 보일 수 있었을까요? 해당 논문에서 제시하는 핵심 수식은 다음과 같습니다.

Sampling 수식[3]

위 수식에 대해 먼저 간략히 설명하자면 샘플링 과정에서 확률론적으로 작동하는 노이즈를 제거해도 이미지의 품질을 유지할 수 있도록 하는 것입니다. 이는 우리가 DDPM의 마지막 부분에서 살펴본 임의의 노이즈를 제거함으로써 생성되는 결과물의 다양성을 확보할 수 있다는 내용에 위배되는 것이기도 합니다. 

위 수식을 간략히 살펴보면 DDPM과 가장 큰 차이를 보이는 것은 t-1 시점의 이미지 x를 만들어내기 위해 t시점뿐 아니라 t-1시점의 변수까지도 고려된다는 것입니다. 즉 t 시점과 그 전 t-1 시점 사이의 관계를 파악해 노이즈의 방향성을 계산함으로써 DDPM과 같이 임의의 노이즈를 각 단계마다 추가하지 않고도, 미세한 변동과 차이를 고려한 이미지를 생성해낼 수 있다는 것입니다.

이를 하나씩 살펴보며 어떻게 이게 가능한 것인지 확인해보면, 첫 번째 'predicted x0'는 현재 상태 xt에서 다음 상태 xt-1로 이동할 때, 예측된 원본 이미지 x0를 기반으로 한다는 것이며 여기서 ϵθ는 모델이 예측한 노이즈를 의미합니다.

다음으로 'direction pointint to xt'는 현재 상태 xt로부터 노이즈를 제거하는 방향을 나타내며 마찬가지로 모델이 예측한 노이즈를 기반으로 합니다.

마지막으로 'random noise'는 DDPM에서와 동일하게 확률론적으로 각 timestep에 따라 추가되는 표준 가우시안 노이즈를 의미하며, 각 샘플링 단계에서 무작위성을 추가함으로써 모델의 유연성을 확보하는 부분이었습니다.

위 수식에서 σt가 어떻게 정의되냐에 따라 확률론적인 DDPM이 될수도, 결정론적인 DDIM이 될 수도 있습니다. 구체적으로 DDPM이 되기 위해서는 σt가 아래의 본 논문 내용과 같이 DDPM 논문에서 α를 통해 정의되었던 βt에 기반해 정의되어야 하며, DDIM에서는 σt를 0으로 둠으로써 forward process를 결정론적으로 만들고 이를 통해 더욱 빠른 샘플링을 가능하도록 합니다.

위 수식이 DDPM or DDIM이 되는 가정[3]

하지만 우리는 위에서 DDPM의 임의의 노이즈를 추가하는 과정을 통해 모델의 다양성을 확보할 수 있음을 확인하였었는데요, DDIM에서는 임의의 노이즈 없이도 모델의 품질을 유지할 수 있도록 x0와 xt에 기반한 예측을 진행하도록 함으로써 이러한 문제를 최소화하고자 했지만 실험적으로 위 수식의 두 번째('direction pointint to xt')에 존재하는 noise의 경우 데이터셋에 따라 변경함으로써 최적화할 수 있는 하이퍼파라미터화하는 것이 더욱 최적임을 확인하였습니다.

따라서 두 번째 수식의 σt를 η로 변경하고, 이를 하이퍼 파라미터화하였으며 이를 실험한 결과는 아래와 같습니다.

η에 따른 성능비교[3]

위 도표에서 가장 마지막 줄에 위치한 σ hat은 DDPM, 즉 η가 1일뿐만 아니라 이후에 임의의 노이즈를 가산하는 경우를 의미합니다. 위 도표를 통해 봤을 때에는 η가 0인, 즉 완벽하게 결정론적인 알고리즘일 경우 적은 Sampling timestep에서 우수했습니다. 또한 timestep이 1,000과 같이 충분히 클 수 있다면 DDPM과 같은 방법론이 오히려 더욱 효과적이었다는 사실도 확인할 수 있습니다.

|Code

이에 대한 내용을 코드로 다시 살펴보도록 하겠습니다. 이에 대한 전체 코드는 여기를 눌러 확인하실 수 있습니다. DDIM에서 사용하는 모델 자체는, 정확히는 이미지 생성을 위해 사용되는 모든 종류의 모델들에 사용되는 Sampler나 Scheduler는 모델의 형태에 영향을 받지 않습니다. 

ddim에 사용되는 모델 및 샘플링에 대한 설명[3]

즉 동일한 모델로 각기 다른 샘플러나 스케줄러를 사용할 수 있으며, 여기서 우리가 살펴보는 DDPM과 DDIM의 차이는 샘플링 영역, 정확히는 어떻게 디노이징 할 것인가에 대한 방법에 대한 차이만이 존재합니다.

|DDPM Sampling

때문에 먼저 DDPM에 대한 디노이징 코드를 다시 살펴보도록 하겠습니다.

DDPM Sampling Code

'denoise_add_noise' 함수는 예측된 노이즈(pred_noise)를 사용해 입력 이미지(x)에서 노이즈를 제거하는 함수입니다. 또한 이 과정에서 임의의 노이즈 z를 추가함으써 중간 샘플을 추가하고, 이를 통해 생성되는 결과물의 다양성을 확보하기도 합니다.

그리고 'sample_ddpm' 함수는 ddpm을 사용해 주어진 수의 샘플을 생성합니다. 구체적으로 표준 정규분포에서 초기 샘플을 무작위로 생성하고, 각 타임스텝마다 노이즈를 예측해 샘플에서 노이즈를 제거합니다. 이러한 노이즈 제거 과정을 지정된 횟수만큼 반복한 뒤 최종적인 결과물을 반환하는 식입니다.

|DDIM Sampling

DDIM Sampling Code

'denoise_ddim' 함수는 DDIM 알고리즘을 이용해 입력 이미지(x)로부터 노이즈를 제거하는 함수입니다. t시점의 누적 알파값과 t-1 시점의 누적 알파값을 가져와 노이즈의 방향성을 파악함으로써, 현재 시간 스텝(t)에 대한 정보만을 통해 노이즈를 제거했던 DDPM과 달리 이전 시간 스텝(t_prev)을 고려함으로써 얻게된 방향성 노이즈를 사용해 디노이징 과정을 진행합니다.

'sample_ddim' 함수는 DDIM을 사용해 주어진 샘플링 횟수를 반복합니다. 이처럼 동작 방식 자체는 DDPM과 동일하며, 사용하는 모델도 동일합니다. 하지만 DDPM은 각 타임스텝마다 노이즈를 제거하고, 새로운 임의의 노이즈를 추가하며 타임스텝을 1씩 감소시키는 방식이지만, DDIM은 예측된 초기 이미지와 방향성 노이즈를 사용해 샘플을 생성하며, 더 큰 스텝 크기를 사용해 빠르게 샘플링하는 것입니다.

|DDIM Result

위까지의 코드를 통해 생성된 결과는 어떨까요? DDPM의 경우, 영웅 캐릭터와 아이템은 제법 괜찮게 생성되었고, 몬스터의 경우 아직 부족한 모습을 보였었습니다.

DDIM Result

25개의 타임스텝을 하나의 단계로 해서 진행하기 때문에 최종적으로 생성된 결과가 순식간에 지나가 확인하기 어려울 수 있는데요, 최종 생성된 결과를 이미지로 살펴보면 아래와 같습니다.

DDIM Result

DDPM과 유사하게 캐릭터와 아이템은 잘 생성하고, 몬스터는 비슷하게 다소 낮은 퀄리티로 생성된 것을 확인할 수 있습니다. 즉, DDPM보다 25배 적은 단계로 구성된 forward process로도 충분히 비슷한 퀄리티의 결과물을 생성해낼 수 있음을 의미합니다.

이를 직접 확인하고 싶으시다면 제 코드를 통해 빠르게 확인할 수 있습니다. DDPM 모델은 사이즈가 작아 github에도 올려두었기 때문에 학습시킨 모델을 그대로 사용하거나, 자체적으로 학습시켜 사용하기에도 큰 부담이 없을 것입니다.

|DDIM vs DDPM speed

마지막으로 DDPM과 DDIM의 생성속도 차이를 비교해보겠습니다. 사용한 모델과 코드는 이전과 동일합니다.

ddpm vs ddim 속도비교

5회씩 10회 반복함으로써 평균적인 속도를 측정하였습니다. 측정 결과 DDIM은 실행에 약 77.6ms의 시간이 필요했고, DDPM은 1,530ms의 시간이 필요했습니다. 간단히 말하자면 DDIM이 DDPM보다 약 18배 빨리 실행된 것을 확인할 수 있는 셈입니다.

|Reference

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

https://arxiv.org/abs/2208.11970

[2] Jonathan Ho, Ajay Jain, Pieter Abbeel. Denoising Diffusion Probabilistic Models. https://arxiv.org/abs/2006.11239

[3] Jiaming Song, Chenlin Meng, Stefano Ermon. Denoising Diffusion Implicit Models. https://arxiv.org/abs/2010.02502

[4] metamath. ml-simple-works/diffusion. https://github.com/metamath1/ml-simple-works/tree/master/diffusion

[5] DLAI. How Diffusion Models Work. https://learn.deeplearning.ai/courses/diffusion-models/lesson/1/introduction