diffusion 모델

Diffusion 모델(DDPM) 수식과 구현, 시각적 이해

아래는 “diffusion 모델 수식”을 처음 접하거나, 수식을 이해하고 싶어하시는 분들을 위해
개념 → 수식 → 구현 순서대로 정리한 가이드입니다.
(이 자료는 바로 실행해볼 수 있는 예시 코드와 수식의 의미를 직관적으로 설명한 도식도 포함합니다.)

필수 전제

• 파이썬 3.7+, PyTorch ≥ 1.7, CUDA ≥ 10.2
• torch, torchvision, numpy 등 기본 패키지 설치

---

1. 개념 정리

단계	내용	비고
STEP1	전진 확산 (Forward Diffusion)	𝑥₀(원본)를 차례로 노이즈에 가두는 과정
STEP2	역 확산 (Reverse Diffusion)	노이즈가 제거되면서 원본으로 복원되는 과정
STEP3	학습 목표	역 확산 네트워크(εθ) 를 학습시켜 ‘노이즈를 예측’하도록 함
STEP4	샘플링	역 확산 과정을 통해 완전히 랜덤한 노이즈에서 이미지 생성

---

2. 수식 상세

2‑1. 전진 확산 과정

q(\mathbf{x}_t \mid \mathbf{x}_{t-1}) = 
\mathcal{N}\!\bigl(
    \mathbf{x}_t;\;
    \sqrt{\alpha_t}\,\mathbf{x}_{t-1},
    \; (1-\alpha_t)\mathbf{I}
\bigr)

• β_t (노이즈 규모) : 0 < β_t < 1
• α_t = 1 − β_t
• ᾱ_t = ∏_{s=1}^{t} α_s

전진 확산을 한 번만 하면

q(\mathbf{x}_t \mid \mathbf{x}_0)
  = \mathcal{N}\!\bigl(
        \mathbf{x}_t;\;
        \sqrt{\bar\alpha_t}\,\mathbf{x}_0,
        \; (1-\bar\alpha_t)\mathbf{I}
    \bigr)

따라서 정확히 한 단계 안에서 x_t 를 직접 샘플할 수 있습니다.

도식

   x0  ──(β1)───>  x1  ──(β2)───>  x2  ── … ──>  xT
   |        |        |
   |        |        +─(노이즈)
   +─────>  (ᾱ_t)    (단계 t에서의 노이즈 비율)

• ᾱ_t 가 0에 가까워질수록 노이즈가 많이 섞인 이미지가 됩니다.
• T 가 1000 정도면 완전한 가우시안 잡음(평균 0, 분산 1)이 됩니다.

2‑2. 역 확산 과정

역 확산은 학습된 파라미터 θ 로 정의됩니다.

p_{\theta}(\mathbf{x}_{t-1}\mid \mathbf{x}_t)
  = \mathcal{N}\!\bigl(
        \mathbf{x}_{t-1};\;
        \mu_{\theta}(\mathbf{x}_t, t),
        \; \sigma_t^2 \mathbf{I}
    \bigr)

여기서

\mu_{\theta}(\mathbf{x}_t, t)
  = \frac{1}{\sqrt{\alpha_t}}
    \Bigl(
        \mathbf{x}_t
        - \frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}}
          \epsilon_{\theta}(\mathbf{x}_t, t)
    \Bigr)

• εθ(·,t) 는 노이즈 예측 네트워크 (예: UNet + CLIP) 가 예측하는
  노이즈 벡터입니다.
• σ_t 는 보통 일정하게 설정하거나, 학습을 위해 σ_t^2 = β_t 로 잡기도 합니다.

2‑3. 학습 목표 (ELBO → L2 Loss)

정상적인 VAE 의 ELBO와 동일한 구조가 등장합니다.
실제 학습에서는 복잡한 KL 항 대신에 단순한 L2 loss 가 사용됩니다:

\mathcal{L}_{θ}
  = \mathbb{E}_{\mathbf{x}_0, t, ε}
    \bigl[
        \|ε - ε_{\theta}(\sqrt{\bar{\alpha}_t}\mathbf{x}_0 + 
                            \sqrt{1-\bar{\alpha}_t} ε, t)\|^2
    \bigr]

