Equivariant flow matching

Leon Klein; Andreas Krämer; Frank Noé

Source-linked AI summary

Equivariant flow matching

Leon Klein, Andreas Krämer, Frank Noé

arXiv:2306.15030v2Jun 26, 2023stat.MLcs.LGphysics.chem-phphysics.comp-ph

TL;DR

Equivariant CNFs offer symmetry-aware Boltzmann generation but are costly to train and sample from, limiting scalability. The paper introduces equivariant flow matching, which uses symmetry-aware optimal-transport training, and reports shorter paths, improved sampling efficiency, and broader scalability, including a Cartesian-coordinate alanine-dipeptide Boltzmann generator.

Problem
Equivariant CNFs preserve physical symmetries for Boltzmann generation but face high computational costs that limit scalability and practical use.
Method
The paper introduces equivariant flow matching, aligning training pairs with particle permutations and rotations before applying optimal-transport flow matching.
Results
The method yields nearly optimal integration paths, supports larger symmetric systems, and enables a Cartesian-coordinate alanine-dipeptide Boltzmann generator with accurate free-energy estimates.
Takeaways & Limitations
Equivariant flow matching extends symmetry-aware CNFs toward larger physical systems and molecular simulation without tailored internal-coordinate featurization.
Takeaways & Limitations
Transferability was not experimentally demonstrated, CNF inference still requires complete vector-field integration, and the Hungarian-algorithm approximation costs O(N^3).

Abstract

from arXiv · show

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.

1 Introduction

Normalizing flows support exact-likelihood generation for physical distributions, but equivariant CNFs remain costly to train and sample from. The paper introduces equivariant flow matching to improve scalability while preserving physical symmetries.

Figure 1 examines LJ13 sample pairs, transport costs, per-particle integration paths, and path-length distributions across flow-matching objectives.
Equivariant CNFs incorporate physical symmetries in Cartesian-coordinate Boltzmann generators but remain computationally expensive during training and inference.
The paper proposes an invariant-density flow-matching objective designed to produce nearly optimal integration paths.
The method compares flow-matching objectives on symmetric many-particle systems and introduces invariant datasets including alanine dipeptide and a large Lennard-Jones cluster.
The resulting Boltzmann generator produces Cartesian-coordinate alanine-dipeptide samples and accurately estimates free-energy differences against umbrella sampling.

2 Related work

Prior molecular flows achieved significant sampling efficiency mainly with system-specific featurization or priors close to the target, while CNF likelihood training remains computationally expensive. Flow matching addresses this training cost by enabling simulation-free CNF training without vector-field integration or Jacobian evaluation.

Molecular flows previously required internal-coordinate featurization or target-like priors to achieve significant sampling efficiency.
Likelihood-based CNF training can be computationally expensive because it requires repeated vector-field evaluations and Jacobian-trace computation.

3 Method

The method combines continuous normalizing flows with equivariance and optimal-transport flow matching. It uses simulation-free vector-field regression and batchwise optimal transport to construct training pairs while incorporating particle, spatial, and translational symmetries.

CNFs model probability paths from a simple prior to a target distribution by integrating a time-dependent vector field.
Flow matching trains the vector field without integrating trajectories or evaluating Jacobians, reducing training cost and enabling larger models and systems.
OT flow matching uses linear conditional paths whose endpoint pairs are obtained from a 2-Wasserstein optimal transport map between prior and target samples.
In each batch, the method computes pairwise transport costs, solves a discrete OT problem, and samples training pairs from the resulting solution.
Equivariance targets permutation, rotation, reflection, and translation symmetries observed in many-particle systems, with translation invariance handled through mean-free systems.
Standard OT flow matching can produce highly curved vector fields that fail to match linear OT paths at finite batch sizes, motivating a modified objective.

4 Equivariant optimal transport flow matching

The method adapts optimal transport flow matching to invariant densities by aligning sample pairs along rotational and permutation symmetry orbits. A sequential approximation makes this symmetry-aware cost computationally tractable while retaining nearly optimal integration paths.

