필사 모드: 모두를 위한 AI 6편 — 111만 파라미터 디퓨전으로 단어에서 숫자 그리기

• 들어가며 — 5편이 남긴 숙제
• 핵심 아이디어 — 망가뜨리기는 쉽고 되돌리기는 어렵다
  • Forward — 노이즈를 섞는 과정
  • 학습 목표 — 이미지가 아니라 노이즈를 맞힌다
• 조건 붙이기 — 단어로 그림 지시하기
• Reverse — 400번 거꾸로 밟기
• 학습 로그
• 결과
• 정리
• 🧠 이해도 체크 퀴즈
• 참고 자료

들어가며 — 5편이 남긴 숙제

5편의 컬러화는 형태를 잘 살렸지만 색이 바랬습니다. 원인은 모델이 아니라 손실이었습니다. 흑백 자동차는 빨강일 수도 파랑일 수도 있는데, L1은 "절대 오차를 최소화하라"고 지시하므로 여러 정답의 중간값인 회색을 찍는 것이 가장 안전한 전략이 됩니다.

이번 편의 문제도 정답이 여럿입니다. "three"라는 단어에 대응하는 손글씨 3은 수천 가지입니다. 만약 5편처럼 회귀 손실로 "3 이미지를 직접 예측하라"고 시키면, 모든 3의 평균인 흐릿한 얼룩이 나올 것입니다.

디퓨전은 이 문제를 정면으로 피해 갑니다. 한 번에 정답을 만들지 않고, 노이즈에서 조금씩 걷어내는 과정으로 바꿉니다.

핵심 아이디어 — 망가뜨리기는 쉽고 되돌리기는 어렵다

깨끗한 이미지에 노이즈를 조금씩 섞어 완전한 잡음으로 만드는 것은 쉽습니다. 공식 한 줄이면 됩니다. 어려운 건 그 반대 방향입니다.

디퓨전의 발상은 이렇습니다. 어려운 문제를 쉬운 문제 400개로 쪼갠다. 잡음에서 이미지를 한 번에 만드는 것은 어렵지만, "아주 조금 더 깨끗하게 만들기"는 배울 만합니다. 그걸 400번 반복합니다.

Forward — 노이즈를 섞는 과정

T = 400
beta = torch.linspace(1e-4, 0.02, T)   # 각 단계에서 섞을 노이즈의 양
alpha = 1 - beta
abar = torch.cumprod(alpha, 0)         # 누적곱

abar[t]는 "시각 t까지 원본이 얼마나 남아 있는가"를 나타냅니다. t=0에서는 거의 1(원본 그대로), t=399에서는 거의 0(완전한 잡음)입니다.

여기서 디퓨전의 첫 번째 장치가 나옵니다. 400단계를 순서대로 밟지 않고 임의의 t로 한 번에 점프할 수 있습니다.

t = torch.randint(0, T, (256,), device=dev)
eps = torch.randn_like(x0)                          # 표준 정규분포 노이즈
a = abar[t][:, None, None, None]
xt = a.sqrt() * x0 + (1 - a).sqrt() * eps           # 한 줄로 t단계 노이즈 적용

이 한 줄 덕분에 학습이 효율적입니다. 배치마다 서로 다른 t를 뽑아 모든 시각을 병렬로 학습시킬 수 있습니다.

학습 목표 — 이미지가 아니라 노이즈를 맞힌다

loss = F.mse_loss(model(xt, t.float(), w), eps)

모델이 예측하는 것은 원본 이미지 x0가 아니라 섞여 들어간 노이즈 eps 입니다.

수학적으로는 둘이 동등합니다. xt와 t를 알면 eps에서 x0를 계산할 수 있고 반대도 성립합니다. 하지만 실제로는 노이즈 예측이 훨씬 잘 작동합니다.

이유는 목표의 성질에 있습니다. x0를 예측하라고 하면, t가 클 때(거의 잡음일 때) 모델은 사실상 정보 없이 이미지를 지어내야 하고 여기서 5편의 평균 회귀가 다시 발생합니다. 반면 eps는 어떤 t에서든 항상 표준 정규분포입니다. 목표의 스케일과 분포가 일정하니 학습이 안정적입니다.

조건 붙이기 — 단어로 그림 지시하기

시각 t와 단어를 같은 벡터 공간에서 더해 U-Net 중앙에 주입합니다.

class CondUNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.temb = nn.Sequential(nn.Linear(64, 128), nn.GELU(), nn.Linear(128, 128))
        self.wemb = nn.Embedding(10, 128)          # zero~nine
        self.e1 = block(1, 48); self.e2 = block(48, 96); self.e3 = block(96, 192)
        self.cond3 = nn.Linear(128, 192)
        self.d2 = block(192 + 96, 96); self.d1 = block(96 + 48, 48)
        self.out = nn.Conv2d(48, 1, 1)

    def forward(self, x, t, w):
        half = torch.exp(-np.log(10000) * torch.arange(32, device=x.device) / 32)
        te = torch.cat([torch.sin(t[:, None] * half), torch.cos(t[:, None] * half)], 1)
        c = self.temb(te) + self.wemb(w)           # 시각 + 단어
        s1 = self.e1(x)
        s2 = self.e2(F.max_pool2d(s1, 2))
        h = self.e3(F.max_pool2d(s2, 2)) + self.cond3(c)[:, :, None, None]
        ...

