面向玻尔兹曼生成器的归一化流高效回归训练方法
Efficient Regression-Based Training of Normalizing Flows for Boltzmann Generators
AI总结:
针对玻尔兹曼生成器中归一化流最大似然训练不稳定、计算成本高的问题,提出基于回归的RegFlow训练方法,采用最优输运耦合或预训练流构造目标,结合前向后向自洽损失提升稳定性,在分子构象采样任务上优于传统方法。
摘要:
无模拟训练框架一直是连续空间生成建模变革的前沿方向,催生了大规模扩散模型和流匹配模型。然而,这类现代生成模型存在推理成本高昂的问题,限制了其在诸多科学场景中的应用,例如需要快速似然计算的分子构象玻尔兹曼生成器(BGs)。本文针对玻尔兹曼生成器场景重新研究经典归一化流:这类模型采样高效、似然计算便捷,但采用最大似然方法训练时往往不稳定,且计算难度大。我们提出归一化流回归训练方法(RegFlow),这是一种新型可扩展的基于回归的训练目标,它规避了传统最大似然训练的数值不稳定性与计算难题,转而采用简单的$\ell_2$回归目标。具体而言,RegFlow将流模型下的先验样本映射至通过最优输运耦合或预训练连续归一化流(CNF)计算得到的目标。为提升数值稳定性,RegFlow采用了多种有效的正则化策略,例如一种实现简便的新型前向后向自洽损失。实验表明,RegFlow可适配更广泛的架构类别,而这些架构此前无法通过最大似然方法用于玻尔兹曼生成器的训练。我们还证实,在丙氨酸二肽、三肽和四肽的笛卡尔坐标平衡采样任务中,RegFlow在性能、计算成本和稳定性方面均优于最大似然训练,展现了其在分子系统中的应用潜力。
英文摘要:
Simulation-free training frameworks have been at the forefront of the generative modelling revolution in continuous spaces, leading to large-scale diffusion and flow matching models. However, such modern generative models suffer from expensive inference, inhibiting their use in numerous scientific applications like Boltzmann Generators (BGs) for molecular conformations that require fast likelihood evaluation. In this paper, we revisit classical normalizing flows in the context of BGs that offer efficient sampling and likelihoods, but whose training via maximum likelihood is often unstable and computationally challenging. We propose Regression Training of Normalizing Flows (RegFlow), a novel and scalable regression-based training objective that bypasses the numerical instability and computational challenge of conventional maximum likelihood training in favour of a simple $\ell_2$-regression objective. Specifically, RegFlow maps prior samples under our flow to targets computed using optimal transport couplings or a pre-trained continuous normalizing flow (CNF). To enhance numerical stability, RegFlow employs effective regularization strategies such as a new forward-backward self-consistency loss that enjoys painless implementation. Empirically, we demonstrate that RegFlow unlocks a broader class of architectures that were previously intractable to train for BGs with maximum likelihood. We also show RegFlow exceeds the performance, computational cost, and stability of maximum likelihood training in equilibrium sampling in Cartesian coordinates of alanine dipeptide, tripeptide, and tetrapeptide, showcasing its potential in molecular systems.
机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。