DDPM(Denosing Diffusion Probabilistic Model) 개념 정리

<요약>

Diffusion Model 개요

Diffusion Model 알고리즘

Diffusion Model Loss 

Experiments

 

---

Diffusion Model 개요

- GAN, VAE 와 같은 생성 모델(Generative Model) 중 하나로써, 2022년에 이슈가 되었던 text-to-image 모델인 Stable-Diffusion, DALL-E-2, Imagen의 기반이 되는 모델

- 입력 이미지에 (정규 분포를 가진) Noise를 여러 단계에 걸쳐 추가하고, 여러 단계에 걸쳐 (정규 분포를 가진) Noise를 제거함으로써, 입력 이미지와 유사한 확률 분포를 가진 결과 이미지를 생성하는 모델

- Forward Diffusion Process에서는 이미지에 고정된(fixed) 정규 분포(=Gaussian분포)로 생성된 Noise가 더해지고,

- Reverse Diffusion Process에서는 이미지를 학습된(learned) 정규 분포로 생성된 Noise를 이미지에서 뺍니다.

그림과 달리 실제 좌측 두개 이미지는 같지 않을 수 있다.

 

- Diffusion Model이 풀려고 하는 문제는, Forward -> Reverse 단계를 거친 '결과 이미지'를 '입력 이미지'의 확률 분포와 유사하게 만드는 것 입니다. 이를 위해 Reverse단계에서, Noise 생성 확률 분포 Parameter인 평균과 표준편차를 업데이트하며 학습이 진행

-  1번째 그림은 최초 Diffusion Model(2015년)에 해당되며, Maximum Log-likelihood Estimation(MLE)을 사용했습니다.

- 2번째 그림은 후속 논문인 DDPM(2020년, Denosing Diffusion Probabilistic Model)에 해당되며, Reverse Diffusion 부분을 CNN계열의 UNet모델을 사용했습니다.

 

Diffusion Model 알고리즘

기본적으로 시간 단계 𝑡의 각각의 새로운 이미지는 조건부 가우스 분포에서 추출됩니다.

 

- 위의 Forward Process 수식은 아래와 같이 x0 수식으로 변환될 수 있다

Diffusion Kernel, 본 식의 유도는 Reparameterization Trick 참고

 

- 본 식은 현재 상태(t)는 이전 상태(t-1)에 의존한다는 마르코프(Markov) 특성을 가정한다. 

- 이전 상태(t-1)가 주어질 때 현재 상태(t)가 될 확률분포q는 평균(μ)과 분산(Σ)으로 구성된 Gaussian Distribution(=Normal Distribution)을 따른다. (이전 상태에서 현재 상태가 될 확률은 조건부 가우시안 분포를 따른다)

- 시간 단계 t의 각각의 새로운(약간 노이즈가 많은) 이미지는 조건부 가우스 분포에서 추출됩니다 (Annotated Diffusion 참고)

- 이 식을 Noise 크기 parameter인 β를 포함하여 다르게 표현하면, 이전 상태(t-1)의 이미지가 β만큼 다른 pixel을 선택하게 하고, root(1-β)만큼 이전 pixel값을 선택하게 하는 식으로 정의 할 수 있습니다. t가 점점 증가함에 따라 β가 커짐으로써 Noise가 강해집니다.

- 𝛽𝑡는 각 시간 단계 𝑡(따라서 아래 첨자)에서 일정하지 않습니다. --- 실제로 선형, 2차, 코사인 등이 될 수 있는 소위 "Variance Schedule(분산 일정)"을 정의합니다. (Learning Rate Schedule 과 약간 비슷함)  (Annotated Diffusion 참고)

 

==> 

