확산 모델(Diffusion Models) — DDPM 직접 구현

Ho, Jain, Abbeel (2020)은 이 분야가 도저히 손에서 놓지 못하는 레시피(recipe)를 만들어냈습니다. 천 개의 작은 스텝(step)에 걸쳐 데이터를 노이즈(noise)로 파괴하고, 하나의 신경망(neural net)이 그 노이즈를 예측하도록 학습시킨 뒤, 추론(inference) 시점에는 그 과정을 거꾸로 되돌립니다. 오늘날 이미지, 비디오, 3D, 음악을 다루는 거의 모든 주류 모델은 이 루프(loop) 위에서 동작하며, 그 위에 플로우 매칭(flow matching)이나 일관성 모델(consistency) 같은 기법을 얹어 쓰기도 합니다.

유형: Build 언어: Python 선수 지식: Phase 3 · 02 (Backprop), Phase 8 · 02 (Autoencoder와 VAE) 예상 시간: 약 75분

문제

우리는 p_data(x)에서 표본을 뽑는 샘플러(sampler)를 원합니다. GAN은 자주 발산하는(diverge) 미니맥스 게임(minimax game)을 풀어야 하고, VAE는 가우시안 디코더(Gaussian decoder) 때문에 흐릿한(blurry) 표본을 만들기 쉽습니다. 우리가 진짜로 원하는 것은 세 가지 조건을 동시에 만족하는 학습 목적(training objective)입니다. 첫째, 안장점(saddle point)이나 미니맥스가 없는 하나의 안정적인 손실 함수(single stable loss). 둘째, log p(x)의 하한(lower bound)을 제공해 우도(likelihood)를 얻을 수 있어야 합니다. 셋째, 최고 수준(SOTA) 품질에 맞먹는 표본을 만들 수 있어야 합니다.

Sohl-Dickstein et al. (2015)은 이론적인 답을 먼저 제시했습니다. 가우시안 노이즈를 점진적으로 더해 가는 마르코프 체인(Markov chain) q(x_t | x_{t-1})을 정의하고, 그 노이즈를 제거하는(denoise) 역방향 체인(reverse chain) p_θ(x_{t-1} | x_t)을 학습하는 방식입니다. 이어서 Ho, Jain, Abbeel (2020)은 손실 함수를 "노이즈를 예측한다"라는 한 줄짜리 식으로 단순화하고 수학을 깔끔하게 정리했습니다. 2020년에는 호기심거리(curiosity) 정도였지만, 2021년에는 당시 최고 수준의 표본을 만들어냈고, 2022년에는 Stable Diffusion으로 이어졌으며, 2026년에는 생성 모델의 기반 토대(substrate)가 되었습니다.

개념

순방향 과정(Forward process) q. T개의 작은 스텝에 걸쳐 가우시안 노이즈를 더합니다. 이 수식이 다루기 쉬운(tractable) 이유는, 누적된 스텝 역시 가우시안이 된다는 닫힌 형태(closed form) 표현을 가지기 때문입니다.

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

여기서 α̅_t = ∏_{s=1..t} (1 - β_s)이고, β_t 스케줄(schedule)을 사용합니다. β_t를 1e-4에서 0.02까지 T=1000 스텝에 걸쳐 선형(linearly)으로 늘리면 x_T는 거의 N(0, I)에 수렴합니다.

역방향 과정(Reverse process) p_θ. 추가된 노이즈를 예측하는 신경망 ε_θ(x_t, t)를 학습합니다. x_t가 주어지면 다음 식으로 노이즈를 제거합니다.

x_{t-1} = (1 / sqrt(α_t)) · ( x_t - (β_t / sqrt(1 - α̅_t)) · ε_θ(x_t, t) )  +  σ_t · z

σ_t는 sqrt(β_t)이거나 학습된 분산(learned variance) 값입니다. 식 자체는 다소 길어 보이지만, 본질은 사후 분포(posterior) q(x_{t-1} | x_t, x_0)에서 x_{t-1}을 풀어 정리한 뒤, x_0 자리에 노이즈로부터 추정한 값(noise-predicted estimate)을 대입한 단순한 대수(algebra)일 뿐입니다.

