만들면서 배우는 생성 AI 8장 : Diffusion 모델

이번 모델은 현재 스테이블 디퓨전을 포함해 대부분의 생성 AI 모델에서 가장 성능이 뛰어난 모델입니다. 앞서 EBM과 유사하게 다른 학문의 이론인 열역학의 확산에서 영감을 얻은 방식입니다. 사실 이 내용 자체는 이전 논문 리뷰 글 (링크) 에서 볼 수 있지만 책 정리 차원에서 오늘은 다시 정리해보려고 한다.

 

 

핵심 아이디어

실제 열이 확산되는 방식처럼 순차적으로 조금씩 원본 이미지에서 잡음을 추가해서 완전히 랜덤한 잡음으로 향하도록 하고, 동일하게 반대 방향으로 랜덤한 잡음에서 조금씩 잡음을 제거해서 원본 이미지로 향하도록 학습시키는 모델인데 이걸 실제로 해보니 잘 작동해서 오늘날 많이 사용하게 되었다.

 

 

추가적으로, 이 노이즈의 비율이 얼마인지, 실제로 넣어서 어떤 이미지가 나왔는지를 토대로 noise를 예측하는 신경망을 학습시킨다. 이를 통해 우리는 랜덤한 노이즈에서 신경망으로 예측된 노이즈를 계산해서 새로운 이미지를 만들어낸다. 

 

 

작동에 대한 개념적 이해

수학적 표현

 

위의 도식을 실제 식으로 표현하면 아래와 같다. 

 

여기서 beta는 일종의 ratio인데 분산을 일정하게 1로 만들기 위해 넣는 요소라고 보면 된다. noise는 평균이 0이고 단위 분산을 가지는 표준 가우스 분포이다.

 

이렇게 분산을 맞춰줌으로써 표준 가우스 분포에 가깝다는 것을 보장시킬 수 있고 이를 통해 추후에 역연산(랜덤한 잡음 -> 이미지)으로 갈 때에도 이를 적용할 수 있게 만들었다.

 

역방향 확산 과정

 

 

재매개변수화 트릭

• 우리가 원하는 건 x0 -> xt로 한번에 가는 것인데 사실 식이 복잡해진다. 이에 대해서는 글 (링크)에서 보다시피 몇가지 변수를 재 정의함으로써 하나의 식으로 한번에 나타낼 수 있다. 자세한 내용은 위의 링크에 있지만 책에 있는 내용도 다시 정리해보겠다.
  
  
• 여기서 beta를 아래와 같이 수정하면 식을 이렇게 전개해볼 수 있다.

1번 식

 

 

2번 식이 이해하기 헷갈리는데 이는 정규분포의 가법성에 따라 독립인 두 확률변수에 대해서 아래와 같은 식을 대입할 수 있어서입니다.

 

 

위 식을 고려해서 다시 식을 전개해보면 두개의 가우스 분포를 더하여 새로운 가우스 분포 하나를 얻을 수 있다는 사실을 이용한다.

 

 

모델 소개

1. 정방향 노이즈 확산

 

• 앞서 말했듯이 원본 이미지에 노이즈를 더하는 작업인데, 이 단계에서는 얼만큼 노이즈를 더해줄지만 정해주면 되는데 이걸 고정적으로 더해주는게 아니라 변동적으로 더해줄 수 있다.

 

확산 스케줄

 

앞서 우리는 베타를 통해서 얼만큼 확산시켜줄 것인지를 조정해주게 된다. 이 베타는 일정하게 고정되는 것이 아니라 t에 따라서 조정할 수 있다. 논문에서는 선형 확산 스케줄로 0.0001에서 0.02까지 선형적으로 증가하는데 이는 이미 잡음이 많은 후반부에는 빠르게 잡음을 추가하고 잡음보다는 이미지(정보)가 많은 초기에는 노이즈를 조금 줘서 신경망이 최대한 정보를 많이 배울 수 있도록 해주는 데에 그 의도가 있다.

 

def linear_diffusion_schedule(diffusion_times):
    min_rate = 0.0001
    max_rate = 0.02
    betas = min_rate + tf.convert_to_tensor(diffusion_times) * (max_rate - min_rate)
    alphas = 1-betas
    alpha_bars = tf.math.cumprod(alphas)#누적 곱인데, 곱해질수록 원래 이미지(신호)에서 얼만큼 소실되었는지를 의미
    signal_rates = alpha_bars
    noise_rates = 1- alpha_bars
    return noise_rates, signal_rates
    
