Diffusion model 노이즈 스케줄러 완전 정복 | Linear Noise Scheduler 개념 & 코드 구현

 

DDPM 노이즈 스케줄러 완전 정복 | Linear Noise Scheduler 개념 & 코드 구현

Diffusion Model 딥러닝 생성 AI  · 읽는 데 약 7분 소요

💡 이 글에서 다루는 내용
Diffusion Model의 Forward Process에서 핵심 역할을 하는 Noise Scheduler의 수식 원리부터, PyTorch로 직접 구현하는 LinearNoiseScheduler 코드까지 단계적으로 설명합니다. Beta schedule → alpha → ᾱ(누적곱) → x_t 계산까지 막힘없이 이해하실 수 있습니다.

📋 목차

1. Noise Scheduler 개념 — β, α, ᾱ 의 관계
2. LinearNoiseScheduler 코드 구현 분석
3. x_t 계산: add_noise() 메서드 완전 분석
4. 핵심 요약

1. Noise Scheduler 개념 — β, α, ᾱ의 관계

Diffusion Model의 학습은 크게 두 단계로 구성됩니다. 원본 데이터에 노이즈를 점진적으로 주입하는 Forward Process와, 노이즈를 제거하며 원본을 복원하는 Reverse Process입니다. 이 중 Forward Process를 수학적으로 제어하는 핵심 요소가 바로 Noise Scheduler입니다.

x_t 공식: 원본에서 노이즈 이미지를 한 번에 계산하기

$t$ 시점의 노이즈 이미지 $x_t$는 원본 $x_0$와 가우시안 노이즈 $\epsilon$의 선형 결합으로 표현됩니다.

$$x_t = \sqrt{\bar{\alpha}_t}\,x_0 + \sqrt{1-\bar{\alpha}_t}\,\epsilon$$

여기서 $\bar{\alpha}_t \in [0,\,1]$이므로, 이 식은 "원본을 조금 남기고, 노이즈를 조금 섞는다"는 직관적인 의미를 가집니다. $t$가 커질수록 $\bar{\alpha}_t$는 0에 가까워지고, 결국 순수한 노이즈가 됩니다.

β (Beta) 스케줄 정의

Forward Process는 조건부 가우시안 분포로 정의됩니다.

$$q(x_t \mid x_{t-1}) = \mathcal{N}\!\left(x_t;\;\sqrt{1-\beta_t}\,x_{t-1},\;\beta_t I\right)$$

여기서 $\beta_t$는 노이즈를 얼마나 주입할지 결정하는 하이퍼파라미터입니다. 일반적으로 Linear schedule에서는 $\beta_0 = 0.001$부터 $\beta_{999} = 0.02$까지 1,000 스텝에 걸쳐 선형으로 증가합니다.

📌 포인트: $\beta_t$가 작을수록 노이즈가 천천히 주입되고, 클수록 빠르게 주입됩니다. 초기에 작은 $\beta$를 사용하는 이유는 Reverse Process가 더 정밀하게 복원할 수 있도록 하기 위해서입니다.

α, ᾱ 유도: 왜 이 치환을 사용하는가?

$x_2$를 구하려면, $x_0$에 다음 세 항이 연속으로 곱해집니다.

$$\sqrt{1-\beta_0} \cdot \sqrt{1-\beta_1} \cdot \sqrt{1-\beta_2}$$

이를 누적곱으로 일반화하면 $\displaystyle\prod_{s=1}^{t}(1-\beta_s)$ 형태가 되어 식이 복잡해집니다. 이를 깔끔하게 표현하기 위해 다음과 같이 정의합니다.

$$\alpha_t = 1 - \beta_t, \qquad \bar{\alpha}_t = \prod_{s=1}^{t}\alpha_s$$

이렇게 정의하면 처음의 $x_t$ 공식이 성립하며, 임의의 타임스텝 $t$에서 $x_0$로부터 단 한 번의 연산으로 $x_t$를 계산할 수 있습니다. 이것이 DDPM의 핵심 수학적 성질 중 하나입니다.

2. LinearNoiseScheduler 코드 구현 분석

위 개념을 PyTorch로 구현한 LinearNoiseScheduler 클래스를 살펴봅니다.

class LinearNoiseScheduler:
    """
    DDPM에서 사용하는 Linear Noise Scheduler 클래스
    """
    def __init__(self, num_timesteps, beta_start, beta_end):
        self.num_timesteps = num_timesteps          # (1) 전체 타임스텝 수 (보통 1000)
        self.beta_start = beta_start                # (2) β 시작값 (보통 0.001)
        self.beta_end = beta_end                    # (3) β 끝값   (보통 0.02)

        self.betas = torch.linspace(               # (4) β 스케줄 생성
            beta_start, beta_end, num_timesteps
        )
        self.alphas = 1. - self.betas              # (5) α = 1 - β

        self.alpha_cum_prod = torch.cumprod(        # (6) ᾱ = cumulative product of α
            self.alphas, dim=0
        )
        self.sqrt_alpha_cum_prod = torch.sqrt(      # (7) √ᾱ
            self.alpha_cum_prod
        )
        self.sqrt_one_minus_alpha_cum_prod = torch.sqrt(  # (8) √(1-ᾱ)
            1 - self.alpha_cum_prod
        )

코드의 흐름을 단계별로 정리하면 다음과 같습니다.

