Diffusion LLM Part 1: Diffusion 기초 -- DDPM에서 Score Matching까지

Diffusion 기반 언어 모델을 이해하려면, 먼저 Diffusion 모델 자체를 이해해야 합니다. 이 글에서는 이미지 생성에서 검증된 Diffusion의 핵심 원리를 다룹니다. 수학이 좀 나오지만, 직관적인 설명을 함께 달아두었으니 수식이 낯설어도 흐름을 따라갈 수 있습니다.

이 글은 Diffusion LLM 시리즈의 첫 번째 편입니다. 시리즈 개요는 Hub 글을 참고하세요.

Diffusion의 핵심 아이디어

Diffusion 모델의 아이디어는 놀랍도록 단순합니다.

1. 깨끗한 데이터에 노이즈를 조금씩 추가해서 완전한 랜덤 노이즈로 만든다 (Forward Process)
2. 그 반대 과정 -- 노이즈에서 깨끗한 데이터를 복원하는 법 -- 을 신경망으로 학습한다 (Reverse Process)

잉크 한 방울을 물에 떨어뜨리면 점차 퍼져서 균일한 색이 됩니다. Forward process는 이 확산 과정입니다. Reverse process는 균일한 색으로 퍼진 물에서 원래의 잉크 방울 모양을 복원하는 것입니다. 물리적으로는 불가능하지만, 신경망이 이 "시간 역행"을 학습할 수 있다는 것이 Diffusion 모델의 핵심 통찰입니다.

Forward Process: 노이즈 추가

Forward process는 원본 데이터 x_0에서 시작해 T 단계에 걸쳐 Gaussian 노이즈를 점진적으로 추가합니다.

각 단계 t에서의 노이즈 추가:

q(x_t | x_{t-1}) = N(x_t; sqrt(1 - beta_t) * x_{t-1}, beta_t * I)

여기서 beta_t는 noise schedule이라 부르는 작은 양수값입니다 (보통 0.0001에서 0.02까지 선형으로 증가). 각 단계마다 원본 신호를 sqrt(1 - beta_t)만큼 줄이고, beta_t만큼의 노이즈를 더합니다.

T가 충분히 크면 (보통 T=1000), x_T는 거의 순수한 Gaussian 노이즈 N(0, I)가 됩니다. 원본 데이터의 정보가 완전히 사라진 상태입니다.

중요한 수학적 트릭이 있습니다. alpha_t = 1 - beta_t, alpha_bar_t = alpha_1 * alpha_2 * ... * alpha_t 로 정의하면, 중간 단계를 모두 건너뛰고 x_0에서 바로 임의의 시점 x_t를 만들 수 있습니다:

q(x_t | x_0) = N(x_t; sqrt(alpha_bar_t) * x_0, (1 - alpha_bar_t) * I)

이것은 reparameterization trick으로 다시 쓸 수 있습니다:

x_t = sqrt(alpha_bar_t) * x_0 + sqrt(1 - alpha_bar_t) * epsilon, epsilon ~ N(0, I)

이 공식이 학습에 핵심적입니다. 매 학습 스텝마다 1부터 T까지 랜덤으로 t를 뽑고, 해당 노이즈 레벨의 x_t를 한 번에 만들 수 있기 때문입니다.

Reverse Process: 노이즈 제거

Reverse process는 forward의 역방향입니다. 순수한 노이즈 x_T에서 시작해 한 단계씩 노이즈를 제거하며 x_0를 복원합니다.

p_theta(x_{t-1} | x_t) = N(x_{t-1}; mu_theta(x_t, t), sigma_t^2 * I)

여기서 mu_theta는 신경망이 예측하는 평균값입니다. 모델은 현재 노이즈가 섞인 상태 x_t와 시간 단계 t를 입력받아, 노이즈를 한 단계 제거한 x_{t-1}의 평균을 예측합니다.

핵심 질문: 모델이 정확히 무엇을 예측하도록 학습해야 할까요?

세 가지 선택지가 있습니다:

• x_0를 직접 예측 (원본 데이터)
• mu를 직접 예측 (reverse step의 평균)
• epsilon을 예측 (추가된 노이즈 자체)

Ho et al. (2020)의 DDPM 논문은 epsilon을 예측하는 것이 가장 잘 작동한다는 것을 발견했습니다. 직관적으로, 모델에게 "원본이 뭐야?"라고 묻는 것보다 "여기 섞인 노이즈가 뭐야?"라고 묻는 것이 더 쉬운 문제라는 뜻입니다.

학습 목표: ELBO에서 Simplified Loss까지

Diffusion 모델의 이론적 학습 목표는 ELBO(Evidence Lower BOund)입니다. VAE와 마찬가지로, 데이터의 로그 가능도 log p(x_0)의 하한을 최대화합니다.

L_ELBO = E[-log p(x_T)] + sum_{t=2}^{T} E[KL(q(x_{t-1}|x_t, x_0) || p_theta(x_{t-1}|x_t))] + E[-log p_theta(x_0|x_1)]

이 식은 세 부분으로 나뉩니다:

• L_T: x_T가 prior N(0,I)에 가까운지 (학습 파라미터 없음)
• L_{t-1}: 각 reverse step이 실제 posterior와 가까운지
• L_0: 최종 복원 품질

수학적으로, q(x_{t-1}|x_t, x_0)는 forward posterior라 부르는데, forward process의 조건부 확률을 Bayes rule로 뒤집은 것입니다. 이것이 "정답"이고, 모델의 예측 p_theta가 여기에 가까워지도록 학습합니다.

