Learning Generation Orders for Masked Discrete Diffusion Models via Variational Inference

掩码离散扩散模型（MDM）是一种颇具前景的新型生成建模方法，其优势在于支持并行化令牌生成，因而相比自回归模型具有更高的计算效率。然而，如何在并行生成能力与样本质量之间取得最优平衡，仍是一个尚未解决的开放性问题。目前主流方法主要依赖固定、启发式的并行采样策略来应对这一挑战。尽管近期已出现若干基于学习的解决方案，但从变分推断视角对该问题进行建模的研究依然十分匮乏。本文提出一种面向MDM的变分推断框架，用于学习最优的并行生成顺序。作为该方法的重要组成部分，我们设计了一种针对生成顺序近似后验分布的参数化形式，使其既能支持并行化操作，又能在训练过程中实现高效采样。基于该方法，我们在GSM8K数据集上开展了初步实验：在高度并行化的生成设定下，本方法的表现可与各类启发式采样策略相媲美。例如，在平均仅需4步生成即可完成推理的情况下，本方法达到了33.1%的准确率，而现有主流对比方法在相同步数下的准确率仅为23.7%–29.0%。我们相信，后续更深入的实验验证与方法分析，将为理解MDM中的并行生成问题提供宝贵洞见。

如何在掩码离散扩散模型（MDMs）中学习最优的并行生成顺序，以在高度并行（极少数步数）条件下兼顾生成效率与样本质量——这是一个尚未被充分建模的新问题，尤其缺乏从变分推断视角的系统性建模。

首次将并行生成顺序的学习形式化为变分推理问题：引入可学习的、支持并行采样的近似后验分布（参数化为顺序-感知的Gumbel-top-k策略），联合优化生成目标与顺序先验，使模型能自适应地发现高质量、低步数的填充路径。

在GSM8K数学推理数据集上验证：仅需平均4步即达33.1%准确率，显著超越同类启发式方法（23.7–29.0%）；方法天然支持训练时并行梯度更新与推理时硬性步数约束；暂未开源代码，但实验设计清晰区分了步数受限下的公平比较；值得深入的方向包括顺序后验的结构归纳偏置设计、跨任务泛化性及与推理链（CoT）解耦分析。

Learning Stepwise Generation Orders for Diffusion Models (ICLR 2024); Masked Autoencoders Are Scalable Vision Learners (MAE, CVPR 2022); Denoising Diffusion Probabilistic Models (DDPM, NeurIPS 2020); Parallel Sampling for Masked Language Modeling (ACL 2023); Variational Diffusion Models (ICML 2022)