Flexible-length Text Infilling for Discrete Diffusion Models

TL;DR

DDOT tackles the key limitation of discrete diffusion for text by enabling flexible-length infilling through diffusion over token positions in addition to token values. It uses sample-level Optimal Transport coupling within prompt and response sets to preserve intra-set order while allowing inter-set rearrangements, combined with a token- and a position-loss objective for end-to-end training. The method supports two initialization schemes and yields efficient, parallelizable inference with near-state-of-the-art non-autoregressive performance on standard infilling benchmarks. Empirical results on One-Billion-Word and Yelp demonstrate strong gains over diffusion baselines and competitive performance relative to non-autoregressive methods, with improved efficiency and scalability for longer prompt scenarios.

Abstract

Discrete diffusion models are a new class of text generators that offer advantages such as bidirectional context use, parallelizable generation, and flexible prompting compared to autoregressive models. However, a critical limitation of discrete diffusion models is their inability to perform flexible-length or flexible-position text infilling without access to ground-truth positional data. We introduce \textbf{DDOT} (\textbf{D}iscrete \textbf{D}iffusion with \textbf{O}ptimal \textbf{T}ransport Position Coupling), the first discrete diffusion model to overcome this challenge. DDOT jointly denoises token values and token positions, employing a novel sample-level Optimal Transport (OT) coupling. This coupling preserves relative token ordering while dynamically adjusting the positions and length of infilled segments, a capability previously missing in text diffusion. Our method is orthogonal to existing discrete text diffusion methods and is compatible with various pretrained text denoisers. Extensive experiments on text infilling benchmarks such as One-Billion-Word and Yelp demonstrate that DDOT outperforms naive diffusion baselines. Furthermore, DDOT achieves performance on par with state-of-the-art non-autoregressive models and enables significant improvements in training efficiency and flexibility.

Figures (7)

• Figure 1: Our diffusion across token positions enables dynamic token movement for infilling. Unlike prior methods, DDOT learns to move masked tokens to appropriate locations, such as to the right of "brown," even if that position was not initially masked. The OT coupling (colored lines) simplifies this learning by drastically reducing the set of possible permutations.
• Figure 2: DDOT learns to vary infilled span lengths and positions, unlike prior fixed-position diffusion methods. (Left) We compute two separate intra-set OT couplings within the prompt positions and the response positions, which drastically simplifies the set of possible permutations. (Right) Given a time step $t$, we predict the token and position.
• Figure 3: Success rate on block datasets. LC and PoP increasingly generate invalid responses (missing or swapping prompt tokens) as the number of prompt tokens grows.
• Figure 4: (a) and (b) show ground-truth token positions over time. Without OT (a), many line crossings indicate unstable permutations, whereas with OT coupling (b), trajectories are nearly straight throughout the denoising process. (c) Performance tends to increase with more sampling steps.
• Figure 5: Ablations on number of sampling steps without OT coupling. Paths curve and cross more often, complicating learning and inference.