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 performant than previously thought. We apply an effective training recipe that improves the performance of masked diffusion models and derive a simplified, Rao-Blackwellized objective that results in additional improvements. Our objective has a simple form -- it is a mixture of classical masked language modeling losses -- and can be used to train encoder-only language models that admit efficient samplers, including ones that can generate arbitrary lengths of text semi-autoregressively like a traditional language model. On language modeling benchmarks, a range of masked diffusion models trained with modern engineering practices achieves a new state-of-the-art among diffusion models, and approaches AR perplexity. We provide the code, along with a blog post and video tutorial on the project page: https://s-sahoo.com/mdlm

한 줄 요약

마스킹된 디퓨전 언어 모델(MDLM)이 기존 AR 모델과 유사한 성능을 보인다.

핵심 기여도

마스킹된 디스크리트 디퓨전 모델이 기존보다 높은 성능을 보임.
Rao-Blackwellized 목적함수를 도출하여 성능 개선.
BERT 스타일 인코더 모델에 생성 능력을 부여.
D3PM 모델을 재구현하여 기존보다 퍼플렉시티 0.1 개선.

핵심 아이디어

기존 연구에서는 디퓨전 모델이 언어 모델링에서 자동회귀(AR) 모델보다 성능이 낮다고 보고되었으나, 본 연구는 단순한 마스킹 디퓨전 언어 모델(MDLM)이 기존 성능 격차를 줄일 수 있음을 보인다. MDLM은 마스킹된 언어 모델링(MLM) 손실의 가중 평균 형태의 목적함수를 사용하며, 이는 BERT 스타일 인코더 모델에 생성 능력을 부여할 수 있다. 또한, SUBS 파라미터화를 통해 역 디퓨전 과정을 단순화하고, Rao-Blackwellized ELBO를 도출함으로써 더 정밀한 하한을 얻는다. 이는 기존 D3PM 모델이 성능이 낮다고 여겨졌지만, 재구현을 통해 기존 최고 성능 모델과 유사한 퍼플렉시티를 달성함을 보여준다.

기술적 접근법

**MDLM**: 마스킹된 디스크리트 디퓨전 언어 모델.
**Rao-Blackwellized ELBO**: 역 디퓨전 과정을 위한 목적함수.
**SUBS 파라미터화**: 단순한 대체 기반 파라미터화.
**T=1000**: 디퓨전 단계 수.
**Zero masking**: 마스킹 연산 중 하나.
**Carry over**: SUBS 파라미터화의 일부.
**Encoder-only 모델**: BERT 스타일 인코더 기반.

주요 결과

**WikiText-103**: MDLM이 기존 디퓨전 모델 중 최고 성능 달성.
**AR 퍼플렉시티 대비**: 15-25% 범위 내 퍼플렉시티 달성.
**D3PM 재구현**: 기존보다 퍼플렉시티 0.1 개선.
**T=1000 vs. Continuous time**: 퍼플렉시티 0.1 증가.
**Carry over 제거**: 퍼플렉시티 1.5 증가.

의의 및 한계

MDLM은 기존 디퓨전 모델의 성능 한계를 극복하고, BERT 스타일 모델에 생성 능력을 부여함으로써 새로운 연구 기회를 열어준다. 특히, D3PM 모델의 재구현을 통해 기존 성능 격차가 잘못된 인식이었다는 점을 밝혀내는 데 기여한다. 그러나, MDLM은 여전히 AR 모델과 완전히 동일한 성능을 보이지는 못하며, 생성 속도나 복잡도 측면에서 추가 연구가 필요하다. 또한, 연구는 언어 모델링에 초점을 맞추고 있어, 다른 이산 데이터 유형에 대한 일반화 가능성은 명시되지 않음.

실용적 활용

MDLM은 효율적인 샘플러를 사용하여 임의 길이의 텍스트를 반자율적으로 생성할 수 있어, 대규모 언어 생성 시스템, 생물학적 시퀀스 모델링, 또는 텍스트 생성을 요구하는 산업 분야에서 활용 가능하다. 특히, DNA 시퀀스 모델링에서 BERT 스타일 모델과 유사하거나 더 높은 성능을 보이며, 생성 능력을 갖춘 새로운 모델로 주목받을 수 있다.