Generalization Bounds for Discrete Diffusion: Statistical Advantage of Masking
Abstract
Discrete diffusion models have recently emerged as a compelling alternative for language generation, enabling efficient non-autoregressive sampling while achieving strong empirical performance. A key design choice in discrete diffusion---absent in most continuous diffusion formulations---is the forward corruption kernel, with masked/absorbing corruption now dominating practice. Despite this empirical preference, there is limited statistical theory explaining when and why masking should outperform alternative kernels such as uniform replacement. In this paper, we take a step toward closing this gap from a statistical learning perspective. Our analysis establishes generalization bounds and, through an explicit comparison across different forward corruption kernels, reveals a central advantage of masking: it scales with the effective data support rather than the full ambient state space, thereby mitigating the curse of state space cardinality. We further derive structure-aware refinements that capture how concentration and sparsity in real sequential data sharpen the sample complexities. Together, these results offer a principled explanation for the empirical strength of masked diffusion and provide guidance for forward-kernel design in discrete generative modeling.
Lay Summary
AI systems that generate text are increasingly built with discrete diffusion, which offers an efficient alternative to the autoregressive generation used by most current language models. These systems learn by first corrupting real text into noise and then learning to undo that corruption. A central design choice is how to corrupt the text. Today's best systems overwhelmingly favor one option: hiding words behind a "mask" token, rather than the alternative of randomly replacing them with other words. Yet there has been little theory explaining why masking should work better. We take a step toward closing this gap. Studying how much training data each approach requires, we show a central advantage of masking: it only needs to learn from text resembling real language, rather than the astronomically larger space of all possible word combinations. We further refine this finding to capture how the structure of real language sharpens it. Together, our results explain the empirical strength of masking and offer guidance for designing future text-generating systems.