T = 1000
diffusion_times = [x/T for x in range(T)]
linear_noise_rates, linear_signal_rates = linear_diffusion_schedule(diffusion_times)

 

이후 논문에서 코사인 확산 스케줄이 원래 선형 스케줄보다 성능이 우수하다는 사실이 밝혀졌다. 코사인 스케줄은 alpha_bars를 아래와 같이 정의한다.

 

아까 베타의 특성이 제곱해서 더했을 때 1이 되어야 하므로 sin제곱 + cos제곱 = 1을 활용하는 개념이다.

 

 

 

#offset과 스케일링이 없는 순수한 코사인 확산 스케줄
def cosine_diffusion_schedule(diffusion_times):
    signal_rates = tf.cos(diffusion_times * math.pi/2)
    noise_rates = tf.sin(diffusion_times * math.pi/2)
    return signal_rates, noise_rates

#오프셋 코사인 확산 스케줄은 잡음 추가 단계 초반에 잡음이 너무 작지 않도록 스케줄을 조정합니다.
def offset_cosine_diffusion_schedule(diffusion_times):
    min_signal_rate = 0.02
    max_signal_rate = 0.95
    start_angle = tf.acos(max_signal_rate)
    end_angle = tf.acos(min_signal_rate)

	diffusion_angles = start_angle + diffusion_times * (end_angle - start_angle)
    
    signal_rates = tf.cos(diffusion_angles)
    noise_rates = tf.sin(diffusion_angles)
    
    return noise_rates, signal_rates

 

 

2. 역방향 확산 과정

 

• 사실 딥러닝 관점으로 보면 신경망이 작동하는 시점은 이 과정에서이다. 잡음 추가 과정을 되돌릴 수 있도록 q(x_{t-1} | x_t)의 역방향 분포를 근사화하는 신경망을 구축하는 것이 모델의 핵심이다.

 

 

 

어떻게 학습시킬 수 있을까?

 

우리는 이전의 확산 과정을 통해서 각 시점별로 노이즈의 비율이 어떻게 되고, 그 결과 노이즈가 낀 이미지 픽셀값이 얼마인지를 알고 있다.  실제로 우리가 알고 싶은 내용은 얼만큼의 noise가 꼈는지이다.

 

noise가 얼만큼 꼈는지를 알 수 있다면, 우리는 원본 이미지를 아래와 같은 식으로 도출할 수 있다.

 

 

따라서 신경망이 학습하려는 대상은 노이즈의 비율과 노이즈가 낀 이미지 픽셀값이 주어졌을 때 실제 노이즈의 크기(Noises)이다. 이를 도식화하면 아래와 같이 정리해볼 수 있다. 각 시점별로 빨간색으로 표시된 값을 넣어서 파란색을 예측하도록 하는 것이다.

 

 

 

직관적으로는 그렇게 예측한 noise를 pred_noises라고 하며 위의 식에 Noises라고 적힌 곳에 pred_noises를 넣으면 우리는 이미지를 생성할 수 있고, loss function 계산할 때 원본 이미지의 픽셀과 pred_image의 픽셀의 차이를 활용할 때도 있다.(보통 우리는 noise를 얼만큼 넣었는지 이미 이전 확산 과정에서 저장했다면 그걸로 그냥 loss function으로 활용할 수도 있다.)

 