학습 손실(Training loss).

L_simple = E_{x_0, t, ε} [ || ε - ε_θ( sqrt(α̅_t) · x_0 + sqrt(1 - α̅_t) · ε,  t ) ||² ]

데이터에서 x_0를 뽑고, 무작위 시점 t를 고른 뒤, ε ~ N(0, I)를 뽑습니다. 닫힌 형태 식을 이용해 잡음이 섞인 x_t를 한 번에 계산하고, 그 노이즈를 회귀(regress)로 맞춥니다. 미니맥스도, KL도, 재매개변수화 기법(reparameterization trick)도 필요하지 않은 단일 손실입니다.

샘플링(Sampling). x_T ~ N(0, I)에서 시작해 t = T부터 1까지 역방향 스텝을 반복하면 끝입니다.

왜 작동하는가

세 가지 직관(intuition)이 있습니다.

1. 노이즈 제거는 쉽고, 생성은 어렵습니다. t=T에서는 데이터가 사실상 순수 노이즈이므로 신경망이 풀어야 할 문제는 자명(trivial)에 가깝습니다. t=0에서는 픽셀 몇 개만 다듬으면 됩니다. 중간 시점(intermediate t)에서는 문제 자체가 어렵지만, 모든 노이즈 수준에서 같은 가중치(weight)에 그래디언트(gradient)가 동시에 흘러 들어옵니다.
2. 사실은 스코어 매칭(score matching)을 다른 방식으로 푸는 것입니다. Vincent (2011)는 노이즈를 예측하는 일이 ∇_x log q(x_t | x_0), 즉 *스코어(score)*를 추정하는 일과 동등함을 증명했습니다. 역방향 확률미분방정식(reverse SDE)은 이 스코어를 따라 밀도 그래디언트(density gradient)를 거슬러 올라가며, 고확률(high-probability) 영역으로 향하는 안내된 무작위 보행(guided random walk)을 수행합니다.
3. 변분 하한(ELBO)이 단순한 MSE로 줄어듭니다. 전체 변분 하한에는 시점(timestep)마다 KL 항(term)이 들어 있습니다. DDPM의 매개변수화(parameterization)에서는 이 KL 항들이 특정한 계수가 붙은, 노이즈 예측에 대한 MSE 형태로 단순화됩니다. Ho는 그 계수를 떼어 낸 "simple" 손실을 사용했고, 그럼에도 품질은 오히려 향상되었습니다.

직접 만들기

code/main.py는 1차원(1-D) DDPM을 구현합니다. 데이터는 두 봉우리를 가진 혼합 분포(two-mode mixture)이고, "신경망"은 (x_t, t)를 입력으로 받아 예측 노이즈(predicted noise)를 출력하는 작은 다층 퍼셉트론(MLP)입니다. 학습은 한 줄짜리 손실이고, 샘플링은 역방향 체인을 반복합니다.

Step 1: 순방향 스케줄(closed form)

betas = [1e-4 + (0.02 - 1e-4) * t / (T - 1) for t in range(T)]
alphas = [1 - b for b in betas]
alpha_bars = []
cum = 1.0
for a in alphas:
    cum *= a
    alpha_bars.append(cum)

Step 2: x_t를 한 번에 표본 추출

def forward_sample(x0, t, alpha_bars, rng):
    a_bar = alpha_bars[t]
    eps = rng.gauss(0, 1)
    x_t = math.sqrt(a_bar) * x0 + math.sqrt(1 - a_bar) * eps
    return x_t, eps

Step 3: 학습 스텝 하나

def train_step(x0, model, alpha_bars, rng):
    t = rng.randrange(T)
    x_t, eps = forward_sample(x0, t, alpha_bars, rng)
    eps_hat = model_forward(model, x_t, t)
    loss = (eps - eps_hat) ** 2
    return loss, gradient_step(model, ...)

Step 4: 역방향 샘플링