(동훈) 이전 이미지에서 현재 이미지로 넘어오는 확률은 이전 이미지(X_t-1과 특정 스케줄링 𝛽𝑡 를 조건으로 하는 가우시안 분포를 따른다고 가정한다. 

 

<Reparameterization Trick 상세>

 

 register_buffer('sqrt_alphas_cumprod', torch.sqrt(alphas_cumprod))

 register_buffer('sqrt_one_minus_alphas_cumprod', torch.sqrt(1. - alphas_cumprod))

<Gaussian Diffusion 클래스>

class GaussianDiffusion(Module):
	def __init__(
    	self,
        model,
        *,
        image_size,
        timesteps = 1000,
        sampling_timesteps = None,
        objective = 'pred_v',
        beta_schedule = 'sigmoid',
        schedule_fn_kwargs = dict(),
        ddim_sampling_eta = 0.,
        auto_normalize = True,
        offset_noise_strength = 0.,  # https://www.crosslabs.org/blog/diffusion-with-offset-noise
        min_snr_loss_weight = False, # https://arxiv.org/abs/2303.09556
        min_snr_gamma = 5
        ):
        	super().__init__()
            '''
            type(self) == GaussianDiffusion: 현재 객체 (self)의 클래스가 GaussianDiffusion인지 확인합니다.
			model.channels != model.out_dim: model 객체의 channels 속성과 out_dim 속성이 같지 않은지 확인합니다.
			이 두 조건이 동시에 참일 경우 assert 문은 실패하게 되고, AssertionError가 발생합니다. 
            이를 통해 GaussianDiffusion 클래스의 인스턴스에서는 model.channels와 model.out_dim이 같아야 함을 보장합니다.
       		
            하나만 False면 된다. 
            GaussianDiffusion 클래스라면, model.channles와 model.out_dim이 동일해야 하며
            GaussianDiffusion 클래스 아니라면, 무관이다.
            '''
            assert not (type(self) == GaussianDiffusion and model.channels != model.out_dim)
        	
            '''
            not A or not B 은 not (A and B) 입니다.
            좌, 우변 중 하나만 False 이면 된다.
            
            model 객체는 random_or_learned_sinusoidal_cond 속성을 가지고 있지 않거나, 가지고 있더라도 그 값이 False이어야 합니다.
            '''
            assert not hasattr(model, 'random_or_learned_sinusoidal_cond') or not model.random_or_learned_sinusoidal_cond
            
            self.model = model
            
            self.channels = self.model.channels
            self.self_condition = self.model.self_condition
            
            if isinstance(image_size, int):
            	image_size = (image_size, image_size)
            
            # 이미지 사이즈는 예를 들어 (256, 256) 이거나 [256, 256] 이어야 한다.
            assert isinstance(image_size, (tuple, list)) and len(image_size) == 2, 'image size must be a integer or a tuple/list of two integers'
        	self.image_size = image_size
            
            self.objective = objective
            
            # 'objective는 pred_noise(노이즈 예측) 또는 pred_x0(이미지 시작 예측) 또는 
            # pred_v(v 예측[progressive distillation paper 부록 D에 정의된 v-매개변수화, imagen-video에서 성공적으로 사용됨]) 중 하나여야 합니다.'
            assert objective in {'pred_noise', 'pred_x0', 'pred_v'}, 'objective must be either pred_noise (predict noise) or pred_x0 (predict image start) or pred_v (predict v [v-parameterization as defined in appendix D of progressive distillation paper, used in imagen-video successfully])'
            
            if beta_schedule == 'linear':
            beta_schedule_fn = linear_beta_schedule
            elif beta_schedule == 'cosine':
                beta_schedule_fn = cosine_beta_schedule
            elif beta_schedule == 'sigmoid': # 디폴트
                beta_schedule_fn = sigmoid_beta_schedule
            else:
                raise ValueError(f'unknown beta schedule {beta_schedule}')
                
            betas = beta_schedule_fn(timesteps, ***schedule_fn_kwargs)
            
            alphas = 1. - betas
            alphas_cumprod = torch.cumprod(alphas, dim=0)
            # alphas_cumprod 텐서의 마지막 원소를 제외한 부분에 1.0을 앞에 추가하여 alphas_cumprod_prev 텐서를 생성하는 것
            # 예를 들어, 누적곱 계산에서 첫 번째 단계의 이전 값이 필요할 때 1.0을 추가하여 이러한 요구를 충족시킬 수 있습니다
            # (1,0), value=1. 의 의미는 앞에 1개, 뒤에 0개 1.0을 추가한다는 의미이다
            alphas_cumprod_prev = F.pad(alphas_cumprod[:-1], (1, 0), value = 1.)
            
            timesteps, = betas.shape # 1000
            self.num_timesteps = int(timesteps)
            
            # 샘플링 관련 파라미터
            self.sampling_timesteps = default(sampling_timesteps, timesteps)
            
            assert self.sampling_timesteps <= timesteps # 1000 <= 1000
            self.is_ddim_sampling = self.sampling_timesteps < timesteps
            self.ddim_sampling_eta = ddim_sampling_eta
            
            # helper function to register buffer from float64 to float32
        	# register_buffer는 lambda 함수는 val를 torch.float32로 변환한다.
            register_buffer = lambda name, val: self.register_buffer(name, val.to(torch.float32))
            
            register_buffer('betas', betas)
        	register_buffer('alphas_cumprod', alphas_cumprod)
        	register_buffer('alphas_cumprod_prev', alphas_cumprod_prev)

 

 

- Reverse과정에서 마지막 상태( $X_T$ )를 반대로 최초 상태( $X_0$ )로 만들고자 합니다.

- (위에서 빨간색으로 표시한) q식의 좌측항(posterior)과 우측항(prior)을 바꾼 조건부 확률 q(xt-1 | xt)을 알 수 있다면, 최초 상태( $X_0$ )로 만들 수 있겠지만 이를 계산 할 수가 없습니다
(각 t시점에서 이미지의 확률 분포q(xt)를 알 수 없기 때문에 베이즈 정리에 의해 계산되지 않습니다).

- 때문에 확률분포 q가 주어졌을 때, 이 확률 분포를 가장 잘 모델링하는 확률 $p_θ$를 찾는 문제로 변환 합니다.

(Maximum Likelihood Estimation 문제로 변환)

 

- 즉 확률분포 q에서 관측한 값으로 확률 $p_θ$ 의 likelihood를 구하였을 때, 그 likelihood값이 최대(Maximum)가 되는 확률분포를 찾는 Maximum Likelihood Estimation 문제

 

- E는 확률 분포 q로 샘플링한 관측값에 대한 평균값(Expectation=기대값=적분값)을 의미합니다. 위 식을 전개 할 경우 아래 수식이 되며, 최초 Diffusion 논문(2015)에서는 Pθ가 가우시안 분포를 따른다고 가정하고 아래 수식이 0이 되도록 확률분포 parameter(θ)를 업데이트

 

- -log를 붙이면 Negative Likelihood가 되고, 이를 최대하 아닌 최소화 하는 식으로 Loss 설계

 

==> 확률 분포 q가 주어졌을 때 확률 분포를 가장 잘 설명하는 확률 $p_θ$를 찾는 것이 목표인데, 이 확률 $p_θ$가 가우시안 분포를 따른다고 가정한다. 

 

모델이 실제 데이터 분포를 잘 모방할수록, 생성된 샘플의 질도 높아집니다. 따라서, "확률 분포를 가장 잘 모델링하는 확률 분포를 찾는다"는 것은, 모델이 실제 데이터 분포와 최대한 유사한 분포를 학습하여, 실제와 구분하기 어려운 고품질의 새로운 샘플을 생성할 수 있도록 하는 것을 목표로 합니다.

 

위 식을 다시 한번 더 정제한 DDPM(2020)에서는 아래 수식을 통해 Loss를 설계하여 UNet을 업데이트 하게 됩니다.

 

수식 유도 과정 참고: [개념 정리] Diffusion Model 과 DDPM 수식 유도 과정 (tistory.com)

 

Diffusion Model Loss 

- 아래 수식은 최초 Diffusion Model(2015) 논문에서 언급된 보기 어려운 Loss 수식을 DDPM(2020)에서 깔끔하게 재정리한 수식

- 확률분포 q를 통해 샘플링 하였을 때, 위 수식의 평균값(=기대값=적분값)을 계산하는 수식으로 Loss가 설계되었으며, 위에서 설명하였듯이 이 값을 최소화 해야 합니다. 

• LT : 마지막 상태(xT)에 대해, 확률분포 q와 p의 KL Divergence를 최소화하여, p와 q 확률분포 차이를 줄일 수 있습니다. Regularization Loss이라 불립니다.
• Lt-1 : 현재 상태(xt)가 주어질 때, 이전 상태(t-1)가 나올 확률 분포 q와 p의 KL Divergence를 최소화하여, p와 q 확률분포 차이를 줄일 수 있습니다. Denoising Process Loss라 불립니다.
• L0 : VAE(Varation Auto Encoder)Loss를 구성하는 Reconstruction Loss식 $ -log(p(x_0|x_T)) $와 대응됩니다. 확률분포 q를 통해 샘플링하였을 때 - log( p(x0|x1) ) 에 대한 기대값을 최소화하여, p와 q 확률분포 차이를 줄일 수 있습니다.

 

-  최초 상태(x0)로 부터 Step t가 경과되고, $ϵ_θ$ 네트워크가 Normalized Gaussian 분포 ϵ 를 따르도록, 기대값을 계산하는 수식이 Loss로 설계되었습니다. 이 값을 최소화 해야 합니다.

---

<Forward Process 수식>

위의 Forward Process 수식은 아래와 같이 x0 수식으로 변환될 수 있다.

 

 

X가 Normal(=Gaussian) Distribution을 따를 때, 평균(μ) + 표준편차(σ) * Normalized 가우시안 분포(ϵ)의 식으로 나타낼 수 있는 아래 Reparameterization Trick을 사용해서, q(xt | x0)를 위 ϵθ 함수의 첫번째 인자로 바꿀 수 있습니다.

 

 

이제 DDPM(2020) Loss 수식의 의미를 보겠습니다. 앞에서 negative log-likelihood를 설명 할 때, 확률분포 q에서 관측한 값으로, 확률분포 pθ의 likelihood를 구하였을 때, 그 likelihood를 최대화시키는 확률 분포를 찾는 문제라고 언급했었습니다.

 

- 여기서는 t, x0, ϵ로 구성된 확률 분포에서 관측한 값으로, 확률분포를 모델링한 네트워크 ϵθ의 likelihood를 구하였을 때, 그 likelihood를 최대화시키는 확률 분포를 찾는 문제가 됩니다.  

 

-log가 아닌, 네트워크 ϵθ와 Normalized Gaussian분포 ϵ가 같아지게하는 Squared Error에 대한 수식으로 바뀌었고
네트워크 ϵθ의 입력은 t, x0, ϵ으로 구성된 확률 분포를 Reparameterization Trick을 통해, 계산 가능한 수식으로 정의되었습니다.

 

<전체 수도 코드>

 

- Forward(Training) 프로세스와 Reverse(Sampling) 프로세스를 거친다.

- Forward(Training) 프로세스: 네트워크ϵθ 학습이 목표 . X0(=입력 이미지)는 forward process에 의해 XT(=노이즈 이미지)를 항상 일정하게 만들게 됩니다.

- Reverse(Sampling) 프로세스:  XT가 학습된 네트워크ϵθ를 통해 X0(=생성 이미지)를 만듭니다

 

 

register_buffer('sqrt_alphas_cumprod', torch.sqrt(alphas_cumprod))

 

 

---

Experiments

- βt 값에 의해 스케쥴링이 가능합니다. 최초 Diffusion 논문에서는 아래식을 사용했습니다.

- DDPM에서는 constant, linear, quadratic schedules을 실험했고 다른 결과를 보였습니다. β1 = 0.0001으로 βT = 0.02로 두고, linear schedule을 사용했습니다.

 

 

 

- T는 1000을 사용하였습니다.

- 최초 Diffusion 연구에서는 CIFAR-10, bark, dead leaves와 같은 natural image와 swis roll, binary sequence, MNIST 숫자를 사용하였고, DDPM에서는 CIFAR10만 실험했습니다.

아래 DDPM의 Sample Quality결과를 표시하였습니다. CIFAR10 데이터셋을 갖고 Inception Score, FID Score, Negative Log Likelihood(NLL)로 측정하였고, DDPM이 가장 좋은 성능을 보였습니다.

 

 

 

 

출처: [개념 정리] Diffusion Model 과 DDPM 수식 유도 과정 (tistory.com)

[개념 정리] Diffusion Model 과 DDPM 수식 유도 과정

 

---

(코드 맛보기)

 

---