직접 해 보기: 2D 장난감 확산 모델을 numpy로 만들기

직접 해 보기: 2D 장난감 확산 모델을 numpy로 만들기

확산 모델의 학습과 생성을 딥러닝 프레임워크 없이 numpy만으로 끝까지 구현합니다. 원 위에 놓인 가우시안 8개에서 점을 뽑아 잡음을 섞고, 작은 MLP가 잡음을 맞히도록 손으로 쓴 역전파 + Adam으로 학습한 뒤, DDPM과 DDIM으로 순수 잡음에서 점을 되살립니다. 실제로 돌린 출력과 함께 단계별로 설명합니다.

들어가기 전에

• 잡음 제거 확산 확률 모델 (DDPM) — 이 코드가 구현하는 식의 출처입니다.
• 확산 모델 실전 — $\hat x_0$ 추정과 DDIM 걸음.
• 신경망과 딥러닝 — 역전파의 네 방정식. 4단계에서 직접 씁니다.

한눈에 보기

• 확산 모델 학습이 "잡음 섞기 → 잡음 맞히기 → 제곱 오차" 세 줄로 끝난다는 것을 코드로 확인합니다.
• 시점 $t$ 를 신경망에 알려 주는 사인 · 코사인 임베딩을 직접 만듭니다.
• 행렬 곱과 ReLU만으로 MLP의 순전파 · 역전파 · Adam 갱신을 구현합니다.
• 학습한 모델로 2,000개 점을 생성해, 순수 잡음(0.5 안쪽 비율 14.5%)이 94.1%로 바뀌는 것을 봅니다.
• 다시 학습하지 않고 DDIM 20단계로도 비슷한 결과(90.9%)가 나오는 것을 확인합니다.

전체 흐름

123456789flowchart LR
  D[데이터 x0<br/>원 위 가우시안 8개] --> N[시점 t · 잡음 ε 뽑기]
  N --> Q[x_t = √ᾱ_t·x0 + √1−ᾱ_t·ε]
  Q --> M[MLP ε_θ x_t, t]
  M --> L[손실 ‖ε − ε_θ‖²]
  L --> B[역전파 + Adam]
  B -->|6,000번| N
  B --> S[생성: 잡음 x_T 에서<br/>T번 거꾸로 걷기]
  S --> E[평가: 중심까지 거리 · 봉우리별 개수]

1단계: 데이터 — 원 위의 가우시안 8개

반지름 2인 원 위에 같은 간격으로 중심 8개를 놓고, 각 중심 주변에 표준편차 0.1로 점을 흩뿌립니다. 봉우리(모드)가 8개라서 "모든 봉우리를 골고루 만드는가"로 모드 붕괴 여부를 쉽게 볼 수 있습니다(GAN 실전 의 Unrolled GAN 실험과 비슷한 설정).

12345678910import numpy as np, time, sys
sys.stdout.reconfigure(encoding="utf-8")

# ---------- 1. 데이터: 반지름 2인 원 위의 가우시안 8개 ----------
rng = np.random.default_rng(0)
K = 8
centers = np.stack([2*np.cos(2*np.pi*np.arange(K)/K), 2*np.sin(2*np.pi*np.arange(K)/K)], 1)
def sample_data(n):
    idx = rng.integers(0, K, n)
    return centers[idx] + 0.1*rng.standard_normal((n, 2))

2단계: 잡음 일정과 순방향 과정

DDPM 원 논문은 $T=1000$, $\beta$ 를 $10^{-4}$ 에서 $0.02$ 까지 선형으로 둡니다. 이 장난감 예에서는 계산을 줄이려고 $T=100$ 으로 줄였고, 그 대신 마지막 $\beta$ 를 $0.1$ 로 키워 끝에서 신호가 충분히 사라지게 했습니다(이 수업의 설정, 논문 값 아님). 실행하면 $\bar\alpha_{T}\approx0.0056$, 즉 원본 진폭이 약 7.5%만 남습니다.

12345678910# ---------- 2. 순방향 과정: 선형 beta 일정 ----------
T = 100
betas = np.linspace(1e-4, 0.1, T)
alphas = 1 - betas
abar = np.cumprod(alphas)
print(f"abar[0]={abar[0]:.4f}  abar[T-1]={abar[-1]:.5f}")

def q_sample(x0, t, eps):
    a = abar[t][:, None]
    return np.sqrt(a)*x0 + np.sqrt(1-a)*eps

