Difusión Enmascarada Simplificada y Generalizada para Datos Discretos

Simplified and Generalized Masked Diffusion for Discrete Data

June 6, 2024

Autores: Jiaxin Shi, Kehang Han, Zhe Wang, Arnaud Doucet, Michalis K. Titsias

cs.AI

Resumen

La difusión enmascarada (o absorbente) se está explorando activamente como una alternativa a los modelos autoregresivos para el modelado generativo de datos discretos. Sin embargo, los trabajos existentes en esta área se han visto obstaculizados por formulaciones de modelos innecesariamente complejas y relaciones poco claras entre diferentes perspectivas, lo que ha llevado a una parametrización subóptima, objetivos de entrenamiento inadecuados y ajustes ad hoc para contrarrestar estos problemas. En este trabajo, nuestro objetivo es proporcionar un marco simple y general que desbloquee todo el potencial de los modelos de difusión enmascarada. Demostramos que el objetivo variacional en tiempo continuo de los modelos de difusión enmascarada es una integral ponderada simple de pérdidas de entropía cruzada. Nuestro marco también permite entrenar modelos generalizados de difusión enmascarada con programaciones de enmascaramiento dependientes del estado. Cuando se evalúan mediante perplejidad, nuestros modelos entrenados en OpenWebText superan a los modelos de difusión de lenguaje previos a escala GPT-2 y demuestran un rendimiento superior en 4 de cada 5 tareas de modelado de lenguaje zero-shot. Además, nuestros modelos superan ampliamente a los modelos de difusión discreta anteriores en el modelado de imágenes a nivel de píxeles, logrando 2.78 (CIFAR-10) y 3.42 (ImageNet 64x64) bits por dimensión, que son comparables o mejores que los modelos autoregresivos de tamaños similares.

English

Masked (or absorbing) diffusion is actively explored as an alternative to autoregressive models for generative modeling of discrete data. However, existing work in this area has been hindered by unnecessarily complex model formulations and unclear relationships between different perspectives, leading to suboptimal parameterization, training objectives, and ad hoc adjustments to counteract these issues. In this work, we aim to provide a simple and general framework that unlocks the full potential of masked diffusion models. We show that the continuous-time variational objective of masked diffusion models is a simple weighted integral of cross-entropy losses. Our framework also enables training generalized masked diffusion models with state-dependent masking schedules. When evaluated by perplexity, our models trained on OpenWebText surpass prior diffusion language models at GPT-2 scale and demonstrate superior performance on 4 out of 5 zero-shot language modeling tasks. Furthermore, our models vastly outperform previous discrete diffusion models on pixel-level image modeling, achieving 2.78~(CIFAR-10) and 3.42 (ImageNet 64times64) bits per dimension that are comparable or better than autoregressive models of similar sizes.