论文

掩码扩散模型在条件混合下的端到端理论保证

End-to-End Guarantees for Masked Diffusion Models under Conditional Mixing

精选理由

这篇理论论文回答了 MDM 为什么并行生成还能靠谱:只要信源满足条件混合,样本量不随序列长度涨,膨胀调度还能把步数从 Ω(N) 压到对数级。做扩散语言模型的人值得看。

arXiv 论文针对掩码扩散模型(MDM)采样界与训练界割裂的问题,给出学习去噪器的端到端保证。核心假设是条件混合(CM),涵盖阶数 r 的马尔可夫源和双侧 Doeblin 隐马尔可夫源,掩码 token 后验随距离按相关长度 τ_c 指数衰减。论文证明经验风险最小化器只需 ε'^(O(τ_c·log|V|)) 条训练序列即可在任意掩码分布上达到每 token 超额风险 ε',与序列长度 N 无关。端到端 KL 被界定为 ε+N·ε',在膨胀调度下去噪步数为 O(τ_c·log(N|V|/ε)),而固定随机顺序需要 Ω(N) 步。论文还证明指数阶 Θ(τ_c·log|V|) 在极小极大意义下最优,并在公开与重训模型上做了实验验证。

原文 · arXiv cs.LG

End-to-End Guarantees for Masked Diffusion Models under Conditional Mixing

Masked diffusion models train a denoiser on masked cross-entropy and generate by unmasking many positions in parallel. Existing MDM sampling bounds assume an exact denoiser or leave its error uncontrolled while training bounds cover only the masking distribution on which the model was trained. We present an end-to-end guarantee for a learned denoiser under conditional mixing (CM): stationary sources (e.g. floored order-$r$ Markov and two-sided-Doeblin hidden-Markov sources) for which revealing distant visible tokens moves a masked token's posterior on average exponentially little in the distance with correlation length $τ_c$. We show the empirical risk minimizer trained under any masking distribution attains per-token excess $\varepsilon'$ on that distribution from $\varepsilon'^{-\mathcal{O}(τ_c\log|\mathcal{V}|)}$ training sequences, independent of sequence length $N$ -- one sample budget serves every decoding schedule when each schedule's denoiser is trained under its own weighting. We further show that the exponent's order $Θ(τ_c\log|\mathcal{V}|)$ is minimax-optimal up to constants; more efficient rates require added structure. With this, we bound the end-to-end KL as $\mathrm{KL}(π\Vert p_{\mathrm{alg}})\leq \varepsilon+N\varepsilon'$ under a dilated schedule of $T=\mathcal{O}(τ_c\log(N|\mathcal{V}|/\varepsilon))$ denoising steps versus $Ω(N)$ steps for a fixed random order on an explicit CM source. Finally, we show that block-autoregressive samplers chain across blocks with no added error and that a $Λ$-looped order-$r$ message-passing class on order-$r$ sources achieves polynomial sample complexity, with approximation error improving per loop until saturation at $Λr\approx τ_c\log(1/\varepsilon')$. Experiments on released and retrained models support our results' qualitative predictions.