arXiv:2502.06768cs.LG2025-02ICML被引 212

改进掩码扩散模型的解码顺序,显著提升其在逻辑谜题上的表现。

Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffusions

  • 通过自适应选择解码顺序,优化推理过程
  • 在数独任务中准确率从不足7%提升至约90%
  • 适用于希望高效生成高质量离散序列的研究者

近年来,掩码扩散模型(MDMs)作为离散域生成建模的一种有前景的替代方法出现。与自回归模型(ARMs)相比,MDMs在训练时付出更高复杂度代价,但在推理时具有更强的灵活性。训练阶段,它们需解决指数级数量的填空问题;推理阶段,可自由选择解码顺序。本文从理论和实证上证明,相较于自回归模型,MDMs确实面临计算上不可行的子问题。同时,我们发现一种合适的自适应解码顺序策略能显著增强模型能力,使其避开困难子问题。在数独等逻辑谜题上,该策略使预训练的MDMs解题准确率从不足7%提升至约90%,甚至超过参数量大7倍且经教师强迫训练的自回归模型。

原文摘要 · Abstract (English)

In recent years, masked diffusion models (MDMs) have emerged as a promising alternative approach for generative modeling over discrete domains. Compared to autoregressive models (ARMs), MDMs trade off complexity at training time with flexibility at inference time. At training time, they must learn to solve an exponentially large number of infilling problems, but at inference time, they can decode tokens in essentially arbitrary order. In this work, we closely examine these two competing effects. On the training front, we theoretically and empirically demonstrate that MDMs indeed train on computationally intractable subproblems compared to their autoregressive counterparts. On the inference front, we show that a suitable strategy for adaptively choosing the token decoding order significantly enhances the capabilities of MDMs, allowing them to sidestep hard subproblems. On logic puzzles like Sudoku, we show that adaptive inference can boost solving accuracy in pretrained MDMs from $<7$% to $\approx 90$%, even outperforming ARMs with $7\times$ as many parameters and that were explicitly trained via teacher forcing to learn the right order of decoding.

掩码扩散生成建模解码顺序数独求解

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。