[论文解读] Equivariant flow matching

Leon Klein, Andreas Krämer+1 more|arXiv (Cornell University)|Jun 26, 2023

Topic Modeling被引用 6

一句话总结

本文提出等变流匹配（equivariant flow matching），一种新颖的训练目标，用于等变连续归一化流（CNFs），该方法利用最优传输（optimal transport）原理，实现高效、无需仿真的训练。通过利用旋转和置换不变性等物理对称性，该方法实现了更短的积分路径、更高的采样效率和更强的可扩展性——在LJ13和丙氨酸二肽等系统上验证，无需依赖定制的内部坐标。

ABSTRACT

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.

研究动机与目标

• 解决在对称多体系统中训练和推断等变连续归一化流（CNFs）时计算成本过高的问题。
• 开发一种无需仿真的训练目标，利用最优传输（OT）映射，同时尊重如旋转和置换不变性等物理对称性。
• 改善玻尔兹曼生成器模型在统计物理应用中的采样效率，并缩短积分路径长度。
• 实现在LJ55和丙氨酸二肽等大型对称系统上的可扩展训练，且不依赖于系统特定的特征化方法。
• 为对称分布提供一种高效替代方案，以替代基于似然的训练和标准OT流匹配方法。

提出的方法

• 提出一种等变最优传输（OT）流匹配目标，强制在先验分布与目标分布之间实现对称性感知的传输路径。
• 将OT流匹配框架适配于不变密度，通过引入尊重目标能量函数对称性的群等变神经网络。
• 使用批量重排和并行预生成最优传输对，以减少训练过程中的计算瓶颈。
• 采用基于冯·米塞斯分布的加权采样策略，以改善低概率区域（如丙氨酸二肽中φ二面角的过渡区域）的密度估计。
• 引入一种对称性感知损失函数，最小化传输成本，同时确保流在相关对称群（如SO(D)和S_N）作用下保持等变性。
• 将该方法应用于在不变分布上训练CNFs，采用带有残差块和等变消息传递的连续归一化流架构。

实验结果

研究问题

• RQ1 能否设计一种流匹配目标，使其在保持对称性的同时，实现在不变分布中的近似最优传输路径？
• RQ2 与标准OT流匹配和基于似然的训练相比，等变流匹配是否能缩短积分路径长度并提升采样效率？
• RQ3 该方法是否能在不依赖内部坐标特征化的情况下，有效扩展至大型对称系统（如LJ55和丙氨酸二肽）？
• RQ4 在传输成本和路径精度方面，等变流匹配与标准OT流匹配相比表现如何？
• RQ5 在具有缓慢过渡的系统（如丙氨酸二肽中的φ二面角）中，对训练数据进行偏置能否改善自由能估计？

主要发现

• 与标准OT流匹配相比，等变流匹配在LJ13和丙氨酸二肽上显著缩短了积分路径，弧长减少高达30%。
• 在LJ13上，训练时间缩短至0.5小时；在丙氨酸二肽上，训练时间缩短至25小时，远快于基于似然的训练（分别为13小时和17小时）。
• 在LJ55系统上，等变流匹配以256的批量大小实现17小时的训练时间，证明了其在更大系统上的可扩展性。
• 在丙氨酸二肽中，该方法在200个周期内实现了1×10⁵的有效样本大小，表现出稳定的收敛性和更高的采样效率。
• 等变OT流匹配的积分路径长度始终更接近理论最优传输路径，尤其在高对称性系统中表现更优。
• 在CPU上，等变OT流匹配的批量重排每批次仅耗时22.4秒，支持高效的并行化，具备向更大系统扩展的能力。