Einfache und effektive maskierte Diffusions-Sprachmodelle

Simple and Effective Masked Diffusion Language Models

June 11, 2024

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

cs.AI

Zusammenfassung

Obwohl Diffusionsmodelle sich hervorragend zur Erzeugung hochwertiger Bilder eignen, berichten frühere Arbeiten über eine signifikante Leistungslücke zwischen Diffusion und autoregressiven (AR) Methoden im Sprachmodellieren. In dieser Arbeit zeigen wir, dass einfache maskierte diskrete Diffusion leistungsstärker ist als bisher angenommen. Wir wenden ein effektives Schulungsrezept an, das die Leistung von maskierten Diffusionsmodellen verbessert, und leiten ein vereinfachtes, Rao-Blackwellisiertes Ziel ab, das zu zusätzlichen Verbesserungen führt. Unser Ziel hat eine einfache Form - es handelt sich um eine Mischung aus klassischen maskierten Sprachmodellierungsverlusten - und kann verwendet werden, um ausschließlich Encoder-Sprachmodelle zu trainieren, die effiziente Sampler zulassen, einschließlich solcher, die halb-autoregressiv beliebige Textlängen generieren können, ähnlich einem traditionellen Sprachmodell. Bei Sprachmodellierungs-Benchmarks erreicht eine Reihe von maskierten Diffusionsmodellen, die mit modernen Ingenieurpraktiken trainiert wurden, einen neuen Stand der Technik unter den Diffusionsmodellen und nähert sich der AR-Perplexität an. Wir veröffentlichen unseren Code unter: 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