Equivariant flow matching

by   Leon Klein, et al.

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 our approach on 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.


page 1

page 2

page 3

page 4


Equivariant Flows: sampling configurations for multi-body systems with symmetric energies

Flows are exact-likelihood generative neural networks that transform sam...

Imitative Planning using Conditional Normalizing Flow

We explore the application of normalizing flows for improving the perfor...

Energy Flows: Towards Determinant-Free Training of Normalizing Flows

Normalizing flows are a popular approach for constructing probabilistic ...

Designing losses for data-free training of normalizing flows on Boltzmann distributions

Generating a Boltzmann distribution in high dimension has recently been ...

SE(3) Equivariant Augmented Coupling Flows

Coupling normalizing flows allow for fast sampling and density evaluatio...

Gaussianization Flows

Iterative Gaussianization is a fixed-point iteration procedure that can ...

Multisample Flow Matching: Straightening Flows with Minibatch Couplings

Simulation-free methods for training continuous-time generative models c...

Please sign up or login with your details

Forgot password? Click here to reset