class DiffusionModel(models.Model):
    def __init__(self):
        super().__init__()
        self.normalizer = layers.Normalization()
        self.network = unet
        self.ema_network = models.clone_model(self.network)
        self.diffusion_schedule = offset_cosine_diffusion_schedule
    
    
    
    ...
    
    
    def denoise(self, noisy_images, noise_rates, signal_rates, training):
        if training:
        	network = self.network
        else:
        	network = self.ema_network
        pred_noises = network([noisy_images, noise_rates**2], training = training)
        pred_images = (noisy_images - noise_rates * pred_noises)/signal_rates
        #위의 pred_images에서 보다시피 우리는 noise_rates와 noisy_images만으로 pred_images를 만들 수 있다.
        
        return pred_noises, pred_images
    
    
    
    def train_step(self, images):
        images = self.normalizer(images, training = True)
        noises = tf.random.normal(shape = tf.shape(images))
        batch_size = tf.shape(images)[0]
        diffusion_times = tf.random.uniform(shape = (batch_size, 1, 1, 1), minval = 0.0, maxval = 1.0)
        noise_rates, signal_rates = self.diffusion_schedule(diffusion_times)
        noisy_images = signal_rates * images + noise_rates + noises
        with tf.GradientTape() as tape:
            pred_noises, pred_images = self.denoise(noisy_images, noise_rates, signal_rates, training = True)
            noise_loss = self.loss(noises, pred_noises)
       	gradients = tape.gradient(noise_loss, self.network.trainable_weights)
        self.optimizer.apply_gradients(zip(gradients, self.network.trainable_weights)
        self.noise_loss_tracker.update_state(noise_loss)
        
        for weight, ema_weight in zip(self.network.weights, self.ema_network.weights):
        	ema_wieght.assign(0.999 * ema_weight + (1 - 0.999) * weight)
        
        return {m.name : m.result for m in self.metrics}

 

어떤 신경망을 쓸까?

 

unet 구조를 사용하는데 아래와 같은 형태를 띈다.

 

unet 구조는 input이 두개이기 때문에 위와 같이 처리해주는데 이 때 잡음의 분산은 1개 값이기 때문에 이걸 최대한 정보를 많이 담을 수 있도록 고차원 벡터로 변환하는 과정이 포함되는데 이를 사인파 임베딩이라고 한다.

 

def sinusoidal_embedding(x):
    frequencies = tf.exp(tf.linspace(
    					tf.math.log(1.0), tf.math.log(1000.0), 16,))
    angular_speeds = 2.0 * math.pi * frequencies
    embeddings = tf.concat([tf.sin(angular_speeds * x), tf.cos(angular_speeds * X)], axis = 3)
    
    return embeddings

 

 

3. 실제 이미지 생성

 

보통 한번에 t -> 0으로 이미지를 생성하기 보다는 한스텝 한스텝씩 이미지를 생성해나가는 방식을 택한다.

#이미지 생성하기
def reverse_diffusion(self, initial_noise, diffusion_steps):
    num_images = initial_noise.shape[0]
    step_size = 1.0 / diffusion_steps
    current_images = initial_noise
    for step in range(diffusion_steps):
        diffusion_times = tf.ones((num_images, 1, 1, 1)) - step * step_size
        noise_rates, signal_rates = self.diffusion_schedule(diffusion_times)
        pred_noises, pred_images = self.denoise(current_images, noise_rates, signal_rates, training = False)
        next_diffusion_times = diffusion_times - step_size
        next_noise_rates, next_signal_rates = self.diffusion_schedule(next_diffusion_times)
        current_images = (next_signal_rates * pred_images + next_noise_rates * pred_noises)
    return pred_images


def denormalize(self, images):
    images = self.normalizer.mean + images * self.normalizer.variance ** 0.5
    return tf.clip_by_value(images, 0.0, 1.0)

def generate(self, num_images, diffusion_steps):
    initial_noise = tf.random.normal(shape = (num_images, 64, 64, 3))
    generated_images = self.reverse_diffusion(initial_noise, diffusion_steps)
    generated_images = self.denormalize(generated_images)
    return generated_images

 

 

4. 이미지 보간

 

추가로 이미지 퀄리티를 높이기 위해 가우스 잠재 공간에 있는 포인트 사이를 보간해주는 작업을 통해서 픽셀 공간에서 이미지 사이를 부드럽게 전환하는 방법을 썼다. 

 

 

5. 데이터셋

bash scripts/downloaders/download_kaggle_data.sh nunenuh pytorch-challenge-flower-dataset

 

 

 

헷갈렸던 점

모든 데이터셋에 대해서 결국 노이즈를 매우 오랜시간동안 더해가면, 동일한 소음 이미지로 수렴될텐데 어떻게 이런 확산모델은 새롭고 다양한 이미지를 만들어내는지?가 처음에 이해가 되지 않았다.

여기서 내가 잘못 안 내용은 랜덤한 잡음을 계속 더해주기 때문에 실은 동일한 T시점(확산 종료 시점)에 가더라도 이미지별로 서로 다른 이미지를 가진다는 점이다.

그리고 DALL-E와 같은 모델은 멀티모달 모델이기 때문에 이런 텍스트 임베딩 정보도 함께 학습시킨다는 점이 차이점이다.

 

 

참고 정보

 

1. 만들면서 배우는 생성 AI 2판

2. https://arxiv.org/abs/2006.11239 논문

3. https://w-log.tistory.com/entry/%EB%94%94%ED%93%A8%EC%A0%84-%EB%AA%A8%EB%8D%B8-%EB%85%BC%EB%AC%B8-%EB%A6%AC%EB%B7%B0-Denoising-Diffusion-Probabilistic-Models