4 Equivariant optimal transport flow matching: Permutation matching scales as N!, whereas available sample pairs scale only with the squared batch size, making ordinary batch transport inadequate for highly symmetric systems.This motivates incorporating symmetry directly into the transport cost rather than relying only on batch reordering.
4 Equivariant optimal transport flow matching: The proposed cost aligns sample pairs consistently with the equivariant optimal transport map, enabling nearly optimal paths for invariant densities.The theoretical result establishes that the optimal transport map remains valid under the symmetry-aware cost.
4 Equivariant optimal transport flow matching: The symmetry-aware cost minimizes squared Euclidean distance over rotations, reflections, and permutations for O(D)- and S(N)-invariant distributions.These symmetries are common in molecular and many-particle systems.
4 Equivariant optimal transport flow matching: Because exact minimization over all rotations and permutations is infeasible, the method sequentially applies Hungarian permutation matching followed by Kabsch rotational alignment.The reordering can be parallelized before or during training, although it adds computational overhead.
4 Equivariant optimal transport flow matching: The approximation produces nearly optimal-transport integration paths for equivariant flows even with small batch sizes.In the infinite-batch limit, the Euclidean cost still yields the correct optimal transport map for invariant densities.

5 Architecture

The architecture parameterizes the flow vector field with an O(D)- and S(N)-equivariant graph neural network. Centering operations ensure that the resulting field conserves the geometric center.

5 Architecture: The vector field vθ(t, x) uses an O(D)- and S(N)-equivariant graph neural network composed of L consecutive layers.The particle updates are computed through the graph-network architecture.
5 Architecture: Table 1 compares flows trained with different objectives, reporting errors computed over three runs.The supplied caption specifies the comparison scope but not the individual table values.
5 Architecture: Geometric-center conservation is enforced directly when particles share a type and otherwise by subtracting the geometric center after the final layer.This construction is combined with a symmetric mean-free prior distribution.

6 Experiments

Experiments on symmetric particle systems and alanine dipeptide show that equivariant OT flow matching improves path efficiency and scalability, while performance varies by system size and metric.

6.1 DW4 and LJ13: For DW4 and LJ13, both flow matching objectives outperform likelihood-based training while being computationally more efficient.Effective sample sizes and negative log likelihoods are comparable between the two flow matching runs.
6.2 LJ55: A three-times speed-up over OT flow matching is achieved for LJ55 using a fixed 20-step integrator, while equivariant OT flow matching outperforms it on all evaluation metrics.Equivariant paths are nearly straight and incur minimal fixed-step integration error.
6.2 LJ55: For LJ55, OT flow matching has about 10 times larger mean batch transport cost and twice as long inference paths as equivariant OT flow matching.The resulting curvature also produces substantially larger fixed-step integration error for OT flow matching.
6.3 Alanine dipeptide: On alanine dipeptide, OT flow matching outperforms equivariant OT flow matching in effective sample size and negative log likelihood, although its integration paths are longer.Reweighting enables asymptotically unbiased samples from the semi-empirical target despite biased training data.
6.3 Alanine dipeptide: The Boltzmann generator’s free-energy difference agrees well with umbrella-sampling simulations, whereas relaxed training data and classical molecular dynamics overestimate it.The comparison concerns alanine dipeptide’s slowest transition.

7 Discussion

The paper introduces equivariant flow matching for invariant densities and demonstrates its value for scaling equivariant CNFs and modeling alanine dipeptide in Cartesian coordinates.

7 Discussion: The proposed objective yields optimal-transport integration paths even with small training batches and extends equivariant flows to the LJ55 cluster.Across four symmetric datasets, accounting for symmetries becomes more important as system size increases.
7 Discussion: The work reports the first successful Cartesian-coordinate Boltzmann generator for alanine dipeptide, with accurate free-energy differences under a semi-empirical force field.The approach uses OT flow matching and equivariant OT flow matching with equivariant flows.

8 Limitations / Future work

The approach remains limited by untested transferability and computational costs during inference and symmetry-aware training. Future work includes more expressive equivariant vector-field architectures.

8 Limitations / Future work: Transferability was not experimentally demonstrated, although the architecture and loss may support transferable models.The authors leave transferability for future research.
8 Limitations / Future work: Inference still requires complete vector-field integration, although equivariant OT flow matching permits faster fixed-step integrators.
8 Limitations / Future work: The Hungarian algorithm used in the approximation has O(N^3) complexity, motivating heuristic approximations.
8 Limitations / Future work: Future work should investigate alternative vector-field architectures to improve equivariant CNF modeling capabilities.

Supplementary Material

The supplementary experiments examine path geometry, memory use, free-energy estimation, molecular integration, minimized structures, reweighting, and OT-cost approximations. Together, they document shorter paths, lower memory requirements, accurate free-energy behavior, and practical approximation choices.

