[2026-1] 정유림 - Simple and EffectiveMasked Diffusion Language Models
paper : https://arxiv.org/abs/2406.07524
Simple and Effective Masked Diffusion Language Models
While diffusion models excel at generating high-quality images, prior work reports a significant performance gap between diffusion and autoregressive (AR) methods in language modeling. In this work, we show that simple masked discrete diffusion is more per
arxiv.org
2024년
260613 기준 : 725회 인용
최근 언어 모델은 대부분 GPT처럼 autoregressive(AR) model, 즉 앞에서부터 한 token씩 순서대로 생성하는 방식에 기반합니다.
이 방식은 생성 성능이 매우 좋지만, 본질적으로 순차 생성이 필요합니다.
반면 diffusion model은 처음부터 완성된 문장을 순서대로 쓰는 것이 아니라, noise가 섞인 상태에서 점진적으로 원래 데이터를 복원하는 방식입니다.
이미지 생성에서는 diffusion model이 큰 성공을 거두었지만, 언어처럼 discrete token을 다루는 영역에서는 autoregressive model보다 성능이 낮다는 평가가 많았습니다. 이 논문은 그 문제의식에서 출발합니다.
논문의 핵심 질문은 다음과 같습니다.
“Discrete diffusion language model은 정말 AR model보다 본질적으로 약한가? 아니면 기존 구현과 학습 objective가 충분히 좋지 않았던 것인가?”
저자들은 후자에 가깝다고 봅니다. 그리고 아주 단순한 masked diffusion language model, 즉 MDLM을 제안합니다.
1. 핵심 아이디어
MDLM은 이름 그대로 masked diffusion language model입니다.
기본 아이디어는 BERT의 masked language modeling과 비슷합니다. 문장 일부를 [MASK]로 가리고, 모델이 원래 token을 맞히도록 학습합니다.
예를 들어 원래 문장이 다음과 같다고 합시다.
Masked diffusion for discrete data
일반적인 BERT-style MLM은 일부 token만 mask합니다.
Masked [MASK] for discrete data
그리고 [MASK] 자리에 들어갈 token을 예측합니다.
MDLM도 겉보기에는 비슷합니다. 하지만 중요한 차이가 있습니다. MDLM은 mask 비율을 고정하지 않고, 다양한 masking rate를 사용합니다.
예를 들어 어떤 경우에는 조금만 mask합니다.
Masked diffusion for discrete [MASK]
다른 경우에는 절반 정도를 mask합니다.
[MASK] diffusion for [MASK] data
또 어떤 경우에는 거의 전부 mask합니다.
[MASK] [MASK] [MASK] [MASK] [MASK]
즉, MDLM은 여러 수준의 masking을 사용해 모델이 원래 문장을 복원하도록 학습합니다.
이 과정을 diffusion 관점에서 보면 다음과 같습니다.
원본 문장
→ 일부 token이 mask됨
→ 더 많은 token이 mask됨
→ 거의 모든 token이 mask됨
이것이 forward diffusion process입니다.
반대로 생성할 때는 완전히 mask된 상태에서 시작해 점점 token을 복원합니다.
[MASK] [MASK] [MASK] [MASK] [MASK]
→ [MASK] diffusion [MASK] [MASK] data
→ Masked diffusion [MASK] discrete data
→ Masked diffusion for discrete data
이것이 reverse diffusion process입니다.
따라서 MDLM은 BERT처럼 mask를 맞히는 모델이면서도, 동시에 diffusion model처럼 생성할 수 있는 모델입니다.
2. “principled”한 objective?
이 논문에서 중요한 표현 중 하나가 principled variational lower bound objective입니다.
말은 어렵지만 핵심은 단순합니다.
일반적인 MLM은 다음과 같은 방법에 가깝습니다.
“문장의 일부를 가리고 맞히도록 학습하면 좋은 representation을 배울 것이다.”
실제로 BERT는 이 방식으로 매우 좋은 성능을 냈습니다.
하지만 일반적인 MLM은 생성 모델의 likelihood를 직접 최적화하는 objective는 아닙니다.
즉, “이 모델이 문장을 생성할 확률을 높이는 방향으로 학습된다”고 깔끔하게 말하기 어렵습니다.
MDLM은 다릅니다.
저자들은 masking 과정을 diffusion process로 정의하고, 그 reverse process를 학습하는 objective를 유도합니다.
이 objective는 diffusion model의 ELBO, 즉 evidence lower bound에서 나옵니다.
ELBO는 생성 모델에서 자주 쓰이는 개념입니다.
모델이 진짜로 최대화하고 싶은 것은 데이터의 likelihood입니다.
log Pθ(x)
하지만 diffusion model처럼 중간 latent step이 많은 모델에서는 이 값을 직접 계산하기 어렵습니다.
그래서 대신 계산 가능한 하한인 ELBO(Evidence Lower Bound)를 최적화합니다.
ELBO ≤ log Pθ(x)
즉, ELBO는 실제 log-likelihood보다 작거나 같은 값입니다.
직관적으로는:
진짜 점수는 직접 못 보지만,
그 점수보다 낮은 계산 가능한 점수를 올린다.
그러면 진짜 점수도 좋아지도록 유도된다.입니다.
ELBO를 높이면 실제 likelihood도 좋아지도록 유도됩니다
MDLM의 흥미로운 점은 이 diffusion ELBO objective가 결과적으로 매우 단순한 형태가 된다는 것입니다.
복잡한 수식을 구현해야 하는 것이 아니라, 실제 학습은 여러 masking rate에서의 MLM loss 평균처럼 수행할 수 있습니다.
즉, 이 논문의 핵심은 다음 문장으로 요약할 수 있습니다.
MDLM은 구현은 MLM처럼 단순하지만, 이론적으로는 diffusion ELBO에 기반한 생성 모델이다.
3. Figure 1
Figure 1은 이 논문의 핵심을 한 장에 담고 있습니다.
왼쪽 그림은 MDLM의 학습 과정을 보여줍니다. 입력 문장 Masked diffusion for discrete data가 있고, diffusion time t에 따라 일부 token이 mask됩니다.
t = 0에 가까울수록 원본 문장에 가깝고, t = 1에 가까울수록 거의 모든 token이 mask됩니다.
모델은 각 masking 상태에서 원래 token을 예측합니다. 그리고 여러 masking rate에서 얻은 unmasking loss를 평균냅니다.
이것이 그림의 제목인:
Diffusion Training: Average of unmasking losses
의 의미입니다.
일반 BERT MLM처럼 고정 mask ratio 하나만 쓰는 것이 아니라, 여러 t, 즉 여러 masking rate에서 loss를 계산합니다.
그래서 MDLM objective는 다음처럼 볼 수 있습니다.
낮은 mask ratio에서의 MLM loss
+ 중간 mask ratio에서의 MLM loss
+ 높은 mask ratio에서의 MLM loss
= 여러 MLM loss의 weighted average논문에서 말하는 핵심이 바로 이것입니다.
일반 BERT MLM은 “가린 단어 맞히기” 훈련이다. MDLM도 겉으로는 가린 단어를 맞히지만, mask 비율을 diffusion time처럼 다루고, 이 복원 과정을 확률적 생성 과정으로 정의한다. 그래서 같은 MLM loss처럼 보여도, MDLM에서는 그것이 diffusion model의 likelihood를 높이는 ELBO objective가 된다.
오른쪽 위에는 MDLM의 특징이 정리되어 있습니다.
첫째, masking rate가 고정되어 있지 않고 random합니다.
둘째, objective가 variational lower bound입니다. 즉, 단순한 MLM heuristic이 아니라 diffusion likelihood objective에서 유도됩니다.
셋째, ancestral sampling을 통해 생성이 가능합니다. 완전히 mask된 문장에서 시작해 점진적으로 token을 복원하는 방식입니다.
넷째, objective는 simple average of MLM losses로 구현할 수 있습니다. 수학적으로는 diffusion objective이지만, 실제 구현은 MLM loss 평균처럼 단순합니다.
오른쪽 아래 그래프는 LM1B benchmark(대규모 텍스트 데이터 기반 모델 평가)에서 perplexity를 비교합니다.
Perplexity는 낮을수록 좋습니다.
기존 diffusion language model들은 AR model보다 성능이 많이 낮았습니다.
예를 들어 Diffusion LM은 118.6, D3PM은 76.9, DiffusionBERT는 63.8, SEDD는 32.8입니다.
반면 MDLM은 23.0까지 낮아집니다. AR baseline은 20.9입니다.
즉, MDLM은 여전히 AR보다 약간 높지만, 기존 diffusion model들과 비교하면 성능을 크게 개선합니다.
Figure 1의 메시지는 명확합니다.
MDLM은 단순한 masked diffusion objective만으로도 기존 diffusion language model보다 훨씬 좋은 성능을 내고, AR model에 가까워질 수 있다.
4. Table 1: LM1B에서의 language modeling 성능
Table 1은 One Billion Words, 즉 LM1B benchmark에서 perplexity를 비교합니다.
이 결과에서 볼 수 있는 것은 세 가지입니다.
첫째, MDLM은 기존 diffusion language model보다 확실히 좋습니다. 같은 33B tokens 기준에서 SEDD가 32.79인데, MDLM은 27.04입니다.
둘째, 학습량을 늘리면 MDLM은 AR model에 더 가까워집니다. 327B tokens에서 MDLM은 23.00이고, AR은 20.86입니다.
셋째, MDLM이 AR을 완전히 이긴 것은 아닙니다. 논문의 주장은 “diffusion이 AR보다 낫다”가 아닙니다. 정확한 주장은 “simple masked diffusion을 잘 설계하면 기존에 생각했던 것보다 훨씬 강하고, AR과의 gap도 크게 줄일 수 있다”입니다.
5. Table 2: OpenWebText에서도 같은 경향이 나타남
Table 2는 OpenWebText에서의 perplexity 결과입니다.
여기서도 AR이 가장 좋습니다. 하지만 MDLM은 SEDD보다 낮은 perplexity를 보입니다.
이 결과는 MDLM의 개선이 LM1B 하나에서만 나온 것이 아니라, 더 큰 자연어 corpus인 OpenWebText에서도 반복된다는 것을 보여줍니다.
즉, MDLM은 기존 diffusion baseline보다 일관되게 나은 language modeling 성능을 보입니다.
6. Table 3: Zero-shot generalization
Table 3은 OpenWebText에서 학습한 모델을 여러 다른 dataset에 바로 평가한 결과입니다.
평가 dataset은 PTB, Wikitext, LM1B, Lambada, AG News, Pubmed, Arxiv입니다. 즉, 학습한 데이터와 다른 분포의 문서에서도 모델이 얼마나 잘 작동하는지 보는 실험입니다.
결과를 보면 MDLM은 모든 dataset에서 SEDD보다 낮은 perplexity를 보입니다.
특히 일부 dataset에서는 MDLM이 AR보다도 낮은 perplexity를 보인다는 것입니다.
7. Table 4: BERT의 representation 성능을 유지하는가?
이 논문에서 중요한 주장은 MDLM이 단순히 generation만 가능한 모델이 아니라, BERT-style encoder의 representation 성능도 유지할 수 있다는 것입니다.
이를 확인하기 위해 Table 4에서는 GLUE benchmark 성능을 비교합니다.
결과적으로 MDLM fine-tuning은 BERT의 GLUE 성능을 떨어뜨리지 않습니다. 오히려 평균 점수는 약간 올라갑니다.
이 결과가 중요한 이유는 다음과 같습니다.
BERT는 원래 representation learning에 강하지만, 생성 모델로 쓰기에는 한계가 있습니다. 반면 AR model은 생성에는 강하지만, GLUE 같은 understanding task에서는 BERT보다 낮은 성능을 보일 수 있습니다.
MDLM은 이 둘 사이의 중간 지점으로, BERT-style encoder-only model의 representation 성능을 유지하면서, diffusion sampling을 통해 generation capability를 추가합니다.
8. Table 5: Semi-autoregressive generation
기존 diffusion language model의 약점 중 하나는 긴 sequence 생성입니다. 보통 diffusion model은 고정 길이 sequence를 생성하는 데 더 자연스럽고, GPT처럼 임의 길이의 text를 계속 생성하는 것은 어렵습니다.
MDLM은 생성 품질도 더 좋고, 속도도 훨씬 빠릅니다. 논문에서는 약 25–30배 빠른 semi-autoregressive decoding을 보인다고 설명합니다.
이 결과는 MDLM이 단순히 이론적으로 generation이 가능한 모델이 아니라, 실제 sampling에서도 효율적일 수 있음을 보여줍니다.
9. Table 6: DNA sequence modeling으로 확장
이 논문은 자연어에서만 MDLM을 평가하지 않습니다. DNA sequence modeling에도 적용합니다.
저자들은 Caduceus라는 DNA language model을 기반으로 MDLM fine-tuning을 수행하고, HG38 human reference genome에서 perplexity를 평가합니다.
여기서도 AR Mamba가 가장 좋은 perplexity를 보입니다. 하지만 diffusion model끼리 비교하면 MDLM이 가장 좋습니다.
Plaid와 SEDD보다 MDLM의 perplexity가 더 낮습니다.
이 결과는 MDLM이 자연어에만 적용되는 아이디어가 아니라, biological sequence 같은 discrete sequence domain에도 확장될 수 있음을 보여줍니다.
10. Table 7: Genomic Benchmarks downstream task
Table 7은 DNA sequence에서 downstream classification task 성능을 평가합니다. 여러 regulatory element classification task에 대해 top-1 accuracy를 비교합니다.
MDLM은 모든 task에서 최고 성능을 내지는 않는다
MDLM은 기존 MLM 기반 representation 성능이 대체로 유지된다. 동시에 generative capability가 추가된다.
이것은 자연어 GLUE 실험인 Table 4와 같은 메시지를 biological sequence domain에서 다시 보여준 것입니다.
11. Table 8: Ablation study
의미
| continuous-time formulation | diffusion step을 1, 2, 3처럼 이산적으로 보지 않고, 연속적인 시간 t ∈ [0,1]로 보는 방식 |
| carry-over unmasking | 이미 복원된 token은 다시 바꾸지 않고 그대로 유지하는 방식 |
| zero masking | 모델이 clean token을 예측할 때 [MASK]를 답으로 내지 못하게 하는 방식 |
여기서 가장 중요한 요소는 carry-over unmasking입니다.
Carry-over unmasking이란 이미 unmasked된 token은 reverse diffusion 과정에서 그대로 유지하는 설계입니다.
예를 들어 현재 상태가 다음과 같다고 합시다.
[MASK] diffusion for [MASK] data
여기서 diffusion, for, data는 이미 보이는 token입니다. MDLM은 이 token들을 다시 바꾸지 않고 그대로 유지합니다. [MASK]인 위치만 새로 예측합니다.
이 설계를 제거하면 perplexity가 27.19에서 28.56으로 크게 나빠집니다.
반면 continuous-time formulation을 제거하면 27.04에서 27.19로 소폭만 나빠집니다. Zero masking을 제거했을 때는 추가 악화가 거의 없습니다.
MDLM의 성능 개선에는 복잡한 이론적 장치보다, carry-over unmasking 구현이 매우 중요하다.
최종 요약
이 논문은 masked language modeling과 diffusion model 사이의 연결을 매우 단순하고 실용적인 방식으로 보여줍니다.
핵심은 다음과 같습니다.
BERT처럼 mask를 맞히는 학습을 한다.
하지만 masking rate를 random하게 바꾼다.
이 과정을 diffusion process로 해석한다.
그 결과 objective는 diffusion ELBO에서 유도된다.
그런데 구현은 MLM loss 평균처럼 단순하다.
MDLM은 기존 diffusion language model보다 훨씬 낮은 perplexity를 달성하고, AR model에 가까운 성능을 보입니다. 또한 BERT-style representation 성능을 유지하면서 generation capability를 추가할 수 있습니다. 더 나아가 DNA sequence modeling에서도 적용 가능성을 보여줍니다.