왜 $T$ 를 줄일 때 $\beta$ 를 키우나. $\bar\alpha_T$ 는 대략 $e^{-\sum\beta_t}$ 입니다. 원 논문 설정은 $\sum\beta\approx10$ 이라 $\bar\alpha_T\approx4\times10^{-5}$ 이지만, $T=100$ 에 같은 $\beta$ 를 쓰면 $\sum\beta\approx1$ 이라 $\bar\alpha_T\approx0.36$ 으로 신호가 많이 남습니다. 그러면 생성의 출발점(순수 잡음)과 학습 때 본 $x_T$ 가 달라집니다(연습 문제 2에서 실제로 확인).

3단계: 시점 임베딩과 MLP

신경망은 "지금 잡음이 얼마나 섞였는지"를 알아야 하므로 $t$ 를 입력에 넣습니다. 정수 하나를 그대로 넣기보다, 주파수가 다른 사인 · 코사인 16개로 펼쳐 넣으면 작은 신경망도 시점마다 다른 행동을 배우기 쉽습니다(트랜스포머 위치 인코딩과 같은 생각). 입력은 점 좌표 2개 + 시점 16개 = 18차원, 은닉층 64칸 3개, 출력은 잡음 2개입니다. 가중치는 He 초기화($\sqrt{2/\text{입력 수}}$)를 씁니다.

1234567891011121314151617181920212223242526272829303132333435# ---------- 3. 잡음 예측 MLP (손으로 쓴 역전파 + Adam) ----------
def t_embed(t):
    # 시점 t 를 사인 · 코사인 16차원으로 바꾸기
    freqs = np.exp(np.linspace(0, np.log(100), 8))
    ang = (t[:, None]/T) * freqs[None, :]
    return np.concatenate([np.sin(ang), np.cos(ang)], 1)

D_IN, H = 2 + 16, 64
def init(i, o):
    return rng.standard_normal((i, o))*np.sqrt(2/i), np.zeros(o)
params = {}
params['W1'], params['b1'] = init(D_IN, H)
params['W2'], params['b2'] = init(H, H)
params['W3'], params['b3'] = init(H, H)
params['W4'], params['b4'] = init(H, 2)

def forward(x, t):
    h0 = np.concatenate([x, t_embed(t)], 1)
    z1 = h0@params['W1'] + params['b1']; h1 = np.maximum(z1, 0)
    z2 = h1@params['W2'] + params['b2']; h2 = np.maximum(z2, 0)
    z3 = h2@params['W3'] + params['b3']; h3 = np.maximum(z3, 0)
    out = h3@params['W4'] + params['b4']
    return out, (h0, z1, h1, z2, h2, z3, h3)

def backward(cache, dout):
    h0, z1, h1, z2, h2, z3, h3 = cache
    g = {}
    g['W4'] = h3.T@dout; g['b4'] = dout.sum(0)
    d3 = (dout@params['W4'].T)*(z3 > 0)
    g['W3'] = h2.T@d3; g['b3'] = d3.sum(0)
    d2 = (d3@params['W3'].T)*(z2 > 0)
    g['W2'] = h1.T@d2; g['b2'] = d2.sum(0)
    d1 = (d2@params['W2'].T)*(z1 > 0)
    g['W1'] = h0.T@d1; g['b1'] = d1.sum(0)
    return g

역전파 읽는 법. 층마다 같은 두 줄입니다. (1) 가중치 기울기 = (그 층 입력)ᵀ × (출력 쪽 기울기), (2) 아래층으로 보낼 기울기 = (출력 쪽 기울기) × Wᵀ 에 ReLU가 켜진 칸만 남기기(z > 0).

4단계: 학습 — L_simple 과 Adam

DDPM 알고리즘 1 그대로입니다. 데이터 $x_0$, 시점 $t$, 잡음 $\epsilon$ 을 뽑아 $x_t$ 를 만들고, 예측 잡음과 실제 잡음의 제곱 오차를 줄입니다. 손실 $\frac1B\sum\|\hat\epsilon-\epsilon\|^2$ 를 출력으로 미분하면 $\frac{2}{B}(\hat\epsilon-\epsilon)$ 입니다. Adam은 학습률 $10^{-3}$, $\beta_1=0.9$, $\beta_2=0.999$ 이고 4,000번째 걸음부터 학습률을 0.3배로 줄였습니다(이 수업의 설정).