def sample(model, alpha_bars, T, rng):
    x = rng.gauss(0, 1)
    for t in range(T - 1, -1, -1):
        eps_hat = model_forward(model, x, t)
        beta_t = 1 - alphas[t]
        x = (x - beta_t / math.sqrt(1 - alpha_bars[t]) * eps_hat) / math.sqrt(alphas[t])
        if t > 0:
            x += math.sqrt(beta_t) * rng.gauss(0, 1)
    return x

1차원 문제라면 시점 40개와 24-유닛 MLP만으로도 약 200 에폭(epoch) 안에 두 봉우리 혼합 분포를 학습합니다.

시점 조건화(Time conditioning)

신경망은 자신이 지금 어느 시점(timestep)을 디노이징(denoising)하고 있는지를 알아야 합니다. 표준적으로 쓰이는 선택지는 두 가지입니다.

• 사인 임베딩(Sinusoidal embedding). 트랜스포머(Transformer)의 위치 인코딩(positional encoding)과 비슷합니다. embed(t) = [sin(t/ω_0), cos(t/ω_0), sin(t/ω_1), ...]을 MLP에 통과시켜 신경망 내부로 브로드캐스트(broadcast)합니다.
• FiLM / 그룹 정규화(group-norm) 조건화. 임베딩(embedding)을 블록(block)별로 채널 단위의 스케일(scale)/바이어스(bias)로 사영(project)합니다.

토이(toy) 예제 코드는 사인 임베딩 후 이어 붙이기(concat) 방식을 사용합니다. 실서비스용 U-Net은 FiLM을 사용합니다.

자주 빠지는 함정(Pitfalls)

• 스케줄(Schedule) 선택이 매우 중요합니다. 선형 β는 DDPM의 기본값이지만, 코사인 스케줄(cosine schedule, Nichol & Dhariwal, 2021)이 같은 계산량 대비 더 나은 FID를 줍니다. 품질이 정체(plateau)되면 스케줄을 바꿔 보세요.
• 시점 임베딩(timestep embedding)은 깨지기 쉽습니다(fragile). 원시 t를 실수(float)로 그대로 넘기는 방식은 토이 1차원에서는 동작하지만 이미지에서는 실패합니다. 항상 제대로 된 임베딩을 사용해야 합니다.
• V-prediction과 ε-prediction의 비교. 시점 t가 아주 작거나 아주 큰 영역(narrow regimes)에서는 ε의 신호 대 잡음비(signal-to-noise)가 나쁩니다. V-prediction(v = α·ε - σ·x)이 더 안정적이며, SDXL, SD3, Flux가 이를 사용합니다.
• 분류기 없는 가이던스(Classifier-free guidance). 추론 시 조건부 ε와 비조건부 ε를 모두 계산한 뒤 ε_cfg = (1 + w) · ε_cond - w · ε_uncond (w ≈ 3-7)를 사용합니다. Lesson 08에서 다룹니다.
• 1000 스텝은 너무 많습니다. 실서비스에서는 DDIM(20-50 스텝), DPM-Solver(10-20 스텝), 또는 증류(distillation, 1-4 스텝)를 사용합니다. Lesson 12를 참고하세요.

사용해보기

디퓨전(Diffusion)은 보편적인 생성 모델 백본(generative backbone)입니다. 플로우 매칭(Flow matching, Lesson 13)은 같은 품질에서 더 빠른 추론 속도(inference speed)를 보여 주며 보통 승리하는, 2024-2026년의 경쟁자입니다.

산출물 만들기

outputs/skill-diffusion-trainer.md를 저장합니다. 이 스킬(skill)은 데이터셋과 컴퓨트 예산(compute budget)을 입력으로 받아 다음을 출력합니다. 스케줄(linear/cosine/sigmoid), 예측 대상(prediction target; ε/v/x), 스텝 수, 가이던스 스케일(guidance scale), 샘플러 계열(sampler family), 그리고 평가 프로토콜(eval protocol)입니다.

연습문제

