Modèles de langage à diffusion masquée simples et efficaces

Simple and Effective Masked Diffusion Language Models

June 11, 2024

Auteurs: Subham Sekhar Sahoo, Marianne Arriola, Yair Schiff, Aaron Gokaslan, Edgar Marroquin, Justin T Chiu, Alexander Rush, Volodymyr Kuleshov

cs.AI

Résumé

Bien que les modèles de diffusion excellent dans la génération d'images de haute qualité, les travaux antérieurs rapportent un écart de performance significatif entre les méthodes de diffusion et les méthodes autorégressives (AR) en modélisation du langage. Dans ce travail, nous montrons que la diffusion discrète masquée simple est plus performante qu'on ne le pensait auparavant. Nous appliquons une recette d'entraînement efficace qui améliore les performances des modèles de diffusion masquée et dérivons un objectif simplifié et rao-blackwellisé qui entraîne des améliorations supplémentaires. Notre objectif a une forme simple — il s'agit d'un mélange de pertes classiques de modélisation du langage masqué — et peut être utilisé pour entraîner des modèles de langage à encodeur uniquement qui admettent des échantillonneurs efficaces, y compris ceux capables de générer des textes de longueur arbitraire de manière semi-autorégressive, comme un modèle de langage traditionnel. Sur les benchmarks de modélisation du langage, une gamme de modèles de diffusion masquée entraînés avec des pratiques d'ingénierie modernes atteint un nouvel état de l'art parmi les modèles de diffusion et s'approche de la perplexité des modèles AR. Nous publions notre code à l'adresse suivante : https://github.com/kuleshov-group/mdlm

English

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