Vereinfachte und verallgemeinerte maskierte Diffusion für diskrete Daten

Simplified and Generalized Masked Diffusion for Discrete Data

June 6, 2024

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

cs.AI

Zusammenfassung

Maskierte (oder absorbierende) Diffusion wird aktiv als Alternative zu autoregressiven Modellen für die generative Modellierung diskreter Daten erforscht. Bisher wurde die Arbeit in diesem Bereich jedoch durch unnötig komplexe Modellformulierungen und unklare Beziehungen zwischen verschiedenen Perspektiven behindert, was zu suboptimaler Parametrisierung, Trainingszielen und ad hoc-Anpassungen zur Gegenwirkung dieser Probleme führt. In dieser Arbeit zielen wir darauf ab, ein einfaches und allgemeines Rahmenwerk bereitzustellen, das das volle Potenzial von maskierten Diffusionsmodellen freisetzt. Wir zeigen, dass das kontinuierliche Variationsziel maskierter Diffusionsmodelle eine einfache gewichtete Integration von Cross-Entropy-Verlusten ist. Unser Rahmenwerk ermöglicht auch das Training von generalisierten maskierten Diffusionsmodellen mit zustandsabhängigen Maskierungsplänen. Bei der Bewertung nach Perplexität übertreffen unsere Modelle, die auf OpenWebText trainiert wurden, frühere Diffusions-Sprachmodelle im GPT-2-Maßstab und zeigen eine überlegene Leistung bei 4 von 5 Zero-Shot-Sprachmodellierungsaufgaben. Darüber hinaus übertreffen unsere Modelle deutlich frühere diskrete Diffusionsmodelle bei der Modellierung von Bilddaten auf Pixel-Ebene und erreichen 2,78 (CIFAR-10) bzw. 3,42 (ImageNet 64x64) Bits pro Dimension, die vergleichbar oder besser sind als autoregressive Modelle ähnlicher Größe.

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.