주석이 달린 확산 모델 – DDPM 구현 상세 안내

TL;DR

Hugging Face는 원본 Denoising Diffusion Probabilistic Model (DDPM)을 구현한 완전한 주석이 달린 PyTorch 노트북을 공개했으며, 순수 Gaussian 노이즈에서 이미지를 복원하는 신경망을 학습하고, 저해상도 데이터셋에서의 학습 및 고품질 이미지 샘플링을 시연합니다. 이 자료는 수학, 네트워크 구조, 학습 루프, 추론 절차를 명확히 설명하여 확산 모델을 실무자들이 쉽게 접근할 수 있도록 합니다.

확산 모델이 하는 일 (핵심 요점)

확산 모델은 무작위 Gaussian 노이즈를 현실적인 데이터로 변환하는 역방향 복원 과정을 학습합니다. 전방 과정은 미리 정의된 스케줄에 따라 노이즈를 추가하고, U‑Net 형태의 신경망이 각 타임스텝에서 추가된 노이즈를 예측하여 정해진 단계 수만큼 이미지 생성을 가능하게 합니다.

전방 확산: 분산 스케줄을 통한 노이즈 추가

Conclusion: 전방 과정은 닫힌 형태의 Gaussian 전이로, 임의의 노이즈가 섞인 타임스텝을 직접 샘플링할 수 있어 반복적인 노이즈 추가가 필요하지 않습니다.

The forward diffusion distribution is

q(x_t | x_{t-1}) = N(x_t ; sqrt(1-β_t)·x_{t-1}, β_t·I)

where the variance schedule (β_1,…,β_T) is monotonic (e.g., linear, cosine, quadratic, or sigmoid). By repeatedly applying this transition, the marginal

q(x_t | x_0) = N(x_t ; sqrt(\bar α_t)·x_0, (1-\bar α_t)·I)

can be sampled directly using pre‑computed (\bar α_t = \prod_{s=1}^t (1-β_s)). This “nice property” lets the training algorithm draw a random timestep (t) and corrupt a real image (x_0) in one step.

역방향 확산: 노이즈 예측을 통한 평균 학습

Conclusion: 모델은 타임스텝 (t)에 추가된 정확한 Gaussian 노이즈 (ε)를 예측하도록 학습되며, 예측된 노이즈는 역 Gaussian 분포의 평균을 계산하는 데 사용됩니다.

The reverse conditional is assumed Gaussian:

p_θ(x_{t-1} | x_t) = N(x_{t-1} ; μ_θ(x_t, t), σ_t^2·I)

DDPM fixes (σ_t^2) to the known forward variance and learns only (μ_θ). By re‑parameterising the mean in terms of the noise predictor (ε_θ):

μ_θ(x_t, t) = (1/√α_t)·(x_t - (β_t/√(1-\bar α_t))·ε_θ(x_t, t))

the training loss reduces to a simple mean‑squared error between the true noise (ε) and the network output (ε_θ):

L_t = || ε - ε_θ(x_t, t) ||^2

각 배치마다 무작위로 (t)를 샘플링하면 변분 하한의 편향되지 않은 추정치를 얻을 수 있습니다.

네트워크 구조: 조건부 U‑Net

Conclusion: DDPM은 시간 조건부 U‑Net을 사용하며, 사인파 위치 임베딩, 잔차 블록, 어텐션, 그룹 정규화를 통해 모든 공간 해상도에서 노이즈를 예측합니다.

Key components:

  • Sinusoidal position embeddings: 타임스텝 (t)를 인코딩하고 작은 MLP를 통해 각 ResNet 블록에 주입합니다.
  • Weight‑standardized convolutions: 그룹 정규화와 결합될 때 학습 안정성을 향상시킵니다.
  • ResNet blocks (두 개의 convolution‑norm‑SiLU 레이어와 시간 임베딩에서 오는 선택적 scale‑shift 포함) : 주요 특징 변환을 제공합니다.
  • Attention modules (전체 멀티‑헤드 또는 선형 어텐션) : 장거리 종속성을 포착합니다.
  • Group normalization: 어텐션 전에 적용됩니다 (PreNorm).
  • Down/upsampling paths: 공간 해상도를 절반 또는 두 배로 조정하면서 채널 깊이를 유지하며, 고전적인 U‑Net 설계를 반영합니다.
  • Final head: 연결된 bottleneck 특징을 이미지 형태로 다시 매핑합니다.

전체 PyTorch 클래스 Unet은 이러한 구성 요소들을 조합하여, 노이즈가 섞인 이미지 텐서와 타임스텝 텐서를 입력받고 동일한 형태의 노이즈 텐서를 반환합니다.

학습 루프: 확률적 타임스텝 샘플링 및 Huber 손실

