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

近年来，掩码扩散模型（MDMs）作为一种有前景的生成建模方法，在离散域上崭露头角。与自回归模型（ARMs）相比，MDMs在训练时权衡了复杂性，而在推理时则提供了更大的灵活性。在训练过程中，它们必须学会解决数量呈指数级增长的填充问题，但在推理时，它们可以几乎以任意顺序解码标记。在这项工作中，我们详细研究了这两种相互竞争的影响。 在训练方面，我们从理论上和实证上证明了与自回归模型相比，MDMs确实在计算上不可行的子问题上进行训练。在推理方面，我们展示了通过采用适当的策略来适应性地选择标记解码顺序，可以显著增强MDMs的能力，使其能够避开困难的子问题。在诸如数独之类的逻辑谜题中，我们证明了自适应推理可以将预训练MDMs的求解准确性从不到7%提升至约90%，甚至超过了参数量多7倍且通过教师强制明确训练以学习正确解码顺序的自回归模型。

该论文试图解决在离散域中生成模型的效率与性能问题，特别是对比自回归模型（ARMs），探讨掩码扩散模型（MDMs）在训练和推理阶段的复杂性及灵活性。这是一个相对较新的问题，特别是在如何平衡训练时的计算复杂性和推理时的灵活性方面。

关键思路在于通过掩码扩散模型（MDMs）在训练时学习解决大量填充问题，而在推理时能够以任意顺序解码令牌，从而提高模型的灵活性。相比传统的自回归模型，MDMs在推理时可以动态选择最优的解码顺序，避免了难以解决的子问题。这种策略尤其在逻辑谜题如数独等任务上表现出色。

论文展示了通过自适应推理策略，预训练的MDMs在数独等逻辑谜题上的解题准确率从不到7%提升到约90%，甚至超过了参数量是其7倍且经过专门训练的自回归模型。实验设计包括理论分析和实证研究，使用了数独等逻辑谜题作为数据集。此外，作者还开源了部分代码，为后续研究提供了基础。值得进一步研究的方向包括如何将这种方法应用于更复杂的离散域任务，以及探索更多类型的自适应推理策略。

最近在这个领域中的相关研究包括：1.《Autoregressive Models for Sequences of Graphs》探讨了图序列的自回归建模；2.《Transformer-Based Masked Language Modeling for Discrete Data Generation》研究了基于Transformer的掩码语言模型在离散数据生成中的应用；3.《Inference Strategies for Masked Diffusion Models in Discrete Domains》讨论了不同推理策略对掩码扩散模型的影响。