T-SCEND: Test-time Scalable MCTS-enhanced Diffusion Model

我们引入了测试时可扩展的MCTS增强扩散模型（T-SCEND），这是一种新颖的框架，通过改进基于能量的训练和扩大测试时计算规模，显著提升了扩散模型的推理能力。首先，我们展示了简单地增加扩散模型的推理预算只会带来边际收益。为了解决这一问题，T-SCEND的训练包括一个新颖的线性回归负对比学习目标，以提高能量景观的性能与能量一致性，并采用KL正则化来减少对抗性采样。在推理过程中，T-SCEND将去噪过程与一种新颖的混合蒙特卡洛树搜索（hMCTS）相结合，该方法在去噪过程中依次执行最佳N随机搜索和蒙特卡洛树搜索。在复杂的迷宫和数独推理任务中，我们展示了T-SCEND训练目标和可扩展推理方法的有效性。特别是，使用最大为6x6大小的迷宫进行训练后，我们的T-SCEND能够解决88%的更大规模（15x15）的迷宫问题，而标准扩散模型则完全失败。用于重现实验的代码可以在https://github.com/AI4Science-WestlakeU/t_scend找到。

该论文试图解决扩散模型在推理任务中的性能提升问题，特别是在增加推理预算时收益递减的问题。这是一个新问题，特别是在结合蒙特卡洛树搜索（MCTS）增强扩散模型的背景下。

T-SCEND的关键思路在于引入了一种新的线性回归负对比学习目标和KL正则化来优化能量景观的一致性和减少对抗性采样。此外，在推理阶段，它结合了最佳N随机搜索和MCTS的混合方法。这种结合不仅提高了推理能力，还使得模型能够在更大规模的任务上表现良好。

论文展示了T-SCEND在Maze和Sudoku等复杂推理任务上的显著效果，特别是对于从未见过的大规模迷宫（15x15），其解题率达到了88%，而传统扩散模型完全失败。此外，作者提供了开源代码，允许他人复现实验结果。未来的研究可以进一步探索不同类型的推理任务以及更广泛的模型应用。

最近在这个领域中，相关的研究包括但不限于： 1. 'Denoising Diffusion Probabilistic Models' - 探索了扩散模型的基本原理及其在生成任务中的应用。 2. 'Monte Carlo Tree Search in Continuous Action Spaces' - 研究了MCTS在连续动作空间中的应用。 3. 'Energy-Based Models: A Survey' - 提供了关于能量基模型的全面综述。 4. 'Scalable Inference for Structured Gaussian Process Models' - 讨论了高斯过程模型中的可扩展推理方法。