최익준
모든 글

Diffusion 언어모델은 어떻게 동작하나

토큰을 하나씩 뽑는 대신 문장 전체를 복원하는 방식. 마스크 기반 diffusion 언어모델의 학습 목표와 샘플링 과정을 수식으로 따라가 봅니다.

  • Diffusion LM
  • 생성모델
  • 수식

언어모델이라고 하면 대부분 자기회귀(autoregressive) 모델을 떠올립니다. 앞 토큰들을 조건으로 다음 토큰의 분포를 예측하고, 하나를 뽑고, 다시 반복합니다. Diffusion 언어모델은 이 순서를 버립니다. 문장 전체를 한꺼번에 놓고, 여러 단계에 걸쳐 점점 "덜 망가진" 문장으로 복원해 갑니다.

이 글은 그중 가장 단순하고 지금 가장 널리 쓰이는 형태인 마스크 기반(masked) diffusion 언어모델을 다룹니다.1 연속 공간의 가우시안 diffusion과는 다르게, 이산 토큰 위에서 "지우기"와 "되살리기"만으로 정의됩니다.

자기회귀와 무엇이 다른가

길이 LL인 문장 x=(x1,…,xL)x = (x_1, \dots, x_L)에 대해 자기회귀 모델은 결합 확률을 이렇게 쪼갭니다.

pθ(x)=∏i=1Lpθ(xi∣x<i)p_\theta(x) = \prod_{i=1}^{L} p_\theta(x_i \mid x_{<i})

생성은 왼쪽에서 오른쪽으로 LL번의 순차 호출입니다. 한 번에 한 토큰이고, 앞에서 뽑은 토큰은 되돌릴 수 없습니다.

Diffusion 언어모델은 결합 분포를 순서로 쪼개지 않습니다. 대신 "완전히 가려진 문장"에서 출발해 "원래 문장"에 도달하는 과정을 학습합니다. 그래서 여러 위치를 동시에 채울 수 있고, 어떤 순서로 채울지는 학습이 아니라 샘플링 단계에서 고를 수 있습니다.

정방향 과정: 토큰을 지운다

연속 diffusion이 데이터에 가우시안 노이즈를 섞듯, 마스크 diffusion은 토큰을 특별한 기호 [M][\mathrm{M}]으로 바꿉니다. 시간 t∈[0,1]t \in [0, 1]에서 각 위치가 독립적으로 마스크될 확률을 tt로 두면, 정방향 과정은 위치별로 이렇게 정의됩니다.