Memory and scalability: Less than 12 GB with batch size 256 versus more than 24 GB with batch size 5 demonstrates flow matching’s LJ55 memory advantage.
Memory and scalability: Less than 3 GB with batch size 256 versus more than 24 GB with batch size 16 demonstrates the same memory advantage for alanine dipeptide.
Training choices: Energy-based training is not used for these datasets because it requires evaluating additional training equations.
Alanine free energy: The alanine free-energy experiments compare umbrella sampling, equivariant-flow samples, and relaxed-MD samples along the slowest φ transition.
Integration paths and errors: Equivariant OT flow matching and OT flow matching produce comparable alanine integration-path differences despite alanine’s 22 atoms.
LJ55 minimized structures: Minimized structures from generated and test samples agree well and both include samples at the global minimum.
LJ55 reweighting: The equivariant OT flow-matching Boltzmann generator has lower LJ55 log-weight variance, which generally corresponds to higher effective sample size.
Approximation methods: Applying the OT approximation multiple times slightly reduces transportation cost, whereas rotating first gives inferior results; the chosen approximation is sufficient for a close OT-map approximation.

A.10 Different data set sizes

Reducing the alanine dipeptide training-set size worsens NLL, while naïve flow matching produces longer paths with similar ESS and likelihoods to OT flow matching.

A.10 Different data set sizes: NLL becomes worse as the alanine dipeptide training set shrinks to 10,000 and 1,000 samples.
A.10 Different data set sizes: Naïve flow matching produces longer integration paths, while ESS and likelihoods remain close to OT flow matching.
A.10 Different data set sizes: A non-equivariant fully connected architecture failed on all systems except DW4 and was not reported further.

B.1 Equivariant OT flow matching

For invariant measures, equivariant optimal transport preserves symmetry and can be approximated by orbit-aligned sample pairs. These results motivate the equivariant flow-matching objective and equivariant normalizing flows.

B.1 Equivariant OT flow matching: Orbit-aware sample pairs have lower average cost than pairs from the ordinary discrete OT plan, with equality approached as sample size grows.
B.1 Equivariant OT flow matching: For G-invariant measures, the OT map is G-equivariant and the corresponding OT plan is G-invariant.
B.1 Equivariant OT flow matching: Using the orbit-aware cost ˜c, the OT map remains optimal and generated pairs are aligned to minimize within-orbit cost.
B.1 Equivariant OT flow matching: The orbit-aligned construction motivates the proposed equivariant flow-matching objective and equivariant normalizing flows.
B.1 Equivariant OT flow matching: The flow preserves the geometric center when all particles have the same type.

C.2 Benchmark systems

The benchmark suite covers DW4, Lennard-Jones clusters, and alanine dipeptide, using system-specific data-generation procedures, model settings, training schedules, and compute resources. Equivariant OT training requires evaluating all sample pairs, but batch generation and reordering can be parallelized to reduce this burden.

Benchmark systems: The benchmarks include DW4, LJ13, LJ55, and alanine dipeptide, with LJ13 and LJ55 generated by MCMC and alanine dipeptide generated from molecular-dynamics trajectories followed by relaxation.The alanine dipeptide data use separate Amber and GFN2-xTB force fields for trajectory generation and relaxation.
Implementation: Training uses dataset-dependent model sizes and schedules, with batches reordered before training except for alanine dipeptide trained with equivariant OT flow matching.The supplied passages refer to Tables 7 and 8 for model hyperparameters and training schedules.
Computational considerations: Equivariant OT flow matching has high batch cost because Equation (16) evaluates every possible sample pair, motivating parallel CPU batch generation and reordering.The reported reordering time is summarized in Table 9, while larger batch sizes can be obtained at little additional cost through parallelization.
Benchmark systems: Alanine dipeptide training samples are reweighted with a von Mises-based bias to increase representation of the less-probable positive-φ state.The weights are defined as ω(φ) = 150 · f_vM(φ|μ = 1, κ = 10) + 1, after which samples are drawn from the weighted distribution.
Experimental setup: The experiments use 5 × 10^5 samples for DW4 and LJ13, 2 × 10^5 for alanine dipeptide, and 1 × 10^5 for LJ55 per model.Reported errors generally use one standard deviation averaged over three runs, with stated exceptions for batch transport and integration-error plots.