Beyond Masked and Unmasked: Discrete Diffusion Models via Partial Masking

TL;DR

This paper introduces Partial Masking, a novel approach for discrete diffusion models that allows tokens to have intermediate states, improving efficiency and performance in generative tasks for text and images.

Contribution

The paper proposes Partial Masking (Prime), enabling intermediate token states in diffusion models, which enhances efficiency and generative quality without autoregressive methods.

Findings

Achieves a perplexity of 15.36 on OpenWebText, outperforming previous models.

Attains FID scores of 3.26 on CIFAR-10 and 6.98 on ImageNet-32, competitive with leading models.

Demonstrates superior performance across diverse generative tasks.

Abstract

Masked diffusion models (MDM) are powerful generative models for discrete data that generate samples by progressively unmasking tokens in a sequence. Each token can take one of two states: masked or unmasked. We observe that token sequences often remain unchanged between consecutive sampling steps; consequently, the model repeatedly processes identical inputs, leading to redundant computation. To address this inefficiency, we propose the Partial masking scheme (Prime), which augments MDM by allowing tokens to take intermediate states interpolated between the masked and unmasked states. This design enables the model to make predictions based on partially observed token information, and facilitates a fine-grained denoising process. We derive a variational training objective and introduce a simple architectural design to accommodate intermediate-state inputs. Our method demonstrates superior performance across a diverse set of generative modeling tasks. On text data, it achieves a perplexity of 15.36 on OpenWebText, outperforming previous MDM (21.52), autoregressive models (17.54), and their hybrid variants (17.58), without relying on an autoregressive formulation. On image data, it attains competitive FID scores of 3.26 on CIFAR-10 and 6.98 on ImageNet-32, comparable to leading continuous generative models.