Modelos de Linguagem de Difusão Mascarada Simples e Eficazes

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

Resumo

Embora os modelos de difusão se destaquem na geração de imagens de alta qualidade, trabalhos anteriores relatam uma lacuna significativa de desempenho entre métodos de difusão e autoregressivos (AR) na modelagem de linguagem. Neste trabalho, mostramos que a difusão discreta mascarada simples é mais eficiente do que se pensava anteriormente. Aplicamos uma receita de treinamento eficaz que melhora o desempenho dos modelos de difusão mascarada e derivamos um objetivo simplificado e Rao-Blackwellizado que resulta em melhorias adicionais. Nosso objetivo tem uma forma simples — é uma mistura de perdas clássicas de modelagem de linguagem mascarada — e pode ser usado para treinar modelos de linguagem apenas com codificadores que admitem amostradores eficientes, incluindo aqueles que podem gerar textos de comprimento arbitrário de forma semi-autoregressiva, como um modelo de linguagem tradicional. Em benchmarks de modelagem de linguagem, uma variedade de modelos de difusão mascarada treinados com práticas modernas de engenharia alcança um novo estado da arte entre os modelos de difusão e se aproxima da perplexidade AR. Disponibilizamos nosso código em: 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