Equivariant flow matching
This addresses the sampling problem in statistical physics for researchers modeling many-body systems like molecules and proteins, offering a more efficient and scalable method, though it is incremental as it builds on existing flow matching techniques.
The paper tackles the computational expense and scalability issues of equivariant continuous normalizing flows (CNFs) for sampling in physics by introducing equivariant flow matching, a new training objective based on optimal transport flow matching. The result is improved sampling efficiency, shorter integration paths, and higher scalability, demonstrated on systems like alanine dipeptide without tailored featurization.
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.