Simplex Relaxation for Discrete Diffusion

2026-08-11

AI Summary

This work addresses the limited expressivity of reverse transitions in uniform discrete diffusion models by proposing Simplax, a novel approach that introduces auxiliary simplex variables to enable precise Dirichlet–categorical augmentation of categorical states. While preserving both the original corruption process and categorical marginal distributions, Simplax achieves the first computationally tractable continuous relaxation of uniform discrete diffusion. The method constructs a Rao–Blackwellized reverse bridge objective coupled with a stochastic reverse sampler, substantially enhancing the expressiveness of the training objective. Empirical results demonstrate improved trade-offs between perplexity and entropy in unconditional generation on OpenWebText, and state-of-the-art performance on Sudoku: models trained with only 30 clues achieve the highest accuracy across all clue densities—including the challenging 17-clue setting—and yield the most effective unconditional generations.

Abstract

Discrete diffusion models for categorical generation are defined by a corruption kernel, which determines the intermediate state space and the associated reverse prediction problem. We study uniform discrete diffusion and ask whether its training objective and reverse transitions can be enriched without changing the underlying categorical corruption process. We introduce Simplax, an exact Dirichlet--categorical augmentation that couples each corrupted categorical state with an auxiliary simplex-valued variable while preserving the original uniform diffusion process as its categorical marginal. This augmentation yields a tractable Rao--Blackwellized reverse-bridge objective and a corresponding stochastic reverse sampler, while retaining the corrupted categorical state as the denoiser input. Empirically, Simplax improves the generative perplexity--entropy tradeoff on unconditional OpenWebText generation. On Sudoku, a model trained exclusively on $30$-clue puzzles achieves the highest accuracy among the compared methods across all evaluated clue densities, including the minimum uniquely solvable $17$-clue regime, and also achieves the highest validity in unconditional generation.

Problem

• discrete diffusion
• categorical generation
• uniform corruption
• reverse transition
• generative modeling

Innovation

• discrete diffusion
• simplex relaxation
• Dirichlet-categorical augmentation
• Rao-Blackwellization
• categorical generation