Simplified and Generalized Masked Diffusion for Discrete Data

TL;DR

This paper introduces a simplified, unified framework for masked diffusion models that improves discrete data modeling, achieving state-of-the-art results in language and image generation tasks by leveraging a continuous-time variational objective.

Contribution

The work presents a simple, general framework for masked diffusion models, clarifies their theoretical foundation, and demonstrates superior performance on language and image benchmarks.

Findings

Outperforms prior diffusion language models on perplexity and zero-shot tasks.

Achieves state-of-the-art bits per dimension on CIFAR-10 and ImageNet 64x64.

Provides a unified, theoretically grounded approach to masked diffusion modeling.

Abstract

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.75 (CIFAR-10) and 3.40 (ImageNet 64x64) bits per dimension that are better than autoregressive models of similar sizes. Our code is available at https://github.com/google-deepmind/md4.