참고1: Annotated Diffusion
참고2: Diffusion Model 과 DDPM 수식 유도 과정

---

이론

• Diffusion이 풀고자 하는 문제는
  forward process와 reverse process를 거친 'output 이미지 $\bold{x_0}$'를 샘플링하는 확률분포 $\bold{p(x_0)}$를
  'input 이미지 $\bold{x_0}$'를 샘플링하는 확률 분포 $\bold{q(x_0)}$와 유사하게 만드는 것이다.

• 직관적으로 생각해보면,
  Forward process를 거치면서
  input 이미지 $\bold{x_0}$에 노이즈를 time step $\bold{T}$만큼 추가해 pure Gaussian noise $\bold{x_T}$를 얻은 후에
  $\bold{q(x_t|x_{t-1})=N(x_t;\sqrt{1-\beta_t}x_{t-1}, \beta_tI)}$
  $\\$
  다시 Reverse process를 거치면서
  $\bold{p(x_{t-1}|x_t)}$를 이용해 Gaussian noise $\bold{x_T}$로부터 $\bold{x_0}$를 샘플링하면 될 것이다.
  

• 하지만, 우리는 확률 분포 $\bold{p}$를 알지 못한다.
  그래서 우리는 신경망을 이용해서, 확률분포 $\bold{p}$를 가장 잘 모델링하는, 확률분포 $\bold{p_{\theta}}$를 예측하고 학습한다.
  

• 즉, 우리가 예측하는 확률분포 $\bold{p_{\theta}}$에서 $\bold{x_0}$ 가 샘플링될 확률이 가장 높게
  = Likelihood $\bold{p_{\theta}(x_0)}$ 가 가장 크도록
  = 확률분포 $\bold{p_{\theta}}$가 $\bold{p}$에 가까워지도록
  신경망을 훈련하면 되겠다!

• $\bold{p_{\theta}(x_0)}$를 최대화 하는 parameter $\theta$는
  $log\bold{p_{\theta}(x_0)}$를 최대화 하는 parameter $\theta$와 동일하다. 수식 전개의 편리함을 위해 $log\bold{p_{\theta}(x_0)}$를 사용한다.
  $\\$
  $log\bold{p_{\theta}(x_0)}$를 최대화 하는 수식을 전개하다보면,
  $\bold{p(x_{t-1}|x_t, x_0)}$과 $\bold{p_{\theta}(x_{t-1}|x_t)}$ 간의 KL-Divergence를 최소화하는 식으로 전개된다.
  $\\$
  이는 두 분포 간의 평균의 차이를 최소화하는 식으로 전개되며,
  $\bold{||\tilde{\mu_t}{(x_t, x_0)} - \mu_{\theta}(x_t, t)||^2}$
  $\\$
  위 식을 수학적인 트릭(Reparameterization trick)을 사용하면 (=수식적으로 전개하다보면)
  결국 time step $t$ 시점에 실제로 추가된 노이즈 $\epsilon$과 신경망이 예측한 $t$ 시점의 노이즈 $\epsilon_{\theta}$ 간의 차이를 최소화하는 식으로 전개된다.
  $||\epsilon - \epsilon_{\theta}(\bold{x_t, t})||^2 = ||\epsilon - \epsilon_{\theta}(\bold{\sqrt{\bar{\alpha_t}}x_0 + \sqrt{1-\bar{\alpha_t}}\epsilon,t})||^2$

