Improving Discrete Diffusion Models via Structured Preferential Generation

TL;DR

This work tackles discrete diffusion models for text by introducing a structured forward process that biases the generation order across token categories. It develops a mutual information schedule to allocate diffusion steps and evaluates several token-order strategies, including a common-first ordering and an information-gain-based approach. Empirical results on toy data and text8 show that common-first generation can improve log-likelihood/perplexity relative to standard absorbing diffusion, though other orders may underperform and larger-vocabulary datasets require further design refinement. Overall, the paper provides a principled framework to inject inductive biases into discrete diffusion and lays groundwork for applying structured diffusion to broader discrete domains.

Abstract

In the domains of image and audio, diffusion models have shown impressive performance. However, their application to discrete data types, such as language, has often been suboptimal compared to autoregressive generative models. This paper tackles the challenge of improving discrete diffusion models by introducing a structured forward process that leverages the inherent information hierarchy in discrete categories, such as words in text. Our approach biases the generative process to produce certain categories before others, resulting in a notable improvement in log-likelihood scores on the text8 dataset. This work paves the way for more advances in discrete diffusion models with potentially significant enhancements in performance.

Figures (7)

• Figure 1: Three approaches to generating the expression: . The words are ordered top-down by how common they are, and the generation order is either left-to-right (a), random (b), or rare-words-first (c).
• Figure 2: Test perplexities on text8 calculated with the ELBO as a proxy for the marginal likelihood.
• Figure 3: Test perplexities with different variants of the common-first model where tokens are grouped together with different strategies on Wikitext-2, as well as the standard absorbing state model.
• Figure 4: Visualization of the forward process for text8 with the diffusion where common tokens are moved to the absorbing state first. The process starts out by diffusing 'e':s and spaces.
• Figure 5: Visualization of the reverse process of the diffusion process where the common tokens are generated first. The trained model is a 12-layer transformer with about 10 million parameters.