Ho et al.의 발견: 이 복잡한 ELBO를 아래의 단순한 loss로 대체해도 (심지어 더 잘) 작동합니다:

L_simple = E_{t, x_0, epsilon}[|| epsilon - epsilon_theta(x_t, t) ||^2]

직역하면: "랜덤 시점 t에서 x_t에 섞인 실제 노이즈 epsilon과, 모델이 예측한 노이즈 epsilon_theta의 차이를 최소화하라."

이것이 DDPM의 학습을 아주 단순하게 만듭니다.

학습 알고리즘:

1. 학습 데이터에서 x_0를 하나 뽑는다
2. 1~T 사이에서 t를 랜덤으로 뽑는다
3. 노이즈 epsilon ~ N(0,I)를 뽑는다
4. x_t = sqrt(alpha_bar_t) * x_0 + sqrt(1 - alpha_bar_t) * epsilon 을 계산한다
5. 모델의 예측 epsilon_theta(x_t, t)와 실제 epsilon의 MSE를 최소화한다

Score Matching과의 연결

DDPM과 거의 같은 시기에, Song과 Ermon의 Score-based 모델이 다른 경로로 같은 목적지에 도달했습니다.

Score function은 데이터 분포의 로그 확률에 대한 기울기(gradient)입니다:

s(x) = nabla_x log p(x) = (1/p(x)) * nabla_x p(x)

수학적으로 풀어 쓰면: p(x)는 데이터의 확률 밀도 함수이고, nabla_x는 x에 대한 편미분입니다. log p(x)를 미분하면 (1/p(x)) * nabla_x p(x)가 됩니다. 이것은 "현재 x 위치에서, 어느 방향으로 움직이면 확률 밀도가 가장 빠르게 올라가는가"를 가리키는 벡터입니다.

직관적으로: 2차원 공간에 데이터 포인트들이 뿌려져 있다고 상상해보세요. 어떤 지점은 데이터가 밀집해 있고(높은 확률), 어떤 지점은 거의 비어 있습니다(낮은 확률). Score function은 현재 위치에서 "데이터가 밀집된 방향"을 가리키는 화살표입니다. 노이즈에 오염된 위치에서 이 화살표를 따라가면, 원본 데이터에 가까워집니다.

이것이 Diffusion의 reverse process와 연결됩니다. 노이즈가 섞인 x_t에서 score function 방향으로 한 걸음씩 이동하면, 점진적으로 깨끗한 x_0에 도달합니다.

Score matching의 핵심 아이디어: 데이터 분포 p(x)를 직접 모델링하는 대신, score function s(x)를 학습합니다. p(x)를 직접 다루려면 정규화 상수 Z = integral p_unnormalized(x) dx를 계산해야 하는데, 고차원에서 이 적분은 사실상 불가능합니다. 하지만 log p(x)의 기울기를 구하면 Z가 상수이므로 미분 시 사라집니다. 이것이 score function이 학습하기 쉬운 근본적 이유입니다.

Song et al. (2021)은 DDPM의 forward process를 SDE(Stochastic Differential Equation)로 일반화했습니다:

dx = f(x, t)dt + g(t)dw (Forward SDE)

여기서 f는 drift coefficient, g는 diffusion coefficient, w는 Wiener process(Brownian motion)입니다.

놀라운 결과: 이 forward SDE에 대응하는 reverse-time SDE가 존재하며, score function만 알면 역방향으로 시간을 돌릴 수 있습니다:

dx = [f(x,t) - g(t)^2 * nabla_x log p_t(x)]dt + g(t)dw_bar (Reverse SDE)

DDPM의 epsilon-prediction과 score matching은 본질적으로 같은 것을 다른 언어로 표현한 것입니다:

nabla_x log p_t(x) = -epsilon / sqrt(1 - alpha_bar_t)

노이즈 예측 = 스코어 추정. 같은 동전의 양면입니다.

왜 텍스트에는 그대로 적용할 수 없는가

지금까지 설명한 Diffusion 모델은 연속(continuous) 공간에서 작동합니다. 이미지 픽셀은 0~255 사이의 실수값이고, Gaussian 노이즈를 자연스럽게 더하고 뺄 수 있습니다.

하지만 텍스트는 이산(discrete) 토큰입니다. "cat"이라는 단어에 Gaussian 노이즈를 더하면 무엇이 될까요? "cat" + 0.3은 의미가 없습니다.

이 문제를 해결하는 두 가지 접근이 있습니다:

1. 토큰을 연속 공간(embedding)으로 옮긴 뒤 Diffusion 적용 (Continuous-space approach)
2. 이산 공간에서 직접 Diffusion 정의 (Discrete Diffusion)

LLaDA가 채택한 것은 두 번째 접근입니다. Gaussian 노이즈 대신 token masking을 사용합니다. 이것이 Part 2의 주제입니다.

핵심 정리

다음 글 예고

Part 2에서는 이 연속 공간의 Diffusion을 이산 토큰 세계로 옮기는 방법을 다룹니다. D3PM의 Transition Matrix, Absorbing State([MASK] 토큰), 그리고 MDLM이 어떻게 이 모든 것을 텍스트 생성에 실용적으로 만들었는지 살펴봅니다.

참고 문헌

• Ho, Jain, Abbeel. "Denoising Diffusion Probabilistic Models." NeurIPS 2020.
• Song et al. "Score-Based Generative Modeling through Stochastic Differential Equations." ICLR 2021.
• Sohl-Dickstein et al. "Deep Unsupervised Learning using Nonequilibrium Thermodynamics." ICML 2015.
• Luo. "Understanding Diffusion Models: A Unified Perspective." arXiv:2208.11970, 2022.