arXiv:2506.16471cs.LGcs.AI2025-06NeurIPS被引 27

提出PITA框架,让扩散模型高效采样分子系统能量分布。

Progressive Inference-Time Annealing of Diffusion Models for Sampling from Boltzmann Densities

  • 分阶段训练扩散模型,从高温到低温逐步优化采样能力。
  • 首次实现原子级粒子系统与三肽在笛卡尔坐标下的平衡采样。
  • 适合需要高效采样复杂分子构型的计算化学研究者。

从目标非归一化概率密度中高效采样仍是众多高影响力科学应用的核心挑战。基于扩散模型的生成式采样器虽有潜力,但现有方法仍无法对简单分子系统进行有效采样。本文提出渐进推理时退火(PITA)框架,结合两种互补插值技术:一是玻尔兹曼分布退火,二是扩散平滑。PITA通过依次在逐步升高的温度下训练一系列扩散模型,利用设计好的路径获得温度退火后目标密度的样本。随后,通过一种新颖的费曼-卡茨偏微分方程与序贯蒙特卡洛相结合的方法,在推理时进行退火,生成低温度下的训练样本以训练下一阶段的模型。实验表明,PITA首次实现了多体粒子系统、丙氨酸二肽及三肽在笛卡尔坐标下的平衡采样,且能量函数评估次数显著减少。代码已公开于 https://github.com/taraak/pita。

原文摘要 · Abstract (English)

Sampling efficiently from a target unnormalized probability density remains a core challenge, with relevance across countless high-impact scientific applications. A promising approach towards this challenge is the design of amortized samplers that borrow key ideas, such as probability path design, from state-of-the-art generative diffusion models. However, all existing diffusion-based samplers remain unable to draw samples from distributions at the scale of even simple molecular systems. In this paper, we propose Progressive Inference-Time Annealing (PITA), a novel framework to learn diffusion-based samplers that combines two complementary interpolation techniques: I.) Annealing of the Boltzmann distribution and II.) Diffusion smoothing. PITA trains a sequence of diffusion models from high to low temperatures by sequentially training each model at progressively higher temperatures, leveraging engineered easy access to samples of the temperature-annealed target density. In the subsequent step, PITA enables simulating the trained diffusion model to procure training samples at a lower temperature for the next diffusion model through inference-time annealing using a novel Feynman-Kac PDE combined with Sequential Monte Carlo. Empirically, PITA enables, for the first time, equilibrium sampling of N-body particle systems, Alanine Dipeptide, and tripeptides in Cartesian coordinates with dramatically lower energy function evaluations. Code available at: https://github.com/taraak/pita

扩散模型分子采样采样效率物理模拟

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