• ε : 표준 가우시안 N(0, I) 로 샘플
• t : 1…T 중 균일하게 샘플
• θ : UNet/ResNet 등 신경망

정리하면
네트워크가 원본에 섞인 노이즈를 예측하도록 학습하고,
역 확산 단계에서는 이 예측값을 이용해 x_{t-1} 를 복원합니다.

---

3. 코드 – 파이썬(PyTorch) 예시

아래 예시는 간단한 MNIST(28×28) 를 대상으로 한 𝑇=1000 단계 DDPM 구현입니다.
전체 구조를 이해하는 데에만 초점을 두었으므로, 실제 고성능 모델을 위해서는
UNet, positional embedding, EMA 등 여러 최적화가 필요합니다.

# --------------------------------------------------
# 1) 노이즈 스케줄 정의 (beta_1 ... beta_T)
# --------------------------------------------------
import torch, torch.nn as nn, torch.optim as optim
import numpy as np

DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
T = 1000
betas = np.linspace(1e-4, 0.02, T)          # β_t : 1e-4 → 0.02
alphas = 1.0 - betas
alphas_bar = np.cumprod(alphas)

alphas_bar = torch.tensor(alphas_bar, dtype=torch.float32, device=DEVICE)

# --------------------------------------------------
# 2) 전진(노이즈) 샘플링 함수
# --------------------------------------------------
def forward_diffuse(x0, t):
    """
    x_t = sqrt(ᾱ_t) * x0 + sqrt(1-ᾱ_t) * ε
    """
    sqrt_alpha_bar = torch.sqrt(alphas_bar[t])
    sqrt_one_minus_alpha_bar = torch.sqrt(1 - alphas_bar[t])
    eps = torch.randn_like(x0)           # 표준 가우시안
    xt = sqrt_alpha_bar * x0 + sqrt_one_minus_alpha_bar * eps
    return xt, eps

# --------------------------------------------------
# 3) ε 예측 네트워크 (간단한 2‑층 MLP 예시)
# --------------------------------------------------
class NoisePredictor(nn.Module):
    def __init__(self, img_size):
        super().__init__()
        self.net = nn.Sequential(
            nn.Flatten(),
            nn.Linear(img_size*img_size, 512),
            nn.ReLU(),
            nn.Linear(512, img_size*img_size),
        )
    def forward(self, xt, t):
        # t를 one‑hot 으로 변환해 입력에 concat
        t_onehot = torch.zeros_like(xt)
        t_onehot[:] = t
        inp = torch.cat([xt, t_onehot], dim=1)
        return self.net(inp)

# --------------------------------------------------
# 4) 역 확산 (DDPM 샘플링)
# --------------------------------------------------
def p_sample(xt, t, eps_pred):
    """
    p_θ(x_{t-1} | x_t) 를 L2 (σ_t=0) 으로 deterministic
    """
    sqrt_alpha_t = torch.sqrt(alphas[t])
    beta_t = betas[t]
    sqrt_one_minus_alpha_bar = torch.sqrt(1 - alphas_bar[t])

    mu = (1/ sqrt_alpha_t) * (xt - beta_t / sqrt_one_minus_alpha_bar * eps_pred)
    # σ_t는 0으로 하면 deterministic
    return mu

def ddim_sample(initial_noise, eps_pred_fn, steps=50):
    """
    DDIM (deterministic) sampling
    """
    x_t = initial_noise
    for t in reversed(range(steps)):
        # t를 0~T 사이의 정수
        eps_pred = eps_pred_fn(x_t, t)
        x_t = p_sample(x_t, t, eps_pred)
    return x_t

# --------------------------------------------------
# 5) 훈련 루프
# --------------------------------------------------
batch_size = 64
img_size = 28
train_loader = torch.utils.data.DataLoader(
    torchvision.datasets.MNIST('.', download=True, transform=transforms.ToTensor()),
    batch_size=batch_size, shuffle=True)

model = NoisePredictor(img_size).to(DEVICE)
optimizer = optim.Adam(model.parameters(), lr=2e-4)

for epoch in range(200):          # 200 epoch 예시
    for imgs, _ in train_loader:
        imgs = imgs.view(batch_size, -1).to(DEVICE)   # (B, 784)
        t = torch.randint(0, T, (batch_size,), device=DEVICE)
        xt, eps = forward_diffuse(imgs, t)

        # 모델이 예측한 노이즈
        eps_pred = model(xt, t)

        loss = ((eps - eps_pred)**2).mean()
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    print(f"Epoch {epoch+1} | loss: {loss.item():.4f}")