• 정리하면, 확률분포 $\bold{p_{\theta}}$가 $\bold{p}$에 가까워지도록 신경망을 학습시키기 위해선
  확률분포 $\bold{p_{\theta}}$와 확률분포 $\bold{p}$간의 loss 함수를 정의해야 한다.

  • Diffusion에서는 확률분포 $\bold{p}$와 확률분포 $\bold{p_{\theta}}$가 정규 분포(gaussian distribution)을 따른다고 설정한다.
    정규분포는 두 가지 정보, 평균($\mu$)과 분산($\sigma$)으로 표현된다.

  • 그렇다면, "확률분포 $\bold{p}$ 의 평균($\bold{\tilde{\mu}}$), 분산"과
    "확률 분포 $\bold{p_{\theta}}$의 평균($\bold{\mu_{\theta}}$), 분산"의 차이를 줄이는 loss 함수를 만들면 되겠다.

  • DDPM에서는 분산은 고정된(fixed) 값으로 두고 평균의 차이를 줄이는 방식을 사용했다.
    이후에 분산까지 학습하는 연구가 나오긴 했으나, 지금은 DDPM을 기준으로 생각하자.

  • 수학적인 트릭(Reparameterization trick)을 사용하면
    (=수식적으로 전개하다보면)
    "확률분포 $\bold{p}$ 의 평균($\bold{\tilde{\mu}}$)"과, "확률 분포 $\bold{p_{\theta}}$의 평균($\bold{\mu_{\theta}}$)"의 차이를 줄이는 loss function 수식을
    "이미지에 실제로 추가된 노이즈($\epsilon$)"과 "신경망이 예측한 노이즈($\epsilon_{\theta}$)"의 차이를 줄이는 loss function 수식으로 전개가 가능하다.

  • 따라서, 신경망은 이미지에 추가된 노이즈를 예측하도록 설계한다.

• Training 과정에서,
  실제 신경망 학습이 어떻게 진행되는지 알고리즘을 정리해보면,
  
  1) real unknown and possibly complex data distribution $\bold{q(x_0)}$로부터
  random sample $\bold{x_0}$를 뽑는다.
  2) $\bold{1}$과 $\bold{T}$ 사이의 random한 noise level (= random time step) $\bold{t}$를 뽑는다.
  3) Gaussian distribution으로부터 noise $\bold{\epsilon}$을 뽑는다.
  4) input image $\bold{x_0}$에 noise level $\bold{t}$만큼 노이즈를 추가한다.
  (known scheldue $\beta_t$ 만큼 noise가 추가된다.)
  5) 신경망은 노이징된 image $\bold{x_t}$에 추가된 noise를 예측하도록 training한다.
  (실제 구현에선, 데이터의 여러 batch들에서 위 과정이 이루어진다. 그리고 stochastic gradient descent를 사용해 신경망을 최적화한다.)

• Inference 과정에서는,
  $\bold{p_\theta}$를 학습한 신경망으로 pure Gaussian noise $\bold{x_T}$로부터 denoising해서 output image $\bold{x_0}$를 sampling한다.
  
  1) time step $T$의 pure noise에서 Gaussian distribution을 sampling한다.
  $\bold{x_t} \sim N(\bold{0,I})$
  2) Denoising
  time step $\bold{T}$부터 0까지, 신경망을 사용해서, 신경망이 학습한 conditional probability를 이용해 점진적으로 denoise한다.
  mean을 reparameterization하여 만든 우리의 noise predictor를 이용해서,
  우리는 조금 더 denoised 된 image $\bold{x_{t−1}}$ 을 얻을 수 있다.
  이 과정을 통해 real data distribution $\bold{q(x_0)}$에서 샘플링한 이미지와 유사한 새로운 image $\bold{x_0} \sim \bold{p_{\theta}}(\bold{x_0})$를 얻게된다.

---

구현

가장 간단한 버전으로 Diffusion을 구현한 코드를 요약하여, 핵심적인 구현 흐름을 확인해보자.

1. 신경망으로 U-Net을 사용한다.

2. Forward Process를 정의한다.

2-1) variance schedule: $\beta_{t}$ scheduling

def linear_beta_schedule(timesteps):
    beta_start = 0.0001
    beta_end = 0.02
    return torch.linspace(beta_start, beta_end, timesteps)

2-2) noising (at time step $\bold{t}$)
$||\epsilon-\epsilon_{\theta}(\bold{x_t, t})||^2=||\epsilon-\epsilon_{\theta}(\sqrt{\bar{\alpha_t}}\bold{x_0}+\sqrt{1-\bar{\alpha_t}}\epsilon, t)||^2$

