Transformer-Based Masked Diffusion Models Provably Learn Distributions with Latent Structure
Abstract
Masked diffusion models (MDMs) generate discrete data by iteratively unmasking multiple tokens in parallel using a learned mask predictor. Their output carries two errors: parallelization error due to parallel sampling, and estimation error from the mask predictor learned with finite samples. Existing theory treats these errors separately, and the available statistical bounds typically scale with the number VL of length-L sequences over a vocabulary of size V, suffering from the curse of dimensionality. In this work, we establish end-to-end statistical guarantees showing that transformer-based MDMs learn structured distributions with a sample complexity governed by their latent structure. Specifically, we model this structure through mixtures of hidden Markov trees, which capture latent dependencies within components and heterogeneity across components. When the data distribution is well approximated by such a mixture, a transformer-based MDM with parallel unmasking schedules achieves ε-accurate sampling in KL divergence with Õ(LKpar/ε) training samples, where Kpar is polynomial in L, V, the number of hidden states, and the number of mixture components. The transformer network and sampling procedure require no knowledge of the tree topologies or component parameters, thereby adapting to the unknown latent structure automatically. Overall, our results provide theoretical justification for the empirical effectiveness of MDMs on real-world data.
Related
- One-Layer Transformer Provably Learns Multiclass One-Nearest Neighbor in Context
- Diffusion Transformers are Provably Optimal In-context Generators
- FHRFormer: A Self-Supervised Masked Transformer Framework for Fetal Heart Rate Time-Series Inpainting and Forecasting
- Litter-Masked Block z-Channel: Designing Cover Distributions Against Public-Design Observers
- Structured Masked Diffusion for Shared-Codebook Unsourced Random Access Decoding
BibTeX
@misc{wu2026transformerbased,
title = {{Transformer-Based Masked Diffusion Models Provably Learn Distributions with Latent Structure}},
author = {Jingda Wu and Changxiao Cai},
howpublished = {arXiv},
year = {2026},
month = oct,
doi = {10.48550/arxiv.2610.12544},
eprint = {2610.12544},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2610.12544},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - Transformer-Based Masked Diffusion Models Provably Learn Distributions with Latent Structure
AU - Jingda Wu
AU - Changxiao Cai
PY - 2026
DA - 2026/10/08
PB - arXiv
DO - 10.48550/arxiv.2610.12544
UR - https://arxiv.org/abs/2610.12544
AB - Masked diffusion models (MDMs) generate discrete data by iteratively unmasking multiple tokens in parallel using a learned mask predictor. Their output carries two errors: parallelization error due to parallel sampling, and estimation error from the mask predictor learned with finite samples. Existing theory treats these errors separately, and the available statistical bounds typically scale with the number V^L of length-L sequences over a vocabulary of size V, suffering from the curse of dimensionality. In this work, we establish end-to-end statistical guarantees showing that transformer-based MDMs learn structured distributions with a sample complexity governed by their latent structure. Specifically, we model this structure through mixtures of hidden Markov trees, which capture latent dependencies within components and heterogeneity across components. When the data distribution is well approximated by such a mixture, a transformer-based MDM with parallel unmasking schedules achieves ε-accurate sampling in KL divergence with Õ(LK_par/ε) training samples, where K_par is polynomial in L, V, the number of hidden states, and the number of mixture components. The transformer network and sampling procedure require no knowledge of the tree topologies or component parameters, thereby adapting to the unknown latent structure automatically. Overall, our results provide theoretical justification for the empirical effectiveness of MDMs on real-world data.
ER -