[论文解读] Smooth Normalizing Flows

Jonas Köhler, Andreas Krämer+1 more|arXiv (Cornell University)|Oct 1, 2021

Protein Structure and Dynamics被引用 15

一句话总结

本论文提出了一种在紧致区间和超环面（hypertori）上建模概率分布的C∞-光滑混合变换正则化流（normalizing flows），特别适用于具有内坐标的分子体系。通过利用反函数定理实现可微分的反演，该方法支持双向训练、力匹配（force matching）以及在分子动力学模拟中的应用，实现了具有良好行为特性的力和高采样效率的精确、平滑势能。

ABSTRACT

Normalizing flows are a promising tool for modeling probability distributions in physical systems. While state-of-the-art flows accurately approximate distributions and energies, applications in physics additionally require smooth energies to compute forces and higher-order derivatives. Furthermore, such densities are often defined on non-trivial topologies. A recent example are Boltzmann Generators for generating 3D-structures of peptides and small proteins. These generative models leverage the space of internal coordinates (dihedrals, angles, and bonds), which is a product of hypertori and compact intervals. In this work, we introduce a class of smooth mixture transformations working on both compact intervals and hypertori. Mixture transformations employ root-finding methods to invert them in practice, which has so far prevented bi-directional flow training. To this end, we show that parameter gradients and forces of such inverses can be computed from forward evaluations via the inverse function theorem. We demonstrate two advantages of such smooth flows: they allow training by force matching to simulation data and can be used as potentials in molecular dynamics simulations.

研究动机与目标

• 解决在紧致区间和超环面上缺乏C∞-光滑正则化流的问题，这对于建模分子内坐标至关重要。
• 通过反函数定理的梯度计算，实现对非解析逆流的双向训练。
• 开发适用于分子动力学模拟的平滑、可微分势能，确保力的精确性。
• 克服先前方法的局限性，如产生不连续的力或需要先验支持知识。

提出的方法

• 提出一类使用可学习形状参数的光滑截断函数（bump functions）的C∞-光滑混合变换，适用于紧致区间和超环面。
• 采用黑箱根求解方法（如多区间搜索）在反向方向上实现流的反演。
• 通过反函数定理计算逆的梯度，从而实现对根求解步骤的反向传播。
• 使用混合损失函数进行训练，结合最大似然估计（MLE）、力匹配（FM）和反向KL散度。
• 在分子动力学模拟中将学习到的流作为势能，采用速度Verlet积分器。
• 使用OpenMM进行能量和力的评估，并通过采样效率和能量守恒性进行验证。

实验结果

研究问题

• RQ1 能否构建C∞-光滑正则化流以同时支持紧致区间和超环面，从而实现对分子内坐标的平滑建模？
• RQ2 是否能够通过反函数定理实现对非解析逆流的梯度反向传播？
• RQ3 结合力匹配与密度估计，能否提升分子体系中学习势能的准确性？
• RQ4 平滑流能否产生与经典力场相当的稳定、能量守恒的分子动力学模拟？
• RQ5 平滑流是否能在采样效率和力精度方面优于基于样条或仿射的流？

主要发现

• 在丙氨酸二肽（alanine dipeptide）上，平滑流模型的采样效率达到42%，优于基于样条的流（38%）和仿射流（<1%）。
• 采用力匹配与MLE联合训练的平滑流，其反向KL散度为0.0003，表明密度匹配近乎完美。
• 使用平滑流的分子动力学模拟中，能量波动与Amber ff99SB-ILDN力场相当，表明具有良好的数值稳定性。
• 由平滑流计算出的力显著优于基于样条的流，测试数据上的力残差更低。
• 该方法通过反函数定理计算梯度，成功实现了对非解析逆流的双向训练，克服了先前工作的关键局限。
• 平滑流架构生成了行为良好的力和能量，可直接用于分子动力学模拟，无力的不连续性。