Flexible-length Text Infilling for Discrete Diffusion Models

TL;DR

This paper introduces DDOT, a novel discrete diffusion model that enables flexible-length text infilling by jointly denoising tokens and their positions using optimal transport, improving flexibility and efficiency over previous methods.

Contribution

DDOT is the first discrete diffusion model to perform flexible-length text infilling by coupling token values and positions with optimal transport, enhancing adaptability and performance.

Findings

Outperforms naive diffusion baselines on text infilling benchmarks.

Achieves performance comparable to state-of-the-art non-autoregressive models.

Enables significant improvements in training efficiency and flexibility.

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.