Conclusion: 학습은 각 배치마다 무작위 타임스텝을 샘플링하고, 닫힌 형태의 전방 과정을 사용해 입력을 손상시킨 뒤, 실제 노이즈와 예측 노이즈 사이의 Huber 손실을 최소화하는 방식으로 진행됩니다.

Pseudo‑code (simplified):

for epoch in range(num_epochs):
    for batch in dataloader:
        t = torch.randint(0, T, (batch_size,)).to(device)          # uniform timestep
        loss = p_losses(model, batch, t, loss_type='huber')        # MSE/Huber on noise
        loss.backward()
        optimizer.step()

헬퍼 p_losses는 q_sample을 호출해 (x_t)를 얻고, 네트워크 출력 model(x_t, t)와의 손실을 계산합니다. 역 과정을 사용한 주기적인 샘플링은 학습 진행 상황을 시각화합니다.

샘플링 (추론): 확산 체인 역전

Conclusion: 생성은 순수 Gaussian 노이즈에서 시작해 학습된 복원 단계를 반복 적용합니다; 각 단계는 예측된 노이즈를 사용해 사후 평균을 계산하고, 알려진 분산 스케줄에 따라 보정된 Gaussian 노이즈를 추가합니다.

Algorithm (Algorithm 2 in the DDPM paper):

x_T = torch.randn(shape)                     # start from noise
for t in reversed(range(T)):
    pred_noise = model(x_t, t)
    mean = (1/√α_t) * (x_t - β_t/√(1-\bar α_t) * pred_noise)
    if t > 0:
        x_{t-1} = mean + √posterior_variance_t * torch.randn_like(x_t)
    else:
        x_0 = mean

제공된 노트북은 p_sample_loop에서 이 루프를 구현하고, 복원 과정을 GIF로 시각화합니다.

Fashion‑MNIST 전체 예제

Conclusion: 이 튜토리얼은 28×28 Fashion‑MNIST 데이터셋에 DDPM을 학습시켜, 몇 천 번의 학습 단계 후에 인식 가능한 의류 아이템을 생성합니다.

  • 데이터 로딩은 🤗 datasets 라이브러리를 사용하며, 실시간 변환(무작위 수평 뒤집기, ([-1,1]) 로 스케일링)을 적용합니다.
  • 모델 하이퍼파라미터: dim=image_size, dim_mults=(1,2,4), 흑백 이미지의 경우 channels=1.
  • 옵티마이저: 학습률 1e‑3인 Adam.
  • 학습은 6 epoch 동안 진행되며, 손실이 빠르게 0.05 이하로 감소합니다.
  • 학습 후 샘플링하면 선명한 티‑셔츠 형태가 생성되어, 구현이 저해상도 데이터에서도 작동함을 확인합니다.

확산 문헌에서의 위치

Conclusion: 이 주석 구현은 원본 DDPM (Ho et al., 2020)을 재현하며, 분산 학습(Nichol et al., 2021), 계층형 확산(Ho et al., 2021), classifier‑free guidance, 대규모 텍스트‑투‑이미지 모델(DALL‑E 2, Imagen) 등 이후 많은 발전의 기반이 됩니다.

블로그 포스트에 나열된 주요 후속 연구:

  • Improved DDPM – 평균과 분산을 모두 학습하여 샘플 품질을 향상시킵니다.
  • Cascaded Diffusion Models – 고해상도 합성을 위해 여러 확산 모델을 순차적으로 쌓습니다.
  • Diffusion Models Beat GANs – 아키텍처 조정 및 classifier guidance를 통해 우수한 FID 점수를 보여줍니다.
  • Classifier‑Free Guidance – 조건부 생성 시 외부 classifier가 필요하지 않게 합니다.
  • DALL‑E 2 & ImageGen – 확산을 CLIP 임베딩이나 대규모 언어 모델과 결합해 텍스트 조건 이미지 합성을 수행합니다.

주된 단점은 여전히 많은 복원 단계(보통 1000 단계)가 필요하다는 점이며, 최근 연구(예: 고차 솔버)는 이를 10 단계 정도로 줄이는 방법을 제시합니다.

개발자를 위한 실용적인 시사점

  • 노트북은 바로 실행 가능한 참고 구현을 제공하므로, 데이터셋을 교체하거나 T를 늘리거나 U‑Net을 더 큰 백본으로 교체해 확장할 수 있습니다.
  • 분산 스케줄(linear_beta_schedule, cosine_beta_schedule 등)을 조정해 샘플 품질과 속도 사이의 균형을 맞출 수 있습니다.
  • 견고성을 위해 Huber 손실을 사용하고, 코드가 L1/L2 손실도 지원합니다.
  • 주기적인 샘플링(save_and_sample_every)은 모드 붕괴나 학습 불안정을 모니터링하는 데 필수적입니다.
  • 더 빠른 추론을 위해 DDIM이나 확률적 샘플러 변형을 구현하거나, 최근의 “few‑step” 솔버를 도입하는 것을 고려하세요.

Sources