torch.cumprod() w.r.t. diffusion noise scheduling

 

 

[cumprod: cumulative product 함수]

 

torch.cumprod 함수는 PyTorch 에서 제공하는 함수로, 텐서의 누적 곱을 요소별로(element-wise) 계산해준다.

사용법은 다음과 같다. 

 

import torch

tensor = torch.tensor([1, 2, 3, 4])
cumprod_result = torch.cumprod(tensor, dim=0)
print(cumprod_result)

tensor([ 1,  2,  6, 24])

 

예를 들어, 'torch.tensor([1,2,3,4])' 텐서에 대해 torch.cumprod 함수를 dim=0 (첫번째 차원)으로 적용하면, 이 함수는 각 요소의 누적곱을 계산하여 [1,2,6,24] 를 반환한다. 

1

1x2=2

1x2x3=6

1x2x3x4=24

 

---

 

[Diffusion model에서 쓰임새]

class Diffusion:
    def __init__(self, noise_steps=1000, beta_start=1e-4, beta_end=0.02, img_size=256, device="cuda"):
        self.noise_steps = noise_steps
        self.beta_start = beta_start
        self.beta_end = beta_end
        self.img_size = img_size
        self.device = device

        self.beta = self.prepare_noise_schedule().to(device)
        self.alpha = 1. - self.beta
        self.alpha_hat = torch.cumprod(self.alpha, dim=0)

    def prepare_noise_schedule(self):
        return torch.linspace(self.beta_start, self.beta_end, self.noise_steps) # 각 timestep 에 해당하는 beta값을 반환

 

 

디퓨전 모델에서는 forward 및 backward 과정에서 각 단계별로 노이즈 수준이 다른데, 전 단계의 이미지에서 점진적으로 노이즈를 추가하거나 제거하기 때문이다. Forward process로 예를 들자면 이전 스텝의 이미지에서 다음 스텝의 이미지로 noise를 가할 때 이미지와 노이즈에는 각각 $\sqrt{\alpha_t}$, $\sqrt{1-\alpha_{t}}$ 라는 계수(coefficient)가 곱해진다. 

$$X_t=\sqrt{\alpha_t}X_{t-1}+\sqrt{1-\alpha_{t}}\epsilon_{t-1}, \  \ \alpha _{t}=1-\beta_{t}$$

그런데 이때 t-1에서 t로 가지 않고 coefficient들의 누적 곱을 사용하여 t=0에서 t=T 로 가는 방법도 있다. 

위의 forward process 식에 t=0, 1, 2, ... 를 대입해보면 아래와 같은 식들을 얻을 수 있다. 

 

 

여기서 식 ②에 식 ①을 대입해보면 아래와 같이 쓸 수 있다. 

 

 

여기서 $\epsilon_{0}$과 $\epsilon_{1}$는 평균이 0, 분산이 1인 가우시안 분포로부터 표본의 개수(=입력 데이터의 shape) 만큼 값을 무작위로 추출하여 얻은 텐서이다. (이론상 각 epsilon이 표준가우시안분포라고 가정) 이때 두 분포 $\epsilon_{0}, \epsilon_{1}$ 를 합칠 수 있다. 분산의 성질에 따라 epsilon 앞에 붙은 계수의 제곱이 분산이 된다. 

즉, $(\alpha_2)(1-\alpha _{1})$ 가 $\epsilon_0$의 분산이 되고 $(1-\alpha_{2})$가 $\epsilon_1$의 분산이 된다. 

두 개의 분산을 가진 가우시안 분포를 합하면 두 분포의 덧셈을 분산으로 가지는 새로운 분포를 만들 수있다. 

 

$$\bar{ \alpha _{t}}=\prod_{i=1}^{t}\alpha _{i}=\prod_{i=1}^{t}(1-\beta _{i})$$

따라서 X_t 는 alpha 의 누적곱들의 제곱근을 X_0 에 곱하고 1-(alpha 누적곱) 의 제곱근을 e_0에 곱하여 얻을 수 있다. 

alpha들의 누적곱을 하기 위해서 torch.cumprod 함수를 사용한다.

 

self.alpha_hat = torch.cumprod(self.alpha, dim=0)

 

이로써 각 단계에서의 노이즈 수준을 결정할 수 있다. 

 

Diffusion의 자세한 수식 과정은 아래 링크를 보면 잘 정리되어 있다. 

https://lilianweng.github.io/posts/2021-07-11-diffusion-models/

What are Diffusion Models?