Diffusione Maschera Semplificata e Generalizzata per Dati Discreti

Simplified and Generalized Masked Diffusion for Discrete Data

June 6, 2024

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

cs.AI

Abstract

La diffusione mascherata (o assorbente) è attivamente esplorata come alternativa ai modelli autoregressivi per la modellazione generativa di dati discreti. Tuttavia, il lavoro esistente in questo ambito è stato ostacolato da formulazioni di modelli eccessivamente complesse e da relazioni poco chiare tra diverse prospettive, portando a una parametrizzazione subottimale, obiettivi di addestramento inadeguati e aggiustamenti ad hoc per contrastare questi problemi. In questo lavoro, ci proponiamo di fornire un framework semplice e generale che sblocchi il pieno potenziale dei modelli di diffusione mascherata. Mostriamo che l'obiettivo variazionale in tempo continuo dei modelli di diffusione mascherata è un semplice integrale ponderato di perdite di entropia incrociata. Il nostro framework consente anche l'addestramento di modelli generalizzati di diffusione mascherata con schemi di mascheramento dipendenti dallo stato. Quando valutati in termini di perplessità, i nostri modelli addestrati su OpenWebText superano i precedenti modelli di diffusione linguistica su scala GPT-2 e dimostrano prestazioni superiori in 4 su 5 task di modellazione linguistica zero-shot. Inoltre, i nostri modelli superano di gran lunga i precedenti modelli di diffusione discretta nella modellazione di immagini a livello di pixel, raggiungendo 2,78 (CIFAR-10) e 3,42 (ImageNet 64×64) bit per dimensione, risultati comparabili o migliori rispetto ai modelli autoregressivi di dimensioni simili.

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.