BERT에서 Diffusion으로: Masked Diffusion Language Model(MDLM) 한 번에 이해하기

TL;DR
한 줄 요약: MDLM은 BERT의 encoder‑only 뼈대에 diffusion을 촘촘히 엮어, 이해(understanding) 능력은 살리고, 문장 생성까지 열어 준 모델입니다. noising은 **mask로 흡수(absorbing)**되는 간단한 규칙, denoising은 mask로 되돌아가지 않는 두 가지 안전장치(zero‑masking probability, carryover unmasking)로 굴러갑니다. 복잡했던 ELBO는 수학적 정리를 거쳐 BERT식 MLM loss의 가중합으로 깔끔히 정리되고, **perplexity(PPL)**는 기존 diffusion LM 대비 최고 수준, 강력한 **autoregressive(AR)**에도 근접합니다. 보너스로 semi‑autoregressive block 생성이 됩니다.
왜 이 글인가
화려한 용어 대신 핵심 그림을 보여주고 싶었습니다. *“BERT가 그대로 글도 쓰면 좋겠다”*라는 마음이 있었다면, MDLM은 꽤나 기분 좋은 해답입니다. 이 글은 문장을 짧게, 리듬감 있게 가져가고, 필요한 영어 용어는 그대로 씁니다. 연구자, 프로덕트 엔지니어, 그리고 호기심 많은 분 모두를 위한 가벼운 가이드가 목표예요.
2024년 6월 헤드라인 포인트
- BERT + diffusion = generation. BERT의 encoder‑only 성격을 유지하면서 generative 능력을 붙였습니다.
- 고정 15% mask → 유동적 masking. BERT가 고정 15% masking을 쓰는 반면, MDLM은 time‑dependent, token‑wise masking을 택합니다.
- Semi‑autoregressive 생성. 토큰을 하나씩(AR) 혹은 한 번에(NAR)가 아니라 block 단위로 씁니다.
- ELBO 단순화. 복잡한 ELBO가 weighted MLM loss로 정리됩니다.
- Perplexity 경쟁력. diffusion 기반 LM 중 SOTA 급 PPL, 강력한 AR과도 격차가 작습니다.
요컨대, 이해는 BERT처럼 유지하고, 생성은 MDLM답게.
Forward(noising): mask로 빨려 들어가는 absorbing 과정
우리가 아는 diffusion은 보통 깨끗한 데이터 x에 noise를 얹는 상상으로 시작하죠. MDLM은 텍스트라는 이crete 세상에 딱 맞는 방식으로 단순화했습니다. **[MASK]**라는 흡수 지점으로 **점진적 흡수(absorbing)**를 하는 구조예요.
- α_t는 scalar schedule, t ∈ [0,1].
- t=1이면 α_t=0, 즉 **모든 token이 [MASK]**가 됩니다.
- 각 위치 i는 고유 **threshold τ_i ∈ (0,1)**를 가집니다.
- 시간이 흐르면서 1−α_t > τ_i가 되는 순간, 그 token은 **[MASK]**로 뒤집히고 이후 계속 mask 상태로 남습니다.
감 잡기 예시. 위치가 1, 2, 3인 세 token을 생각해 봅시다. τ₁=0.5, τ₂=0.3, τ₃=0.3. 1−α_t가 0.5를 넘으면 token 1이 먼저 mask, 이후 0.3을 넘으면 token 2, 3도 mask. 결국 t=1에는 **전부 [MASK]**로 흡수됩니다.
중요한 귀결. 임의의 시점에서 각 latent token은 두 가지 중 하나예요.
- 원래 token의 one‑hot 벡터, 또는
- **[MASK]**의 one‑hot 벡터.
가우시안 구름 같은 중간 상태가 없습니다. 이 세계는 선명합니다: 진짜 token 아니면 mask.
One‑hot embedding + mask 차원
Vocabulary가 V라고 할 때, **[MASK]**를 위해 +1 차원을 더해 **(V+1)**개의 one‑hot feature를 씁니다. 일반 token은 자기 인덱스에 1, **[MASK]**는 마지막 칸이 1인 one‑hot. 이때 저자들이 강조한 포인트:
어떤 clean token x도 mask one‑hot m과의 내적 x·m = 0.
두 상태가 직교하니 noising/denoising이 깔끔하게 분리됩니다.
Reverse(denoising): 간단하지만 강력한 두 규칙
복원 과정은 mask를 걷어내며 실제 token을 드러냅니다. 핵심은 두 가지:
A) Zero‑masking probability
예측 시 [MASK] logit을 **−∞**로 밀어 넣습니다. softmax를 통과하면 [MASK] 확률은 정확히 0. 즉, 한 번 드러난 자리를 다시 mask로 돌아가게 두지 않기.
B) Carryover unmasking
어떤 위치가 한 번 unmask되면 끝까지 unmask 유지. 이 단조 규칙 덕분에 계산량이 줄고, 추론 경로가 안정됩니다. *“확신이 생겼다면 그대로 밀고 가자”*는 태도죠.
A와 B를 묶으면, denoising은 mask → real로만 움직이는 단방향 트랙이 됩니다.
Low‑discrepancy timestep 샘플링(Sobol): 구간을 고르게, 분산은 낮게
Diffusion 학습에서 timestep을 뽑는 방식은 의외로 중요합니다. 무작위 샘플링은 한쪽으로 몰릴 수 있죠. MDLM은 low‑discrepancy 시퀀스(예: Sobol)로 [0,1] 구간을 균등히 커버합니다.
- 장점: ELBO 추정의 분산 감소, 학습 곡선이 차분해집니다.
- 보너스: low dimension에서 특히 효과적인데, one‑hot 위주의 MDLM 세팅과 잘 맞아떨어집니다.
Objective: ELBO가 MLM처럼 들리기 시작할 때
원래 ELBO는 복잡해 보이죠. 그런데 MDLM은 연속 시간 관점에서 정리해 두 항을 없애고, 남은 항을 BERT식 MLM loss의 가중합으로 깔끔하게 바꿉니다.
- t=0의 reconstruction 항은 사실상 사라지고,
- t=1의 prior 항도 **모든 token이 [MASK]**라는 설계 덕에 0이 됩니다.
- 결국 남는 diffusion 항은 time‑aware weight를 얹은 MLM loss로 귀결됩니다.
해석: MLM을 시간 가중치로 훈련하는 느낌. 단순함이 성능으로 이어지는 순간입니다.
Noise schedule: 성능엔 둔감, 안정성엔 민감
흥미로운 이론 결과가 있습니다. **γ = log(1 − α_t)**로 치환하고 적분 정리를 거치면, noise schedule의 모양이 최종 품질 식에서 소거됩니다. 실제로 보면:
- 최종 성능은 합리적인 schedule들 사이에서 큰 차이가 없습니다.
- 대신 training stability와 loss variance에는 영향을 줍니다.
즉, schedule은 성능 미세 튜닝보다 학습을 매끈하게 만들기 위해 고르면 됩니다.
결과 한눈에
단출한 규칙만으로도 MDLM은 이렇게 보고합니다.
- Perplexity(PPL): 기존 diffusion LMs 중 최고 수준, 강력한 AR과도 근접.
- Semi‑autoregressive generation: block 생성에서 SSD‑LM 등 diffusion 계열 대비 우수.
- Zero‑shot: **OpenWebText(OWT)**로만 학습해도 LAMBADA, WinoGrande, ARC 등 OOD 벤치마크에서 선전, 간혹 우위.
- Downstream BERT tasks: 이해 중심 과제에서 BERT와 대등 수준을 유지하면서 생성 능력을 획득.
- 감 잡는 숫자: 어떤 코퍼스 기준 AR ≈ 22 PPL, MDLM ≈ 35 PPL 언급이 있습니다(세부 설정은 다름). 중요한 건 단일 숫자가 아니라 diffusion과 AR의 간극을 크게 좁혔다는 흐름입니다.
핵심은 간명합니다. BERT의 장점은 그대로, 거기에 생성의 목소리를 더했습니다.
Semi‑autoregressive decoding: 토큰과 전체 사이, 그 중간의 길
생성은 이분법일 필요가 없습니다. MDLM은 **semi‑autoregressive(SAR)**로 block 단위 생성을 합니다.
- Autoregressive: 한 번에 1 token(느리지만 의존성 강함).
- Non‑autoregressive: 한 번에 전체(빠르지만 품질 기복).
- MDLM SAR: 문장 덩어리(예: “나는 학교에” → block decode → “갔다 오늘”). 속도와 일관성 사이의 균형점을 노립니다.
실험에서는 이 접근이 SSD‑LM 같은 diffusion baseline을 앞서는 결과를 보입니다.
Ablation: 무엇이 실제로 기여했나
기능을 빼보면 사실이 드러납니다.
- Zero‑masking probability 또는 carryover unmasking을 제거 → 성능 유의미 하락.
- Continuous time에서 reverse step을 늘릴수록 → 점진적 향상.
메시지는 분명합니다. 간단한 규칙, 큰 효과.
Walk‑through: 손에 잡히는 미니 시뮬레이션
세 자리 문장을 떠올려 보죠(위치 1, 2, 3). threshold는 각각 0.5, 0.3, 0.3.
- Forward(absorbing): 1−α_t가 0.5를 넘으면 token 1이 먼저 [MASK]. 이후 0.3을 넘길 때 token 2, 3도 mask. 마지막엔 전부 [MASK].
- Reverse(denoising): 모델이 실제 token을 채워 넣습니다. 한 번 unmask된 위치는 끝까지 유지. [MASK] 확률은 항상 0.
- 마무리: 문장이 매끈히 복원됩니다. SAR 설정이라면 block으로 나눠 한 번에 몇 토큰씩 생성하며 속도/품질을 조율합니다.
이 밑그림이 MDLM의 심장입니다.
실무 관점에서 MDLM이 바꾸는 것들
- BERT 애호가라면: 분류/검색/QA 같은 understanding 작업을 유지한 채 generation을 붙일 수 있습니다.
- 콘텐츠 UX를 설계한다면: SAR block은 latency와 coherence의 신선한 절충안입니다.
- 대규모 학습을 돌린다면: weighted MLM objective, low‑variance timestep sampling, mask‑only 규칙 덕에 파이프라인이 단순해집니다.
- 강건성이 걱정이라면: OWT로만 학습해도 OOD에서 준수한 zero‑shot 성질을 확인할 수 있습니다.
FAQ (검색 친화 요약)
Masked diffusion language model이 뭔가요?
[MASK] 상태로 diffuse해 들어가고, 그 흐름을 reverse하며 token을 드러내는 텍스트 모델입니다.
BERT와 뭐가 다른가요?
BERT는 고정 15% masking + non‑generative. MDLM은 time‑varying, token‑wise masking을 쓰고 generation(특히 SAR)을 지원합니다.
Autoregressive LM을 대체하나요?
아니요. AR은 여전히 강력합니다. 다만 MDLM은 PPL 격차를 줄이고, latency/robustness/objective 단순성 같은 새로운 장점을 제안합니다.
왜 [MASK] 확률을 0으로 막나요?
backtracking을 금지해 reverse chain을 단순화합니다. 한 번 드러난 토큰은 되돌리지 않습니다.
Noise schedule은 중요하지 않다는데요?
최종 성능엔 둔감하지만, 안정성과 loss variance엔 영향을 줍니다. 학습이 부드럽게 가는 쪽을 고르세요.
Sobol 같은 low‑discrepancy를 왜 쓰죠?
timestep을 균일 커버해 ELBO 추정의 variance를 줄입니다. one‑hot 중심의 MDLM 세팅에 특히 잘 맞습니다.
Zero‑shot도 되나요?
네. OWT만으로도 LAMBADA/WinoGrande/ARC 등에서 준수한 결과를 보입니다.
BERT의 downstream 성능은 유지되나요?
네. 보고된 수치는 BERT와 비슷한 이해 성능을 보이면서 generation을 추가로 얻었다고 말합니다.
마무리: 단순함이 주는 품격
MDLM의 매력은 억지로 BERT를 다른 존재로 만들지 않는 데 있습니다. 그저 조용히 귓속말을 건넵니다. “mask로 소음을 흡수하고, 의미는 천천히 돌아오게 해.” 수식은 가벼워지고, objective는 익숙해지고, sampler는 고르게 숨 쉽니다. 후회 없는 generation—현장에서 오래 함께할 만한 태도죠.
제품을 만드는 사람에게 이 감성은 중요합니다. block 단위로 성능과 속도를 맞추고, 학습은 안정적이며, 이해 성능은 견고합니다. 끝은 아니지만, 인간 친화적인 다음 걸음임은 분명합니다.