(4) Beta 스케줄: torch.linspace(0.001, 0.02, 1000)으로 $\beta$를 선형으로 생성합니다. 결과적으로 0.001에서 0.02까지 일정한 간격으로 증가하는 1,000개의 값이 만들어집니다.

(5) Alpha: $\alpha_t = 1 - \beta_t$이므로, $\beta$의 정반대 형태를 가집니다. $\beta$가 증가하면 $\alpha$는 감소합니다.

(6) ᾱ 누적곱: torch.cumprod()를 통해 $\bar{\alpha}_t = \prod_{s=1}^{t}\alpha_s$를 계산합니다. 선형이었던 $\alpha$의 누적곱은 sigmoid 형태로 변환됩니다 — 초기에는 느리게 감소하다가 후반부에 급격히 0에 수렴합니다. 이것이 Linear schedule의 특징적인 $\bar{\alpha}_t$ 곡선입니다.

(7), (8): 최종적으로 $x_t$ 공식에 직접 사용할 $\sqrt{\bar{\alpha}_t}$와 $\sqrt{1-\bar{\alpha}_t}$를 미리 계산해 저장합니다.

💡 설계 의도: __init__에서 모든 타임스텝의 값을 미리 계산해두는 이유는, 학습 중 매 배치마다 이 값들을 반복 계산하는 오버헤드를 줄이기 위해서입니다.

3. x_t 계산: add_noise() 메서드 완전 분석

$x_t$를 실제로 계산하는 add_noise() 메서드를 살펴봅니다. 다음 수식이 코드로 어떻게 구현되는지 확인합니다.

$$x_t = \sqrt{\bar{\alpha}_t}\,x_0 + \sqrt{1-\bar{\alpha}_t}\,\epsilon$$

def add_noise(self, original, noise, t):
    """
    original : 원본 이미지 x_0  (shape: B, C, H, W)
    noise    : 가우시안 노이즈 ε (torch.randn_like(original))
    t        : 배치별 타임스텝  (shape: B,)
    """
    original_shape = original.shape
    batch_size = original_shape[0]

    # (1) 배치 내 각 샘플의 타임스텝 t에 해당하는 √ᾱ_t 추출
    sqrt_alpha_cum_prod = (
        self.sqrt_alpha_cum_prod.to(original.device)[t]
        .reshape(batch_size)
    )

    # (2) 마찬가지로 √(1-ᾱ_t) 추출
    sqrt_one_minus_alpha_cum_prod = (
        self.sqrt_one_minus_alpha_cum_prod.to(original.device)[t]
        .reshape(batch_size)
    )

    # (3) 브로드캐스팅을 위한 shape 맞추기: (B,) → (B, 1, 1, 1)
    for _ in range(len(original_shape) - 1):
        sqrt_alpha_cum_prod = sqrt_alpha_cum_prod.unsqueeze(-1)
    for _ in range(len(original_shape) - 1):
        sqrt_one_minus_alpha_cum_prod = sqrt_one_minus_alpha_cum_prod.unsqueeze(-1)

    # (4) x_t = √ᾱ_t * x_0 + √(1-ᾱ_t) * ε
    return (
        sqrt_alpha_cum_prod * original
        + sqrt_one_minus_alpha_cum_prod * noise
    )

코드 포인트별 해설

(1), (2) 인덱싱: 배치 내 각 샘플은 서로 다른 타임스텝 $t$를 가질 수 있습니다. [t] 인덱싱으로 각 샘플에 맞는 $\sqrt{\bar{\alpha}_t}$를 정확히 추출합니다.

(3) Unsqueeze: 이미지 텐서의 shape가 (B, C, H, W)인 반면, 추출된 스케줄러 값은 (B,)입니다. 원소별 곱셈을 위해 unsqueeze(-1)을 3번 적용하여 (B, 1, 1, 1)로 맞춥니다. PyTorch의 브로드캐스팅이 나머지를 처리합니다.

(4) 최종 연산: 수식 $x_t = \sqrt{\bar{\alpha}_t}x_0 + \sqrt{1-\bar{\alpha}_t}\epsilon$이 코드 한 줄로 구현됩니다. 이것이 DDPM Forward Process의 핵심 연산입니다.

📌 실제 사용 예시:
noise = torch.randn_like(x0)
t = torch.randint(0, 1000, (batch_size,))
x_t = scheduler.add_noise(x0, noise, t)

4. 핵심 요약

✅ 이 글의 핵심 포인트

• β (Beta): 타임스텝마다 얼마나 노이즈를 주입할지 결정하는 스케줄. Linear schedule에서는 0.001 → 0.02 선형 증가.
• α = 1 - β: β의 보수값. 원본 정보가 얼마나 보존되는지를 나타냄.
• ᾱ = cumulative product(α): t 스텝 동안의 누적 보존률. sigmoid 형태로 감소.
• x_t 공식: $x_t = \sqrt{\bar{\alpha}_t}x_0 + \sqrt{1-\bar{\alpha}_t}\epsilon$ — 임의의 t에서 한 번에 계산 가능.
• add_noise(): 위 공식을 배치 연산으로 구현. unsqueeze로 shape를 맞춰 브로드캐스팅 활용.

이 글이 도움이 되셨다면 구독과 댓글 남겨주세요 🙌

DDPM Diffusion Model Noise Scheduler 딥러닝 생성 AI PyTorch Forward Process Beta Schedule 머신러닝