Normalizing flows are a class of deep generative models that are especially interesting for modeling probability distributions in physics, where the exact likelihood of flows allows reweighting to known target energy functions and computing unbiased observables. For instance, Boltzmann generators tackle the long-standing sampling problem in statistical physics by training flows to produce equilibrium samples of many-body systems such as small molecules and proteins. To build effective models for such systems, it is crucial to incorporate the symmetries of the target energy into the model, which can be achieved by equivariant continuous normalizing flows (CNFs). However, CNFs can be computationally expensive to train and generate samples from, which has hampered their scalability and practical application. In this paper, we introduce equivariant flow matching, a new training objective for equivariant CNFs that is based on the recently proposed optimal transport flow matching. Equivariant flow matching exploits the physical symmetries of the target energy for efficient, simulation-free training of equivariant CNFs. We demonstrate the effectiveness of flow matching on rotation and permutation invariant many-particle systems and a small molecule, alanine dipeptide, where for the first time we obtain a Boltzmann generator with significant sampling efficiency without relying on tailored internal coordinate featurization. Our results show that the equivariant flow matching objective yields flows with shorter integration paths, improved sampling efficiency, and higher scalability compared to existing methods.
翻译:归一化流是一类深度生成模型,在物理学中用于建模概率分布时尤为有趣,其精确似然性允许对已知目标能量函数进行重加权并计算无偏观测量。例如,玻尔兹曼生成器通过训练流来产生小分子和蛋白质等多体系统的平衡样本,从而解决统计物理学中长期存在的采样问题。要为这类系统构建有效模型,必须将目标能量的对称性融入模型中,这可通过等变连续归一化流(CNF)实现。然而,CNF的训练和采样计算成本高昂,这限制了其可扩展性和实际应用。本文提出等变流匹配——基于最新提出的最优输运流匹配方法为等变CNF设计的新型训练目标。该目标利用目标能量的物理对称性,实现无需模拟的等变CNF高效训练。我们在旋转置换不变多粒子系统及小分子丙氨酸二肽上验证了流匹配的有效性,首次在不依赖定制内坐标特征化的情况下,获得了具有显著采样效率的玻尔兹曼生成器。实验结果表明,与现有方法相比,等变流匹配目标能够生成积分路径更短、采样效率更高且可扩展性更强的流。