def q_sample(x_start, t, noise=None):
	if noise is None:
    	noise = torch.randn_like(x_start)

	sqrt_alphas_cumprod_t = extract(sqrt_alphas_cumprod, t, x_start.shape)
    sqrt_one_minus_alphas_cumprod_t = extract(
    	sqert_one_minus_alphas_cumprod, t, x_start.shape
    )
    
    return sqrt_alphas_cumpord_t * x_start + sqrt_one_minus_alphas_cumprod_t * noise

2-3) Loss Function을 정의한다.

def p_losses(denoise_model, x_start, t, noise=None, loss_type="l1"):

	if noise is None:
    	noise = torch.randn_like(x_start)

	# noised image
	x_noisy = q_sample(x_start=x_start, t=t, noise=noise)

	# predicted noise. 여기서 denoise_model은 U-Net
	predicted_noise = denoise_model(x_noisy, t)

	# l1_loss를 사용할 경우
	loss = F.l1_loss(noise, predicted_noise)

return loss

3. Reverse Process를 정의한다.

1) time step $T$의 pure noise에서 Gaussian distribution을 sampling한다.
$\bold{x_t} \sim N(\bold{0,I})$

2) denoising
time step T부터 0까지, 신경망을 사용해서, 신경망이 학습한 conditional probability를 이용해 점진적으로 denoise한다.
mean을reparameterization하여 만든 우리의 noise predictor를 이용해서,
우리는 조금 더 denoised 된 image $\bold{x_{t−1}}$ 을 얻을 수 있다.
이 과정을 통해 real data distribution $\bold{q(x_0)}$에서 생성된 것과 유사한 새로운 image를 얻게된다.

@torch.no_grad()
def p_sample(model, x, t, t_index):
	betas_t = extract(betas, t, x.shape)
    sqrt_one_minus_alphas_cumprod_t = extract(
    	sqrt_one_minus_alphas_cumprod, t, x.shape
    )
    
    # 1/sqrt(\alpha_t)
    sqrt_recip_alphas_t = extract(sqrt_recip_alphas, t, x.shape)
    
    # Equation 11 in the paper
    # Use our model (noise predictor) to predict the mean
    model_mean = sqrt_recip_alphas_t * (
    	x - betas_t * model(x, t) / sqrt_one_minus_alphas_cumprod_t
    )
    
    if t_index == 0:
    	return model_mean
	else:
    	posterior_variance_t = extract(
        	posterior_variance, t, x.shape
        )
        noise = torch.randn_like(x)
        # Algorithm2 line 4:
        return model_mean + torch.sqrt(posterior_variance_t) * noise

4. Model Training

epochs = 6

for epoch in range(epochs):
	for step, batch in enumerate(dataloader):
    	optimizer.zero_grad()

        # Algorithm 1 line 3: sample t uniformally for every example in the batch
        t = torch.randint(0, timesteps, (batch_size,), device=device).long()

		loss = p_losses(denoise_model=model, x_start=batch, t=t, loss_type="huber")
        
        if step % 100 == 0:
        	print("Loss:", loss.item())
		
        loss.backward()
        optimizer.step()

5. Sampling(Inference)

# Algorithm 2 (includint returning all images)
@torch.no_grad()
def p_sample_loop(model, shape):
	device = next(model.parameters()).device
    
    b = shape[0]
    # start from pure noise (for each example in the batch)
    img = torch.randn(shape, device=device)
    imgs = []
    
    for i in tqdm(reversed(range(0, timesteps)), desc = 'sampling loop time step', total=timesteps):
    	img = p_sample(model, img, torch.full((b,), i, device=device, dtype=torch.long), i)
    	imgs.append(img.cpu().numpy())
	return imgs

@torch.no_grad()
def sample(model, image_size, batch_size=16, channels=3):
	return p_sample_loop(model, shape=(batch_size, channels, image_size, image_size))

결과 시각화

# show a random one
random_index = 5
plt.imshow(samples[-1][random_index].reshape(image_size, image_size, channels), cmap="gray")