基于噪声轨迹搜索的扩散模型测试时扩展

扩散模型的迭代性与随机性使其具备测试时扩展能力，即在去噪过程中投入额外计算量可生成更高保真度的样本。增加去噪步数是主要的扩展维度，但其收益递减极快。转而优化噪声轨迹——即注入的噪声向量序列——更具前景，因为具体的噪声实现严重影响样本质量；然而由于高维搜索空间、复杂的噪声-结果交互以及高昂的轨迹评估代价，这极具挑战性。我们首先将扩散过程建模为带有终端奖励的马尔可夫决策过程（MDP），表明蒙特卡洛树搜索（MCTS）等树搜索方法虽有意义但不可行，以此应对该挑战。为平衡性能与效率，我们随后采用 MDP 的松弛策略，将去噪视为一系列独立的上下文赌博机问题。这使得我们能够引入一种 $\epsilon$-greedy 搜索算法，在极端时间步进行全局探索，并在发生解混的中间步骤进行局部利用。在 EDM 和 Stable Diffusion 上的实验表明，该方法在类条件生成/文本到图像生成中取得了 SOTA 分数，较基线提升最高达 $164\%$，且性能匹配或超越 MCTS。据我们所知，这是首个针对任意（不可微）奖励的测试时噪声轨迹优化的实用方法。