Diffusion LLM Part 2: Discrete Diffusion -- 텍스트에 노이즈를 어떻게 추가하나

Part 1에서 우리는 연속 공간에서 작동하는 Diffusion의 원리를 살펴봤습니다. 이미지 픽셀에는 Gaussian 노이즈를 자연스럽게 더할 수 있지만, 텍스트 토큰은 이산(discrete) 데이터입니다. "hello"에 노이즈 0.3을 더하면 무엇이 될까요?

이 글에서는 Diffusion을 이산 공간으로 옮기는 방법을 다룹니다. D3PM의 Transition Matrix에서 시작해, MDLM의 Masked Diffusion까지 -- LLaDA의 직접적인 선조들입니다.

D3PM: 이산 공간의 Diffusion

Austin et al. (2021)은 D3PM(Discrete Denoising Diffusion Probabilistic Models)에서 핵심 질문을 제기합니다: Gaussian 노이즈를 더할 수 없는 이산 데이터에서 forward process를 어떻게 정의할 것인가?

답: Transition Matrix를 사용합니다.

연속 Diffusion에서 Gaussian 노이즈가 하는 역할을, 이산 Diffusion에서는 transition matrix Q_t가 대신합니다. 각 단계 t에서 토큰 x_{t-1}이 x_t로 바뀔 확률을 행렬로 정의합니다:

q(x_t | x_{t-1}) = Cat(x_t; p = x_{t-1} * Q_t)

여기서 Cat은 Categorical 분포이고, Q_t는 K x K 행렬입니다 (K는 vocabulary 크기). Q_t[i][j]는 토큰 i가 토큰 j로 바뀔 확률을 나타냅니다.

연속 Diffusion과의 대응 관계:

Transition Matrix의 세 가지 선택

D3PM은 Q_t를 설계하는 여러 가지 방법을 제안합니다.

Uniform Transition: 모든 토큰이 동일한 확률로 다른 아무 토큰으로 바뀝니다.

Q_t[i][j] = (1 - beta_t) if i == j, beta_t / (K-1) otherwise

이것은 연속 Diffusion의 isotropic Gaussian과 가장 비슷합니다. 하지만 vocabulary가 수만 개인 텍스트에서는 비효율적입니다. "cat"이 랜덤으로 "the"나 "quantum"으로 바뀌는 것은 학습에 유용한 신호를 주지 않습니다.

Gaussian-like Transition: 토큰 embedding 공간에서 가까운 토큰으로 바뀔 확률이 높습니다. 연속 공간의 국소적 노이즈를 이산 공간에서 모방합니다.

Absorbing State Transition: 모든 토큰이 하나의 특수한 "흡수 상태"로 바뀝니다. 이 흡수 상태가 바로 [MASK] 토큰입니다.

Q_t[i][MASK] = beta_t, Q_t[i][i] = 1 - beta_t, Q_t[i][j] = 0 (j != i, j != MASK)

각 단계마다 토큰이 원래 값을 유지하거나(확률 1-beta_t), [MASK]로 바뀝니다(확률 beta_t). 한번 [MASK]가 되면 절대 돌아오지 않습니다. 충분한 단계를 거치면 모든 토큰이 [MASK]가 됩니다.

이 세 번째 방식이 텍스트에 가장 잘 작동합니다. 그리고 이것은 BERT의 Masked Language Modeling과 놀라울 정도로 비슷합니다.

[MASK]와 BERT의 연결

Absorbing state transition을 다시 보면:

• Forward: 토큰들이 점점 [MASK]로 바뀜
• Reverse: [MASK] 위치에 원래 토큰을 예측하여 복원

이것은 BERT의 학습과 거의 같습니다. BERT도 입력의 15%를 [MASK]로 가리고, 원래 토큰을 맞추도록 학습합니다.

하지만 결정적인 차이가 있습니다:

BERT는 고정된 15% masking으로 표현을 학습하는 인코더입니다. Absorbing state Diffusion은 masking 비율을 연속적으로 변화시키며 likelihood를 최적화하는 생성 모델입니다.

이 차이가 왜 중요한지는 Part 3의 LLaDA에서 자세히 다룹니다.

MDLM: Masked Diffusion의 단순화

Sahoo et al. (2024)의 MDLM(Simple and Effective Masked Diffusion Language Models)은 D3PM의 absorbing state 방식을 텍스트에 최적화했습니다. 핵심 기여는 단순화입니다.

