Diffusion은 GAN과 함께 10년 넘게 Image Generation방법 중 영향력이 큰 생성모델링 기법이다. (GAN의 최대 유행: 2017~2020) Diffusion모델은 여러 Benchmark Dataset에서 기존 최신 GAN보다 성능이 뛰어난데, 특히 txt2img의 OpenAI DALL∙E, Google Imagen등에서 많이 사용된다.
Diffusion모델의 핵심기반 아이디어는 VAE, EBM과 비슷한 점이 존재한다. Contrastive Diffusion(= Score 함수) 사용 대신, Log분포의 Gradient를 직접 추정해 모델을 학습하는 EBM분야인 "Score-based Generative Model"에서 매우 중요하다.
NCSN(Noise Conditional Score Network)라는 모델이 원시 data에 여러 scale의 noise를 적용, Data밀도가 낮은 영역에서도 잘 작동하였다.
img x0를 많은 step(ex. T=1000)동안 점진적으로 손상시켜 최종적으로 "Standard Gaussian Noise"와 구별불가능하게 만든다 가정하자.
img xt-1에 분산 βt를 갖는 소량의 Gaussian Noise를 추가. 새로운 img xt를 생성하는 함수 q를 정의하자. 이 q를 적용하면, 아래처럼 점진적으로 잡음이 커지는 img sequence (x0, ... , xT)를 생성할 수 있다.
이때, update과정은 아래와 같이 수학적으로 표기가능하다.
Reparameterization Trick
🤔 q를 t번 적용하지 않고 이미지 x0+Noise인 xt로 바로 Skip가능하다면 유용하지 않을까?
위 식의 2번째 줄은 2개의 Gaussian분포를 더해 새로운 Gaussian분포를 하나 얻을 수 있다:
따라서 원본 img x0에서 정방향 diffusion과정의 어느단계로든 건너뛸 수 있게 되었다. 또한, 기존의 βt대신, āt값을 사용해 diffusion schedule을 정의할 수 있다. āt: signal(= 원본 x0)로 인한 분산 1-āt: noise(ε)로 인한 분산
∴ Foward Diffusion과정 q는 아래와 같다:
Diffusion Schedule
추가적으로 각 time step마다 다른 β를 자유로이 선택가능하다. 즉, β 나 ā 값이 t에 따라 변하는 방식을 "Diffusion Schedule"이라 한다.
[Ho et al;2020]에서는 Linear Diffusion Schedule을 선택했다. 위 논문에서는 β1=0.0001부터 ,βT=0.02까지 선형적으로 증가한다.
이후 논문에서는 Cosine Diffusion Schedule이 도입되었다. 이때, 코사인 스케줄은 ā를 아래와 같이 정의한다:
위 그림을 보면 Cosine Diffusion Schedule이 더 느리게 상승했음을 확인가능하다. 즉, img에 noise를 점진적으로 추가해 train효율성 및 생성품질을 향상한다.
Reverse Diffusion
noise추가과정을 되돌리는 신경망 pθ(xt-1|xt)를 구축하자. 이때, 신경망 pθ(xt-1|xt)는 q(xt-1|xt)의 역방향 분포를 근사화하는 신경망이다.
이를 통해 N(0, I)에서 random noise sampling 후, reverse diffusion과정을 여러번 적용해 새로운 img를생성할 수 있다:
Train Algorithm
[Reverse Diffusion과 VAE간 비교] ∙ 목표: 신경망으로 random noise를 의미있는 출력으로 변환하는 것. ∙ 차이점: - VAE는 정방향(= img2noise)과정이 모델의 일부 (= 학습됨) - Diffusion: 이에 대한 parameter가 없이 진행.
그렇기에, VAE와 동일한 Loss function을 적용한다.
⭐️ 주의점: Diffusion모델이 실제로 2개의 신경망복사본을 유지한다는 점. 경사하강법으로 train된 신경망과 이전 train step의 신경망 가중치의 EMA(지수이동평균)을 사용하는 또다른 EMA신경망이다.
2. DDPM with U-Net
DDPM with UNet
앞서 신경망의 종류를 확인했다: img에 추가된 noise를 예측하는 신경망 이제, 이에 사용할 신경망 구조를 살펴보자. Skip Connection을 통해 정보가 신경망 일부를 건너뛰고 후속층으로 signal을 흘러보낼 수 있다. 특히나 "출력이 입력크기와 같아야할 때, U-Net이 유용하다".
신경망 후속층에서 사용가능하도록 noise분산값(= 스칼라값)을 더 복잡한 표현이 가능한 고차원 벡터로 변환한다. NeRF논문에서는 문장에서 단어의 이산적인 위치를 벡터로 encoding하는 것이 아닌, 연속적인 값으로 확장했다:
보통, noise embedding길이의 절반이 되게 L=16으로 선택하고 주파수 f의 최대 scaling계수로 ln(1000) / (L-1) 을 택한다.
Residual Block구조는 아래와 같다.
일부 Residual Block에서 Block의 출력과 channel수를 일치시켜야한다. 그렇기에 Skip Connection에 kernel_size=1인 Conv2D층을 추가한다.
Down Block & UpBlock
∙ Down Block
Residual Block으로 channel수를 늘린다. 또한, img_size를 줄이려고 마지막에 AvgPooling층을 적용한다. (UpBlock과의 Skip Connection을 위해 각 Residual Block에 list를 추가해야함.)
∙ UpBlock
UpSampling2D를 진행한다. (보통 ConvTranspose나 Interpolation을 적용.) 연속된 UpBlock은 channel수를 줄이면서 DownBlock의 출력과 연결한다.
import torch
import torch.nn as nn
import torch.nn.functional as F
class SwishActivation(nn.Module):
def forward(self, x):
return x * torch.sigmoid(x)
class ResidualBlock(nn.Module):
def __init__(self, width):
super(ResidualBlock, self).__init__()
self.width = width
self.conv1 = nn.Conv2d(in_channels=width, out_channels=width, kernel_size=3, padding=1)
self.bn1 = nn.BatchNorm2d(width)
self.activation = SwishActivation()
self.conv2 = nn.Conv2d(in_channels=width, out_channels=width, kernel_size=3, padding=1)
self.bn2 = nn.BatchNorm2d(width)
def forward(self, x):
input_width = x.size(1)
if input_width == self.width:
residual = x
else:
residual = nn.Conv2d(in_channels=input_width, out_channels=self.width, kernel_size=1)(x)
x = self.bn1(x)
x = self.activation(x)
x = self.conv1(x)
x = self.bn2(x)
x = self.activation(x)
x = self.conv2(x)
x = x + residual
return x
class DownBlock(nn.Module):
def __init__(self, width, block_depth):
super(DownBlock, self).__init__()
self.width = width
self.block_depth = block_depth
self.res_blocks = nn.ModuleList([ResidualBlock(width) for _ in range(block_depth)])
self.avg_pool = nn.AvgPool2d(kernel_size=2)
def forward(self, x):
x, skips = x
for res_block in self.res_blocks:
x = res_block(x)
skips.append(x)
x = self.avg_pool(x)
return x
class UpBlock(nn.Module):
def __init__(self, width, block_depth):
super(UpBlock, self).__init__()
self.width = width
self.block_depth = block_depth
self.up_sampling = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False)
self.res_blocks = nn.ModuleList([ResidualBlock(width) for _ in range(block_depth)])
def forward(self, x):
x, skips = x
x = self.up_sampling(x)
for res_block in self.res_blocks:
x = torch.cat([x, skips.pop()], dim=1)
x = res_block(x)
return x
3. DDPM Sampling과 DDIM
DDPM을 이용한 Sampling
훈련된 모델에서 img를 sampling해야한다. 이를 위해 reverse diffusion과정이 필요하다. 즉, random noise에서 시작해 원본이 남을 때 까지 모델로 noise를 점진적으로 제거해야한다.
모델은 random noise추가과정의 마지막 time_step에서 추가된 noise뿐만아니라 trainset의 img에 추가된 noise총량을 예측하게 훈련된다. (다만, 완전한 random noise에서 한번에 img예측은 불가능하다.)
그렇기에 두 단계를 거쳐 xt에서 xt-1로 이동한다. Step 1. 모델의 noise예측을 사용, x0의 추정치 계산 Step 2. 예측 noise를 t-1 step까지만 다시 적용해 xt-1을 생성. 이 과정을 여러단계에 걸쳐 반복하면 조금씩 점진적으로 x0에 대한 추정치로 결국 돌아갈 수 있다.. 아래 식은 이 과정을 수학적으로 보여준다: 위 식을 해석하면 다음과 같다: 첫번째 괄호:신경망 ε(t)θ에 의해 예측된 noise로 계산된 추정이미지 x0 이후 t-1 signal비율인 √āt-1로 scale을 조정, 예측된 noise를 재적용.
두번째 괄호: t-1 noise 비율인 √1 - āt-1 - σ2t로 scale 조정.
세번째 괄호: 추가적인 Gaussian Noise를 더한다.
DDIM (Denoising Diffusion Implicit Model)
DDPM에서 모든 t에 대해 σt=0인 특수한 경우. 즉, DDIM을 사용하면 생성과정이 완전히 "결정론적(deterministic)"하다. 그렇기에 Random Noise가 같다면 항상 동일한 출력을 만든다.
이는 latent space의 sample과 pixel space에서 생성된 출력사이에 잘 정의된 mapping이 있음을 나타낸다.
4. 요약
DDPM이후 DDIM논문의 아이디어는 생성과정을 완전히 결정론적으로 만들었다.
DDPM은 2가지 과정으로 나뉜다.
① 정방향 Diffusion - 일련의 작은 단계로 train data에 noise를 추가. 이때, reparameterization trick으로 어느단계에 해당하는 noise_img라도 계산가능하다. parameter선택 schedule 선택 또한 중요하다.
② 역방향 Diffusion
- 추가된 noise를 예측하는 모델로 구성. noise_img와 해당단계 noise비율이 주어지면 각 time step에서 noise를 예측하는 U-Net으로 구현된다. UpBlock: img 크기를 늘리고 channel수를 줄임. DownBlock: img 크기를 줄이고 channel수를 늘림. noise비율은 사인파임베딩(sinusoidal embedding)으로 encoding
역방향 Diffusion 단계를 늘리면 속도는 느려지나 img생성품질은 향상된다. 또한, 두 img 사이를 보간하기 위해 latent space연산을 수행한다.
