Diffusion Language Models are Provably Optimal Parallel Samplers

TL;DR

This work provides a theoretical foundation for diffusion language models (DLMs) as optimal parallel samplers by modeling their computation with circuit depth and width as time and space resources. It proves that DLMs endowed with chain-of-thought (CoT) can realize any sampling procedure realized by a circuit of depth $d$ in exactly $d$ decoding rounds, in contrast to autoregressive models whose sequential cost scales with circuit size. With remasking or revision, DLMs further achieve minimal memory, and the authors establish a strict expressivity gap: DLMs with these features can sample certain distributions in constant steps that are impossible for DLMs without them. Additionally, the paper presents upper and lower bounds showing that parity-based distributions can be efficiently sampled with revision/remasking, while such efficiency is impossible without these mechanisms. Altogether, the results position DLMs as the most efficient parallel samplers and motivate incorporating revision/remasking into forward processes and training data.

Abstract

Diffusion language models (DLMs) have emerged as a promising alternative to autoregressive models for faster inference via parallel token generation. We provide a rigorous foundation for this advantage by formalizing a model of parallel sampling and showing that DLMs augmented with polynomial-length chain-of-thought (CoT) can simulate any parallel sampling algorithm using an optimal number of sequential steps. Consequently, whenever a target distribution can be generated using a small number of sequential steps, a DLM can be used to generate the distribution using the same number of optimal sequential steps. However, without the ability to modify previously revealed tokens, DLMs with CoT can still incur large intermediate footprints. We prove that enabling remasking (converting unmasked tokens to masks) or revision (converting unmasked tokens to other unmasked tokens) together with CoT further allows DLMs to simulate any parallel sampling algorithm with optimal space complexity. We further justify the advantage of revision by establishing a strict expressivity gap: DLMs with revision or remasking are strictly more expressive than those without. Our results not only provide a theoretical justification for the promise of DLMs as the most efficient parallel sampler, but also advocate for enabling revision in DLMs.

Theorems & Definitions (16)

• Theorem 3.1
• Definition 1
• proof : Proof of \ref{['thm:cot']}.
• Theorem 3.2
• proof : Proof of \ref{['thm:cotr']}.
• Theorem 3.3
• Definition 2
• Definition 3
• Theorem 4.1
• proof