Adjoint Sampling: Highly Scalable Diffusion Samplers via Adjoint Matching

我们提出了伴随采样（Adjoint Sampling），这是一种用于学习从非归一化密度或能量函数中采样的扩散过程的高效且可高度扩展的算法。这是首个允许比能量评估次数和模型样本数量多得多的梯度更新的策略内方法，使我们能够扩展到比以往类似方法探索过的更大规模的问题设置。我们的框架在理论上基于随机最优控制，并与伴随匹配（Adjoint Matching）具有相同的理论保证，能够在无需使用将样本推向目标分布的校正措施的情况下进行训练。我们展示了如何在建模分子时结合关键对称性以及周期性边界条件，适用于笛卡尔坐标和扭转坐标中的分子表示。通过在经典能量函数上的广泛实验，我们验证了该方法的有效性，并进一步扩展到基于神经网络的能量模型，在多个分子系统中实现了摊销构象生成。为了鼓励开发高度可扩展的采样方法的研究，我们计划开源这些具有挑战性的基准测试，成功的采样方法将直接推动计算化学领域的进展。

该论文试图解决从非标准化密度（或能量函数）中高效采样扩散过程的问题，特别是在大规模分子建模场景下的应用。这是一个长期存在的问题，但现有的方法通常受限于梯度更新和能量评估之间的效率瓶颈。

Adjoint Sampling是一种基于策略（on-policy）的方法，其核心思想是通过结合随机最优控制理论，实现比能量评估次数更多的梯度更新。这种方法突破了传统方法的限制，允许在更大规模的问题设置中进行训练，同时无需额外的修正措施来确保样本接近目标分布。

论文展示了如何将关键对称性和周期性边界条件整合到模型中，用于分子的笛卡尔和扭转坐标建模。实验涵盖了经典能量函数以及神经网络驱动的能量模型，并成功应用于多个分子系统的构象生成任务。此外，作者计划开源具有挑战性的基准测试数据集，以推动高效采样方法的研究进展。

与本研究相关的其他工作包括：1) Adjoint Matching，它提供了类似的理论基础；2) 基于能量的模型（EBM），如“Energy-based Models: Structured Learning Beyond Likelihoods”；3) 分子动力学模拟中的扩散模型，如“Diffusion probabilistic modeling of molecular conformations”。这些研究共同推动了从复杂能量函数中采样的技术发展。