Adaptation to Intrinsic Dependence in Diffusion Language Models
Yunxiao Zhao, Changxiao Cai
TL;DR
This work analyzes diffusion language models (DLMs) that generate text tokens in parallel by unmasking subsets of positions during sampling. It introduces a distribution-agnostic, randomized unmasking framework with two concrete schemes that adapt to the target data’s intrinsic dependence measured by total correlation (TC) and dual total correlation (DTC). Theoretical guarantees show KL convergence rates scaling with $ ext{TC}(X)$ or $ ext{DTC}(X)$, modulo logarithmic terms, in the practically important regime $K<L$, and without distribution-specific knowledge or hyperparameter tuning. Empirically, the results demonstrate that randomized unmasking sizes can yield substantial speedups for low-complexity distributions, offering principled guidance for designing DLM inference schedules with structure-adaptive benefits.
Abstract
Diffusion language models (DLMs) have recently emerged as a promising alternative to autoregressive (AR) approaches, enabling parallel token generation beyond a rigid left-to-right order. Despite growing empirical success, the theoretical understanding of how unmasking schedules -- which specify the order and size of unmasked tokens during sampling -- affect generation quality remains limited. In this work, we introduce a distribution-agnostic unmasking schedule for DLMs that adapts to the (unknown) dependence structure of the target data distribution, without requiring any prior knowledge or hyperparameter tuning. In contrast to prior deterministic procedures that fix unmasking sizes, our method randomizes the number of tokens revealed at each iteration. We show that, for two specific parameter choices, the sampling convergence guarantees -- measured by Kullback-Leibler (KL) divergence -- scale as $\widetilde O(\mathsf{TC}/K)$ and $\widetilde O(\mathsf{DTC}/K)$ respectively. Here, $K$ is the number of iterations, and $\mathsf{TC}$ and $\mathsf{DTC}$ are the total correlation and dual total correlation of the target distribution, capturing the intrinsic dependence structure underlying the data. Importantly, our guarantees hold in the practically relevant parallel-sampling regime $K<L$ where $L$ is the token sequence length. These results significantly improve upon prior convergence theories and yield substantial sampling acceleration for low-complexity distributions. Overall, our findings unveil the adaptivity of DLMs to intrinsic data structures and shed light on the benefit of randomized unmasking sizes in inference schedule design.