arXiv：cs.LG· Jingda Wu, Changxiao Cai·· 11 小时前
基于 Transformer 的掩码扩散模型可证明学习具有潜在结构的分布
Transformer-Based Masked Diffusion Models Provably Learn Distributions with Latent Structure
AI 导读
研究证明基于 Transformer 的掩码扩散模型（MDM）能以受潜在结构约束的样本复杂度学习结构化分布，当数据分布可由隐马尔可夫树混合良好近似时，并行解掩码调度可在 KL 散度下实现 ε 精度采样，所需训练样本为 Õ(LK_par/ε)，其中 K_par 是 L、V、隐藏状态数与混合分量数的多项式。
正文
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 $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 $\varepsilon$-accurate sampling in KL divergence with $\widetilde{O}(LK_{\mathrm{par}}/\varepsilon)$ training samples, where $K_{\mathrm{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.
| Subjects: | Machine Learning (cs.LG); Information Theory (cs.IT); Statistics Theory (math.ST); Machine Learning (stat.ML) |
| Cite as: | arXiv:2610.12544 [cs.LG] |
| (or arXiv:2610.12544v1 [cs.LG] for this version) | |
| https://doi.org/10.48550/arXiv.2610.12544 arXiv-issued DOI via DataCite (pending registration) |
Submission history
From: Jingda Wu [view email]
[v1]
Thu, 8 Oct 2026 17:54:49 UTC (90 KB)
来源：arXiv：cs.LG · arxiv.org