# --------------------------------------------------
# 6) 샘플링 (전역 노이즈 → 이미지 생성)
# --------------------------------------------------
@torch.no_grad()
def sample_one():
    xt = torch.randn(batch_size, img_size*img_size, device=DEVICE)
    for t in reversed(range(T)):
        eps_pred = model(xt, t)
        xt = p_sample(xt, t, eps_pred)
    return xt.view(batch_size, 1, img_size, img_size)

generated_imgs = sample_one()

핵심 포인트

• forward_diffuse 에서 한 단계만으로 x_t 를 샘플할 수 있음 → 전진 과정이 deterministic
• 역 확산 p_sample 은 εθ 를 이용해 x_{t-1} 를 예측 → εθ 가 노이즈를 예측하도록 학습됨
• 손실은 L2 (노이즈 예측) 으로 매우 단순화

---

3. 시각적 설명 – “타임라인 + 노이즈 흐름”

t=0   ── β1 ──>  t=1   ── β2 ──>  t=2   ── … ──>  t=T
  |        |        |                     |
  |        |        +─(노이즈 삽입)        |
  +─────────────>  (ᾱ_t)  (노이즈 비율)   |
          𝑥0 = 깨끗한 이미지
          𝑥T = 완전한 가우시안 잡음

• 전진 확산:
  x₀ 에서 x_T 로 가는 길은 노이즈가 누적되는 마르코프 체인입니다.
  각 단계마다 가우시안 분포가 적용되고,
  ᾱ_t 가 0으로 가면 x_t 는 𝒩(0,1) 잡음과 거의 동일해집니다.

• 역 확산:
  x_T (노이즈) 에서 시작해 x₀ (깨끗한 이미지) 로 되돌아갑니다.
  이때 εθ 가 “현재 이미지가 얼마나 ‘잡음’인지를 예측”하도록 학습되면,
  역 확산이 노이즈를 점진적으로 제거할 수 있습니다.

---

4. 더 나아가기 – DDIM, Score‑Based 모델 등

4‑1. DDIM (Deterministic Diffusion Implicit Models)

DDPM에서 샘플링 단계는 ε 를 재현할 때마다 σ_t 를 넣어 stochastic하게
x_{t-1} 을 생성합니다.
DDIM 은 deterministic 경로를 만들어 샘플링 속도를 크게 끌어올립니다.

DDPM	DDIM
x_{t-1} = μ_θ(x_t, t) + σ_t * ε	x_{t-1} = μ_θ(x_t, t) (σ=0)
샘플링 횟수: T	샘플링 횟수: s (보통 50~100)

4‑2. Score‑Based Generative Models

Score‑matching 방식은 εθ 대신 score (∇ log p(x)) 를 예측합니다.
이때 β_t 가 매우 작은 경우(ε ≈ 0)에서는 εθ 가 ∇log p(x)` 와 동일합니다.
Song, J. et al. 2020 논문이 대표적입니다.

---

5. 실전 활용 팁

상황	참고 자료
HuggingFace Diffusers	pip install diffusers → from diffusers import DDPMPipeline
대용량 이미지	T 를 1000 → 300 으로 줄이면 속도 ↑, betas 를 Cosine 스케줄 사용
텍스트‑가이드	εθ 를 CLIP‑Embedding + UNet 으로 연결 (diffusers 에서 StableDiffusionPipeline 제공)
연산량 최적화	EMA (Exponential Moving Average) 모델 파라미터 사용, torch.compile 활용

---

6. 정리

1. Noise Predictor(εθ) 가 노이즈를 예측하도록 학습
2. 전진 확산 은 한 단계만으로 x_t 를 생성 → 전진 과정 deterministic
3. 역 확산 은 εθ 를 이용해 x_{t-1} 를 복원 → 역 확산이 노이즈를 제거
4. 손실은 노이즈 예측을 위한 L2 로 단순화

위 흐름을 파이썬 코드 한 줄 한 줄 따라가 보시면,
diffusion model이 어떻게 “노이즈를 덜어내는 일방향 흐름”을 만드는지 명확히 이해할 수 있습니다.