Search papers, labs, and topics across Lattice.
This paper introduces a novel, distribution-agnostic unmasking schedule for diffusion language models (DLMs) that adapts to the intrinsic dependence structure of the target data without requiring prior knowledge or hyperparameter tuning. The method randomizes the number of tokens revealed at each iteration, unlike prior deterministic approaches with fixed unmasking sizes. Theoretical analysis demonstrates that the sampling convergence, measured by KL divergence, scales as $\widetilde O(\mathsf{TC}/K)$ and $\widetilde O(\mathsf{DTC}/K)$ for two parameter choices, where $\mathsf{TC}$ and $\mathsf{DTC}$ represent total correlation and dual total correlation, respectively, showing improved convergence guarantees and sampling acceleration, especially for low-complexity distributions.
Diffusion language models can now adapt to the underlying data structure without hyperparameter tuning, thanks to a new randomized unmasking schedule that provably accelerates sampling.
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.