12345678910111213141516171819202122232425m = {k: np.zeros_like(v) for k, v in params.items()}
v = {k: np.zeros_like(v) for k, v in params.items()}
lr, b1, b2, ae = 1e-3, 0.9, 0.999, 1e-8

# ---------- 4. 학습: L_simple = || eps - eps_theta(x_t, t) ||^2 ----------
STEPS, B = 6000, 256
t0 = time.time()
for step in range(1, STEPS+1):
    x0 = sample_data(B)
    t = rng.integers(0, T, B)
    eps = rng.standard_normal((B, 2))
    xt = q_sample(x0, t, eps)
    pred, cache = forward(xt, t)
    loss = np.mean(np.sum((pred-eps)**2, 1))
    dout = 2*(pred-eps)/B
    g = backward(cache, dout)
    lr_t = lr if step < 4000 else lr*0.3
    for k in params:
        m[k] = b1*m[k] + (1-b1)*g[k]
        v[k] = b2*v[k] + (1-b2)*g[k]**2
        mh = m[k]/(1-b1**step); vh = v[k]/(1-b2**step)
        params[k] -= lr_t*mh/(np.sqrt(vh)+ae)
    if step % 1000 == 0:
        print(f"step {step:5d}  loss {loss:.4f}")
print(f"학습 시간 {time.time()-t0:.1f}초")

손실이 0으로 가지 않는 이유. 잡음이 많이 섞인 큰 $t$ 에서는 $x_t$ 만 보고 어떤 잡음이 섞였는지 정확히 알 수 없습니다(여러 원본이 같은 $x_t$ 를 만들 수 있음). 그래서 손실은 2차원 잡음의 분산 합 2보다는 한참 작지만 0은 아닌 값(아래 출력에서 0.6~0.9)에서 오르내립니다. 배치마다 뽑는 $t$ 가 달라 값이 들쭉날쭉한 것도 정상입니다.

5단계: DDPM으로 생성하기

DDPM 알고리즘 2입니다. 순수 잡음에서 시작해 $t=T-1$ 부터 0까지 평균을 계산하고, 마지막 걸음을 빼고는 $\sigma_t^2=\beta_t$ 크기의 잡음을 더합니다.

123456789101112# ---------- 5. DDPM 표본 생성 (sigma_t^2 = beta_t) ----------
def sample(n):
    x = rng.standard_normal((n, 2))
    for t in reversed(range(T)):
        tt = np.full(n, t)
        eps_hat, _ = forward(x, tt)
        mean = (x - betas[t]/np.sqrt(1-abar[t])*eps_hat)/np.sqrt(alphas[t])
        if t > 0:
            x = mean + np.sqrt(betas[t])*rng.standard_normal((n, 2))
        else:
            x = mean
    return x

6단계: 평가 — 봉우리를 다 만들었나

생성한 점마다 가장 가까운 중심까지의 거리를 재고, 0.5 안쪽에 들어온 점이 어느 봉우리에 몇 개인지 셉니다. "거리"는 품질(정밀도), "봉우리별 개수"는 다양성(재현율 · 모드 붕괴 여부)을 대신 봅니다. 비교를 위해 진짜 데이터와 순수 잡음도 같은 지표로 잽니다.

1234567891011# ---------- 6. 평가 ----------
def evaluate(x, name):
    d = np.linalg.norm(x[:, None, :]-centers[None], axis=2)
    nearest = d.min(1); which = d.argmin(1)
    good = nearest < 0.5
    counts = np.bincount(which[good], minlength=K)
    print(f"[{name}] 가까운 중심까지 평균 거리 {nearest.mean():.3f} | 0.5 안쪽 비율 {good.mean()*100:.1f}% | 중심별 개수 {counts.tolist()}")

evaluate(sample_data(2000), "진짜 데이터")
evaluate(rng.standard_normal((2000, 2)), "순수 잡음")
evaluate(sample(2000), "확산 모델 표본")

7단계(확장): 같은 모델로 DDIM 20단계

확산 모델 실전 의 DDIM 식($\eta=0$)을 그대로 씁니다. 100개 시점 중 20개만 골라 큰 걸음으로 내려갑니다. 모델은 다시 학습하지 않습니다.