Continuous-time 정의: D3PM이 T개의 이산 단계를 사용하는 대신, MDLM은 연속 시간 t (0에서 1 사이)로 forward process를 정의합니다. 시간 t에서 토큰이 [MASK]가 될 확률은:

P(x_t = MASK | x_0 != MASK) = 1 - alpha(t)

alpha(t)는 1에서 0으로 감소하는 noise schedule입니다. t=0이면 모든 토큰이 원본이고, t=1이면 모든 토큰이 [MASK]입니다.

단순한 학습 목표: MDLM의 loss는 놀랍도록 간단합니다.

L = E_{t, x_0}[ -alpha'(t)/alpha(t) * sum_{i: x_t^i = MASK} log p_theta(x_0^i | x_t) ]

풀어 쓰면: "마스킹된 위치에서 원래 토큰을 맞추는 cross-entropy loss에, noise schedule의 가중치를 곱한 것."

BERT의 MLM loss와 비교하면, 가중치만 추가된 형태입니다. 이것이 Diffusion의 ELBO를 만족시킵니다.

핵심: 왜 Masking이 Diffusion인가

Masked Diffusion이 "진짜" Diffusion인 이유를 정리합니다.

Forward process 존재: 원본 텍스트에서 시작해 점진적으로 [MASK]가 늘어남. t가 커질수록 정보가 사라짐. 최종 상태(t=1)는 완전한 [MASK] 시퀀스 = 순수 노이즈에 해당.

Reverse process 학습: 신경망이 [MASK] 위치의 원래 토큰을 예측하며 정보를 복원. 여러 step에 걸쳐 점진적으로 수행.

이론적 보장: 학습 목표가 데이터 로그 가능도의 variational lower bound(ELBO).

연속 Diffusion에서 "깨끗한 이미지 -> Gaussian 노이즈"가 forward process인 것처럼, Masked Diffusion에서 "온전한 텍스트 -> 모든 토큰이 [MASK]"가 forward process입니다.

텍스트 생성 과정

Masked Diffusion으로 텍스트를 생성하는 과정은 직관적입니다:

1. 모든 위치가 [MASK]인 시퀀스로 시작: [MASK] [MASK] [MASK] [MASK] [MASK]
2. 모델이 각 [MASK] 위치에 대해 원래 토큰의 확률 분포를 예측
3. 가장 확신이 높은 위치부터 토큰을 채워 넣음
4. 아직 [MASK]인 위치에 대해 다시 예측
5. 모든 [MASK]가 실제 토큰으로 채워질 때까지 반복

Step 1: [MASK] [MASK] [MASK] [MASK] [MASK]

Step 2: [MASK] [MASK] is [MASK] [MASK]

Step 3: The [MASK] is [MASK] [MASK]

Step 4: The cat is on [MASK]

Step 5: The cat is on mat

이 과정에서 모델은 양방향 context를 활용합니다. Step 3에서 "is"를 이미 알고 있으므로, 왼쪽의 "The"와 오른쪽의 "is"를 모두 참고해 빈칸을 채울 수 있습니다. AR 모델은 이것이 불가능합니다.

Autoregressive와의 결정적 차이

LLaDA로 가는 길

D3PM -> MDLM의 흐름을 정리하면:

D3PM (2021): "Transition Matrix로 이산 Diffusion을 정의할 수 있다. 특히 Absorbing State([MASK])가 텍스트에 효과적이다."

MDLM (2024): "Continuous-time으로 단순화하면 학습이 더 안정적이고, BERT MLM과 유사한 형태의 loss가 된다."

LLaDA (2025): "이 원리를 8B 파라미터까지 스케일링하면 진짜 LLM 수준의 성능이 나온다. Variable masking ratio가 핵심이다."

Part 3에서는 LLaDA가 이 Masked Diffusion을 어떻게 8B 스케일로 확장했는지, 그리고 AR 모델과의 직접 비교 결과를 다룹니다.

핵심 정리

참고 문헌

• Austin et al. "Structured Denoising Diffusion Models in Discrete State-Spaces." NeurIPS 2021.
• Sahoo et al. "Simple and Effective Masked Diffusion Language Models." NeurIPS 2024.
• Lou et al. "Discrete Diffusion Modeling by Estimating the Ratios of the Data Distribution." ICML 2024.
• Devlin et al. "BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding." NAACL 2019.