Simple and Effective Masked Diffusion Language Models
Authors: name: Subham Sekhar Sahoo; name: Marianne Arriola; name: Yair Schiff; name: Aaron Gokaslan; name: Edgar Marroquin; name: Justin T Chiu; name: Alexander Rush; name: Volodymyr Kuleshov
Published: 2024-06-11 17:51:40

Abstract: 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 release our code at:
https://github.com/kuleshov-group/mdlm