123456789101112131415# ---------- 7. (확장) 같은 모델로 DDIM 20단계 ----------
def sample_ddim(n, steps=20):
    ts = np.linspace(T-1, 0, steps).round().astype(int)
    x = rng.standard_normal((n, 2))
    for i, t in enumerate(ts):
        eps_hat, _ = forward(x, np.full(n, t))
        x0_hat = (x - np.sqrt(1-abar[t])*eps_hat)/np.sqrt(abar[t])
        if i == len(ts)-1:
            x = x0_hat
        else:
            s = ts[i+1]
            x = np.sqrt(abar[s])*x0_hat + np.sqrt(1-abar[s])*eps_hat
    return x

evaluate(sample_ddim(2000), "DDIM 20단계")

실제 실행 결과

위 코드를 한 파일로 이어 붙여 python toy_diffusion.py 로 실행한 출력입니다(시드 0, 학습 시간은 실행할 때마다 다름).

123456789101112abar[0]=0.9999  abar[T-1]=0.00562
step  1000  loss 0.8654
step  2000  loss 0.7179
step  3000  loss 0.7328
step  4000  loss 0.6413
step  5000  loss 0.6885
step  6000  loss 0.7495
학습 시간 163.0초
[진짜 데이터] 가까운 중심까지 평균 거리 0.126 | 0.5 안쪽 비율 100.0% | 중심별 개수 [257, 249, 255, 230, 245, 244, 258, 262]
[순수 잡음] 가까운 중심까지 평균 거리 0.960 | 0.5 안쪽 비율 14.5% | 중심별 개수 [33, 39, 38, 37, 28, 45, 37, 33]
[확산 모델 표본] 가까운 중심까지 평균 거리 0.201 | 0.5 안쪽 비율 94.1% | 중심별 개수 [236, 275, 209, 252, 200, 217, 242, 251]
[DDIM 20단계] 가까운 중심까지 평균 거리 0.183 | 0.5 안쪽 비율 90.9% | 중심별 개수 [251, 254, 217, 269, 200, 185, 222, 220]

결과 읽기.

• 순수 잡음은 14.5%만 봉우리 근처에 있었는데, 학습한 모델로 거꾸로 걷자 94.1%가 봉우리 0.5 안쪽에 들어왔습니다.
• 8개 봉우리 모두 200~275개씩 나왔습니다. 한두 봉우리만 만드는 모드 붕괴는 보이지 않습니다.
• 평균 거리 0.201은 진짜 데이터(0.126)보다 큽니다. 점 구름이 진짜보다 조금 퍼져 있다는 뜻으로, 작은 MLP · 짧은 학습 · 줄인 $T$ 의 한계입니다.
• DDIM 20단계는 신경망 호출을 100번에서 20번으로 줄였는데도 90.9%를 유지했습니다. 다만 봉우리별 개수는 185~269로 조금 더 고르지 않습니다.

이 수치는 이 장난감 설정에서 얻은 것이며 논문 결과와 비교하는 값이 아닙니다.

연습 문제(확장 과제)

1. STEPS 를 1,000으로 줄여 다시 돌리면 결과가 어떻게 바뀔지 예상하고 확인하세요.
   • 풀이: 잡음 예측이 덜 정확해져 품질이 떨어집니다. 이 수업 작성 때 실제로 돌려 보니(시드 0) DDPM 표본의 평균 거리가 0.451, 0.5 안쪽 비율이 57.9%로, 6,000걸음일 때(0.201, 94.1%)보다 크게 나빴습니다.
2. 2단계에서 betas = np.linspace(1e-4, 0.02, T) (원 논문 범위)로 바꾸면 $\bar\alpha_T$ 는 얼마이고, 무엇이 걱정되나요? 실제 결과는?
   • 풀이: 출력은 $\bar\alpha_T\approx0.364$ 입니다. 학습 때 본 $x_T$ 에는 원본 신호가 꽤 남아 있는데 생성은 순수 잡음에서 시작하므로 "학습-추론 불일치"가 생깁니다. 그런데 실제로 돌려 보니 이 장난감 지표로는 오히려 0.5 안쪽 비율 98.4%, 평균 거리 0.158로 더 좋게 나왔습니다. 우리 지표가 "봉우리에 가까운가 · 봉우리별 개수"만 보기 때문에, 신경망이 아무 출발점이나 가까운 봉우리로 끌어당기면 좋은 점수가 나옵니다. 불일치의 영향은 데이터와 지표에 따라 다르며, Lin 외(2023)는 Stable Diffusion에서 마지막 시점의 신호 대 잡음비가 0이 아니어서 생기는 같은 종류의 불일치 때문에 아주 밝거나 아주 어두운 이미지를 잘 못 만든다고 보고했습니다. 장난감 실험의 좋은 점수가 일반 원칙을 뒤집지는 않는다는 교훈입니다.