1. 쉬움. code/main.py에서 T를 40에서 10으로 줄여 봅니다. 출력 히스토그램(histogram)에서 표본 품질이 어떻게 나빠지는지 관찰하세요. 어느 T 값에서 두 봉우리 구조(two-mode structure)가 무너지는지 확인합니다.
2. 중간. ε-prediction을 v-prediction으로 바꿔 보세요. 역방향 스텝(reverse step)을 다시 유도(derive)하고, 최종 표본 품질을 비교합니다.
3. 어려움. 분류기 없는 가이던스(Classifier-free guidance)를 추가합니다. 클래스 라벨 c ∈ {0, 1}로 조건화하고, 학습 도중 10% 확률로 라벨을 드롭(drop)합니다. 샘플링에서는 ε = (1+w)·ε_cond - w·ε_uncond를 사용하고 w = 0, 1, 3, 7에서 조건부 모드 적중률(conditional-mode-hit rate)을 측정하세요.

핵심 용어

실서비스 메모: 디퓨전 추론은 결국 스텝 수(step-count) 문제다

DDPM 논문은 T=1000번의 역방향 스텝을 실행하지만, 실서비스(production)에서는 누구도 그대로 배포하지 않습니다. 실제 추론 스택은 세 가지 전략 중 하나를 고르며, 각각은 LLM에서 흔히 이야기하는 "지연(latency)이 어디서 오는가"라는 관점과 정확히 맞아떨어집니다.

1. 같은 모델, 더 빠른 샘플러. DDIM(20-50 스텝), DPM-Solver++(10-20), UniPC(8-16). 학습된 ε_θ 가중치는 그대로 둔 채 역방향 루프(reverse loop)만 갈아 끼우는 대체(drop-in replacement) 방식이며, 지연 시간을 20-50배 줄입니다.
2. 증류(Distillation). 더 적은 스텝으로 교사(teacher)를 모사하는 학생(student) 모델을 학습합니다. Progressive Distillation(2 → 1), Consistency Models(임의 → 1-4), LCM, SDXL-Turbo, SD3-Turbo가 여기에 속합니다. 지연 시간을 추가로 5-10배 줄이지만, 재학습(retraining)이 필요합니다.
3. 캐싱(Caching)과 컴파일(Compilation). torch.compile(unet, mode="reduce-overhead"), TensorRT-LLM의 디퓨전 백엔드, xformers/SDPA 어텐션(attention), bf16 가중치 같은 기법을 씁니다. 스텝당 지연을 약 2배 줄여 주며, 위의 (1), (2)와 함께 누적 효과를 낼 수 있습니다.

실서비스 디퓨전 서버에서 예산(budget) 논의는 LLM 실서비스 문헌과 동일한 구조입니다. 지연은 num_steps × step_cost + VAE_decode, 처리량(throughput)은 batch_size × (num_steps × step_cost)^-1로 정리됩니다. 사용자 관점에서 이미지 생성은 한 번에 모두 결과가 나오는(all-at-once) 방식이므로, TTFT(첫 토큰까지의 시간)는 한 스텝 정도로 짧고, TPOT(토큰당 시간)에 해당하는 값은 사실상 전체 응답 시간 그 자체입니다.

더 읽을거리

• Sohl-Dickstein et al. (2015). Deep Unsupervised Learning using Nonequilibrium Thermodynamics — 시대를 앞서간 디퓨전의 이론적 출발점입니다.
• Ho, Jain, Abbeel (2020). Denoising Diffusion Probabilistic Models — DDPM 논문입니다.
• Song, Meng, Ermon (2021). Denoising Diffusion Implicit Models — DDIM, 더 적은 스텝으로 표본을 만듭니다.
• Nichol & Dhariwal (2021). Improved DDPM — 코사인 스케줄과 학습된 분산입니다.
• Dhariwal & Nichol (2021). Diffusion Models Beat GANs on Image Synthesis — 분류기 가이던스(classifier guidance)입니다.
• Ho & Salimans (2022). Classifier-Free Diffusion Guidance — CFG 논문입니다.
• Karras et al. (2022). Elucidating the Design Space of Diffusion-Based Generative Models (EDM) — 통일된 표기법으로 정리한 가장 깔끔한 레시피입니다.