Modelos de Lenguaje de Difusión Enmascarada Simples y Efectivos

Simple and Effective Masked Diffusion Language Models

June 11, 2024

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

cs.AI

Resumen

Si bien los modelos de difusión sobresalen en la generación de imágenes de alta calidad, trabajos previos reportan una brecha significativa en el rendimiento entre los métodos de difusión y los modelos autoregresivos (AR) en el modelado del lenguaje. En este trabajo, demostramos que la difusión discreta enmascarada es más eficiente de lo que se pensaba anteriormente. Aplicamos una receta de entrenamiento efectiva que mejora el rendimiento de los modelos de difusión enmascarada y derivamos un objetivo simplificado y Rao-Blackwellizado que resulta en mejoras adicionales. Nuestro objetivo tiene una forma simple — es una mezcla de pérdidas clásicas de modelado del lenguaje enmascarado — y puede usarse para entrenar modelos de lenguaje que solo incluyen codificadores y admiten muestreadores eficientes, incluyendo aquellos que pueden generar textos de longitud arbitraria de manera semi-autoregresiva, como un modelo de lenguaje tradicional. En benchmarks de modelado del lenguaje, una variedad de modelos de difusión enmascarada entrenados con prácticas de ingeniería modernas alcanza un nuevo estado del arte entre los modelos de difusión y se acerca a la perplejidad de los modelos AR. Publicamos nuestro código en: 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