3. 5단계의 마지막 걸음($t=0$)에 잡음을 더하지 않는 이유는?
   • 풀이: $t=0$ 의 결과가 최종 표본입니다. 여기에 잡음을 더하면 학습한 분포보다 일부러 흐트린 점을 내놓게 됩니다. DDPM 알고리즘 2도 $t=1$ 에서 $z=0$ 으로 둡니다.
4. 평가 함수에 "가장 많이 나온 봉우리 개수 ÷ 가장 적게 나온 봉우리 개수"를 추가하면, 위 출력에서 DDPM과 DDIM 값은?
   • 풀이: DDPM $275/200=1.375$, DDIM $269/185\approx1.454$ 입니다. 진짜 데이터는 $262/230\approx1.139$ 로, 1에 가까울수록 골고루입니다.
5. 조건부 생성(봉우리 번호 $c$ 를 입력)과 분류기 없는 가이던스를 붙이려면 무엇을 바꾸나요?
   • 풀이: 입력에 $c$ 의 원-핫(8차원)과 "빈 조건" 칸을 더하고, 학습 때 확률 $p_{\text{uncond}}$ (예: 0.1)로 조건을 빈 조건으로 바꿉니다. 생성 때는 조건 있음 · 없음 두 번 예측해 $(1+w)\epsilon_c-w\epsilon_u$ 로 섞습니다.

스스로 점검

1. 이 코드에서 "확산 모델 학습"에 해당하는 핵심 세 줄은?
   • 답: xt = q_sample(x0, t, eps) (잡음 섞기), pred, cache = forward(xt, t) (잡음 예측), loss = mean(sum((pred-eps)**2)) (제곱 오차)입니다.
2. 시점 임베딩이 없으면 무슨 문제가 생기나요?
   • 답: 신경망이 지금 잡음이 얼마나 섞였는지 모른 채 같은 함수로 모든 시점을 처리해야 해서, 거의 깨끗한 점과 거의 잡음인 점을 구분해 다르게 고치기 어렵습니다.
3. DDIM 20단계가 DDPM 100단계와 비슷한 품질을 낸 이유는?
   • 답: 같은 잡음 예측 모델로 매 걸음 원본 $\hat x_0$ 를 추정하고 다음 시점의 잡음 몫만 다시 섞는 결정적 걸음이라, 큰 걸음으로 건너뛰어도 경로가 크게 어긋나지 않기 때문입니다.

용어 정리

더 깊이 보려면

• Hugging Face Diffusion Models Course Unit 1 — "Diffusion Models from Scratch" 노트북(PyTorch): https://huggingface.co/learn/diffusion-course/unit1/1
• 확산 정책 — 같은 원리로 2차원 점 대신 로봇 행동 묶음을 생성

다음에 볼 것

• 정규화 흐름과 플로우 매칭 — 같은 장난감 문제를 직선 경로로 푸는 방법
• 생성 모델 평가 — 오늘 쓴 거리 · 개수 지표를 FID · 정밀도/재현율로 일반화

라이선스

본문과 코드는 KAITABASE가 이 수업을 위해 작성했습니다. 코드는 자유롭게 고쳐 써도 됩니다. 알고리즘은 DDPM · DDIM 논문의 식을 따랐고, 데이터 설정 · 하이퍼파라미터 · 실행 결과는 이 수업의 해설용 설정입니다.

출처

• Ho · Jain · Abbeel, Denoising Diffusion Probabilistic Models(2020) 알고리즘 1 · 2: https://arxiv.org/abs/2006.11239
• Song · Meng · Ermon, Denoising Diffusion Implicit Models(2021) 식 (12): https://arxiv.org/abs/2010.02502
• Metz 외, Unrolled GAN(2016) 8개 가우시안 장난감 실험: https://arxiv.org/abs/1611.02163
• Lin 외, Common Diffusion Noise Schedules and Sample Steps are Flawed(2023): https://arxiv.org/abs/2305.08891
• Hugging Face Diffusion Models Course: https://huggingface.co/learn/diffusion-course/unit1/1