Diffusion masquée simplifiée et généralisée pour les données discrètes

Simplified and Generalized Masked Diffusion for Discrete Data

June 6, 2024

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

cs.AI

Résumé

La diffusion masquée (ou absorbante) est activement explorée comme alternative aux modèles autorégressifs pour la modélisation générative de données discrètes. Cependant, les travaux existants dans ce domaine ont été entravés par des formulations de modèles inutilement complexes et des relations obscures entre différentes perspectives, conduisant à une paramétrisation sous-optimale, des objectifs d'entraînement peu clairs et des ajustements ad hoc pour contrer ces problèmes. Dans ce travail, nous visons à fournir un cadre simple et général qui libère tout le potentiel des modèles de diffusion masquée. Nous montrons que l'objectif variationnel en temps continu des modèles de diffusion masquée est une simple intégrale pondérée de pertes d'entropie croisée. Notre cadre permet également d'entraîner des modèles de diffusion masquée généralisés avec des plans de masquage dépendants de l'état. Évalués en termes de perplexité, nos modèles entraînés sur OpenWebText surpassent les précédents modèles de diffusion linguistique à l'échelle GPT-2 et démontrent des performances supérieures sur 4 des 5 tâches de modélisation linguistique en zero-shot. De plus, nos modèles surpassent largement les précédents modèles de diffusion discrète sur la modélisation d'images au niveau des pixels, atteignant 2,78 (CIFAR-10) et 3,42 (ImageNet 64×64) bits par dimension, des résultats comparables ou supérieurs à ceux des modèles autorégressifs de tailles similaires.

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.