Masked diffusion models no longer pay for sequence length: latent structure sets the sample count

The sample count for a masked diffusion model can be read off the data's latent structure instead of the size of the whole sequence space. Earlier statistical bounds scaled either with the target support size or with the full sequence-space size, and could stay exponential in sequence length even when the data admit a compact latent representation. When the data distribution is well approximated by a mixture of hidden Markov trees and the parallelization error is sufficiently small, a transformer-based masked diffusion model reaches a target KL accuracy with a number of training samples whose structural complexity is polynomial in sequence length, vocabulary size, number of hidden states, and number of mixture components.

The guarantee puts training and sampling into one error decomposition, counting both the error from parallel unmasking and the error from finite-sample training. Existing theory usually takes the mask predictor as given or with prescribed error and analyzes sampling alone, or analyzes score estimation alone, without characterizing the statistical cost of learning the predictor. The end-to-end sampling error has three parts: a first-order parallelization error that depends only on the data distribution and the unmasking schedule, plus joint-KL terms obtained by multiplying the mask-predictor estimation error by the sequence length; the parallelization error can be characterized by dual total correlation, which for star-shaped components can be far smaller than the sequence length.

The mask predictor is realized by a sparse transformer that returns conditional marginals for all masked positions in a single forward pass. Learning the joint conditionals of masked positions directly is computationally prohibitive, so practical models learn conditional marginals and sample from their product. The construction first approximates the conditional logits of a single hidden Markov tree inside transformer blocks using parallel tree contraction and reconstruction, then combines components with posterior mixture weights computed from context; the architectural budgets depend only on parameters such as sequence length and vocabulary size, not on the specific tree topology or probability tables.

The guarantee applies when the data distribution is well approximated by a mixture of hidden Markov trees and the parallelization error is sufficiently small, and it analyzes transformer-based masked diffusion models under fixed-cardinality unmasking schedules. The architectural budgets depend only on parameters such as sequence length, vocabulary size, number of hidden states, and number of components, not on the specific tree topologies, component parameters, or hidden states, so the model adapts to unknown latent structure. For trees of bounded height, a separate construction based on upward and downward message passing has depth proportional to tree height.

The analysis assumes a predictor obtained by empirical risk minimization, while practical training performs finitely many parameter updates, so optimization error is not yet covered. The parallelization-error term depends on the data distribution and the unmasking schedule, and its computability depends on estimating quantities such as dual total correlation. Extension to more general structures, including latent graphical models with cycles and confidence-based unmasking, remains open.