q(xt(i)∣x0(i))={x0(i)확률 1−t[M]확률 tq(x_t^{(i)} \mid x_0^{(i)}) = \begin{cases} x_0^{(i)} & \text{확률 } 1 - t \\ [\mathrm{M}] & \text{확률 } t \end{cases}

t=0t = 0이면 원문 그대로, t=1t = 1이면 전부 마스크입니다. 중요한 성질은 두 가지입니다. 첫째, 한번 마스크된 위치는 tt가 커져도 계속 마스크입니다(흡수 상태).1 둘째, 각 위치가 독립이라서 임의의 tt에 대한 xtx_t를 x0x_0에서 한 번에 샘플링할 수 있습니다. 학습 때 매 스텝 tt 하나를 뽑아 바로 마스크하면 됩니다.

역방향 과정: 가려진 자리를 맞춘다

모델 pθp_\theta가 배우는 것은 역방향, 즉 마스크된 문장 xtx_t를 보고 가려진 위치의 원래 토큰을 예측하는 일입니다.

pθ(x0(i)∣xt),i∈{ j:xt(j)=[M] }p_\theta(x_0^{(i)} \mid x_t), \quad i \in \{\, j : x_t^{(j)} = [\mathrm{M}] \,\}

구조는 BERT2와 같은 양방향 트랜스포머입니다. 모든 위치를 동시에 보고, 마스크된 자리마다 어휘 전체에 대한 분포를 내놓습니다. 차이는 마스크 비율이 15%로 고정된 게 아니라 00부터 11 사이에서 매번 달라진다는 점입니다. 거의 다 가려진 입력에서도, 거의 다 드러난 입력에서도 동작해야 합니다.

학습 목표

연속 diffusion의 변분 하한(ELBO)을 마스크 과정에 대해 전개하면 놀랄 만큼 단순한 식이 나옵니다.34

L(θ)=Et∼U(0,1)  Ext∼q(⋅∣x0)[1t∑i: xt(i)=[M]−log⁡pθ ⁣(x0(i)∣xt)]\mathcal{L}(\theta) = \mathbb{E}_{t \sim U(0,1)}\; \mathbb{E}_{x_t \sim q(\cdot \mid x_0)} \left[ \frac{1}{t} \sum_{i:\, x_t^{(i)} = [\mathrm{M}]} -\log p_\theta\!\left(x_0^{(i)} \mid x_t\right) \right]

읽으면 이렇습니다. 마스크 비율 tt를 뽑고, 그 비율로 문장을 가리고, 가려진 위치들의 교차 엔트로피를 더한 뒤, 1/t1/t로 나눕니다.

1/t1/t 가중치가 핵심입니다. tt가 작으면 가려진 토큰이 몇 개 없어서 손실 합이 작아지는데, 이를 1/t1/t로 보정해야 모든 tt에서 손실의 기대 크기가 비슷해집니다. 이 가중치를 빼면 식은 그냥 "마스크 비율을 랜덤하게 바꾼 BERT"가 되고, 더 이상 우도의 하한이 아닙니다. 실제로 가중치를 빼면 perplexity가 눈에 띄게 나빠집니다.3

이 식에는 자기회귀 모델에 없는 장점이 하나 숨어 있습니다. tt가 11에 가까울 때 모델은 거의 빈 문장에서 토큰을 맞춰야 하고, tt가 00에 가까울 때는 문맥이 거의 다 주어진 채 몇 개만 맞춥니다. 한 모델이 "처음 쓰기"와 "빈칸 채우기"를 동시에 배우는 셈입니다.

샘플링: 몇 단계로, 어떤 순서로

생성은 x1x_1(전부 마스크)에서 출발합니다. 단계 수 TT를 정하고 tt를 11에서 00으로 내리면서, 매 단계 다음을 반복합니다.

  1. 현재 xtx_t를 모델에 넣어 마스크된 모든 위치의 분포 pθ(x0(i)∣xt)p_\theta(x_0^{(i)} \mid x_t)를 얻는다.
  2. 각 위치에서 토큰을 하나씩 뽑는다.
  3. 뽑은 토큰 중 일부만 확정하고, 나머지는 다시 [M][\mathrm{M}]으로 돌린다.

3번에서 "일부"를 얼마나, 어떻게 고르느냐가 샘플러의 전부입니다. 비율은 보통 스케줄로 정합니다. tt에서 s<ts < t로 갈 때 마스크로 남길 확률은

p(남김)=stp(\text{남김}) = \frac{s}{t}

이고, 이렇게 하면 역방향 전이가 정방향 과정과 정확히 짝이 맞습니다. 순서는 선택입니다.

전략어떤 위치를 먼저 확정하나특징
무작위임의이론과 가장 가깝고, 품질은 기준선
확신도 우선5예측 확률이 높은 위치부터품질 개선, 반복 문구 증가
확신도 + 노이즈확신도에 Gumbel 노이즈를 더해 순서를 흔듦둘의 절충, 실무에서 가장 흔함

TT가 작을수록 빠르지만 한 단계에서 여러 위치를 동시에 확정하게 되고, 동시에 확정된 토큰끼리는 서로를 보지 못합니다. 그래서 단계 수를 줄이면 문법은 맞는데 앞뒤가 안 맞는 문장이 늘어납니다. "몇 단계까지 줄일 수 있나"가 이 분야의 핵심 질문 중 하나인 이유입니다.

왜 이 방식에 관심을 갖나

  • 병렬 생성. T≪LT \ll L이면 자기회귀보다 호출 횟수가 적습니다. 긴 문장일수록 차이가 큽니다.
  • 양방향 문맥. 문장 중간을 고치거나, 앞뒤가 주어진 빈칸을 채우는 일이 자연스럽습니다. 자기회귀 모델은 이를 위해 별도 학습이 필요합니다.
  • 제어. 특정 위치를 처음부터 고정해 두면 그 제약을 만족하는 문장이 생성됩니다. 분류기 지도(guidance)를 붙이기도 쉽습니다.

반대로 아직 약한 부분도 분명합니다. 같은 파라미터 수에서 perplexity는 자기회귀가 앞서고, 샘플링 단계를 줄이면 품질이 빠르게 떨어지며, 길이를 미리 정해야 합니다. 다만 수십억 파라미터 규모로 키웠을 때 자기회귀 모델과 견줄 만하다는 보고가 나오기 시작했습니다.6

정리

마스크 diffusion 언어모델은 세 줄로 요약됩니다. 정방향은 토큰을 확률 tt로 지우고, 모델은 지워진 자리를 맞추며, 손실은 마스크된 위치의 교차 엔트로피에 1/t1/t를 곱한 것입니다. 생성은 전부 가려진 문장에서 출발해 몇 단계에 걸쳐 일부씩 확정해 가는 과정이고, 그 "몇 단계"와 "어떤 순서"가 품질과 속도를 가릅니다.

수식이 단순한 만큼 변주할 자리도 많습니다. 다음 글에서는 그중 샘플링 단계를 줄이면서 품질을 지키는 방법들을 다룰 생각입니다.

참고문헌

처음 읽는다면 이 순서를 권합니다. 마스크 diffusion의 정식화는 2와 3이 가장 깔끔하고, 1은 그 뿌리, 4는 샘플링 순서에 대한 직관, 5는 규모를 키웠을 때의 결과입니다.

  1. Austin, J., Johnson, D. D., Ho, J., Tarlow, D., & van den Berg, R. (2021). Structured Denoising Diffusion Models in Discrete State-Spaces. NeurIPS. arXiv:2107.03006
  2. Sahoo, S. S., Arriola, M., Schiff, Y., Gokaslan, A., Marroquin, E., Chiu, J. T., Rush, A., & Kuleshov, V. (2024). Simple and Effective Masked Diffusion Language Models. NeurIPS. arXiv:2406.07524
  3. Shi, J., Han, K., Wang, Z., Doucet, A., & Titsias, M. K. (2024). Simplified and Generalized Masked Diffusion for Discrete Data. NeurIPS. arXiv:2406.04329
  4. Chang, H., Zhang, H., Jiang, L., Liu, C., & Freeman, W. T. (2022). MaskGIT: Masked Generative Image Transformer. CVPR. arXiv:2202.04200
  5. Nie, S., Zhu, F., You, Z., Zhang, X., Ou, J., Hu, J., Zhou, J., Lin, Y., Wen, J.-R., & Li, C. (2025). Large Language Diffusion Models. arXiv:2502.09992

Footnotes

  1. 흡수 상태(absorbing state)를 갖는 이산 diffusion은 Austin et al. (2021)의 D3PM에서 정식화되었습니다. arXiv:2107.03006 ↩ ↩2

  2. Devlin et al. (2019). BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. NAACL. arXiv:1810.04805 ↩

  3. Sahoo et al. (2024). Simple and Effective Masked Diffusion Language Models. 본문의 1/t1/t 가중 목표와, 가중치를 뺐을 때의 성능 비교는 이 논문을 따릅니다. arXiv:2406.07524 ↩ ↩2

  4. Shi et al. (2024). Simplified and Generalized Masked Diffusion for Discrete Data. 같은 목표를 독립적으로 유도하고 일반화했습니다. arXiv:2406.04329 ↩

  5. 확신도 순으로 위치를 확정하는 샘플링은 이미지 생성의 MaskGIT에서 왔습니다. Chang et al. (2022). arXiv:2202.04200 ↩

  6. Nie et al. (2025). Large Language Diffusion Models. 80억 파라미터 마스크 diffusion 모델을 처음부터 학습해 비슷한 규모의 자기회귀 모델과 비교했습니다. arXiv:2502.09992 ↩