시각 임베딩에 사인·코사인을 쓰는 것은 트랜스포머의 위치 인코딩과 같은 발상입니다. 정수 t를 그대로 넣으면 신경망이 다루기 어려운 스케일이지만, 서로 다른 주파수의 사인파로 펼치면 인접한 시각끼리 비슷한 벡터를 갖게 됩니다.

주입 위치가 가장 깊은 층(e3 출력) 인 것도 의도적입니다. 해상도가 8×8로 가장 작아 채널당 정보 밀도가 높고, 여기서 조건을 걸면 이후 디코더 전체가 그 영향을 받습니다.

self.cond3(c)[:, :, None, None] 의 None 두 개는 (배치, 채널) 벡터를 (배치, 채널, 1, 1)로 만들어 공간 차원 전체에 브로드캐스트하기 위한 것입니다. 즉 조건이 모든 위치에 동일하게 더해집니다.

Reverse — 400번 거꾸로 밟기

@torch.no_grad()
def sample(words):
    x = torch.randn(len(words), 1, 28, 28, device=dev)      # 순수 잡음에서 시작
    w = torch.tensor([WORDS.index(s) for s in words], device=dev)
    for ti in reversed(range(T)):                            # 399 → 0
        t = torch.full((len(words),), ti, device=dev, dtype=torch.float32)
        eps = model(x, t, w)                                 # 노이즈 예측
        a, ab = alpha[ti], abar[ti]
        x = (x - (1 - a) / (1 - ab).sqrt() * eps) / a.sqrt() # 한 단계 제거
        if ti > 0:
            x = x + beta[ti].sqrt() * torch.randn_like(x)    # 새 노이즈 재주입
    return ((x.clamp(-1, 1) + 1) / 2).squeeze(1).cpu().numpy()

마지막 줄의 노이즈 재주입이 이상해 보일 수 있습니다. 애써 걷어낸 노이즈를 왜 다시 넣을까요.

그것이 다양성의 원천이기 때문입니다. 매 단계 새 난수를 넣으므로 같은 단어를 넣어도 매번 다른 3이 나옵니다. 5편에서 회귀가 하나의 평균값에 갇혔던 것과 정반대입니다. 디퓨전은 "여러 정답 중 하나를 확률적으로 골라내는" 구조입니다. 마지막 단계(ti == 0)에서만 재주입을 생략해 깨끗한 결과를 얻습니다.

학습 로그

[    9.2s] params {"params": 1111777, "T": 400, "n": 60000}
[   10.1s] step      0 mse=1.1502
[   12.2s] step    100 mse=0.0807
[  959.7s] step  44200 mse=0.0324
[  960.6s] SUMMARY {"steps": 44215, "final_mse": 0.0320}

100스텝, 2초 만에 MSE가 1.15에서 0.081로 떨어졌습니다. 예측 대상이 표준 정규분포라 초기 손실이 1 근처에서 시작하는 것도 자연스럽습니다. 이후 완만하게 0.032까지 내려갔습니다.

결과

각 행은 "zero"부터 "nine"까지 열 단어이고, 두 행은 같은 단어에 대한 서로 다른 샘플입니다.

1행은 0부터 9까지 전부 정확합니다. 획의 두께와 기울기도 손글씨답습니다.

2행에서 두 개가 틀렸습니다. "five" 자리에 9에 가까운 형태가, "nine" 자리에 7에 가까운 형태가 나왔습니다. 10개 중 8개 정확, 80%입니다.

두 행을 비교하면 디퓨전의 성질이 드러납니다. 같은 단어인데 글씨체가 다릅니다. 1행의 2와 2행의 2는 곡선의 굽은 정도가 다르고, 1행의 7과 2행의 7은 가로획 유무가 다릅니다. 평균으로 뭉개지지 않고 매번 하나의 구체적인 손글씨를 만들어 낸다는 뜻입니다. 5편의 바랜 색과 대비되는 지점입니다.

틀린 5와 9는 형태적으로 서로 닮은 숫자입니다. 111만 파라미터, 16분 학습으로는 이 정도 혼동이 남습니다.

정리

핵심을 세 줄로 요약하면 이렇습니다.

디퓨전은 어려운 생성 문제를 쉬운 노이즈 제거 400개로 쪼갭니다. forward는 공식 한 줄이고 학습은 임의의 t로 점프해 병렬화됩니다.

모델은 이미지가 아니라 노이즈를 예측합니다. 목표가 어떤 t에서든 표준 정규분포로 일정하기 때문에 학습이 안정적이고, x0를 직접 예측할 때 발생하는 평균 회귀를 피합니다.

reverse 과정의 노이즈 재주입이 다양성을 만듭니다. 같은 조건에서도 매번 다른 결과가 나오는 것이 회귀 모델과의 결정적 차이입니다.

다만 이번 구현은 기본형입니다. 400단계를 전부 밟아야 해서 느리고, 조건을 얼마나 강하게 걸지 조절할 수단도 없습니다. 시리즈 11편에서 이 두 가지를 해결하는 Classifier-Free Guidance와 DDIM을 다루면서, 두 기법이 서로 독립이 아니라는 것도 실측으로 보게 됩니다.

🧠 이해도 체크 퀴즈

참고 자료

• MNIST 데이터셋
• Denoising Diffusion Probabilistic Models (DDPM)
• Diffusion Models Beat GANs on Image Synthesis
• Attention Is All You Need — 사인·코사인 위치 인코딩