PaperlayerView exact revision on arXiv

Source-linked AI summary

Equivariant Flows: Exact Likelihood Generative Learning for Symmetric Densities

Jonas Köhler, Leon Klein, Frank Noé

arXiv:2006.02425v2stat.MLcs.LGphysics.chem-phphysics.comp-ph

TL;DR

The paper addresses how to build natural symmetries of physical and chemical many-body densities into exact-likelihood generative flows. It establishes that equivariant flows preserve corresponding density invariances and develops gradient-flow constructions, with experiments showing better generalization than non-equivariant flows. Numerical integration and divergence estimation remain important implementation constraints.

  • Problem

    Ignoring rotations and permutations in flow-based density estimation and enhanced sampling can lead to inferior results, motivating symmetry-preserving architectures.

  • Method

    The paper constructs equivariant normalizing flows using symmetry-preserving diffeomorphisms and gradient fields derived from invariant potential functions.

  • Results

    Equivariant flows approximate symmetric particle-system densities and generalize beyond biased data, whereas non-equivariant normalizing flows cannot.

  • Takeaways & Limitations

    Incorporating target-distribution symmetries into density-estimation and sampling mechanisms can improve generalization and be critical for successful learning.

  • Takeaways & Limitations

    Divergence estimation with off-the-shelf automatic differentiation can impose infeasible high-dimensional overhead, while Hutchinson estimation can make importance weighting ineffective.

Abstract

from arXiv · show

Normalizing flows are exact-likelihood generative neural networks which approximately transform samples from a simple prior distribution to samples of the probability distribution of interest. Recent work showed that such generative models can be utilized in statistical mechanics to sample equilibrium states of many-body systems in physics and chemistry. To scale and generalize these results, it is essential that the natural symmetries in the probability density -- in physics defined by the invariances of the target potential -- are built into the flow. We provide a theoretical sufficient criterion showing that the distribution generated by \textit{equivariant} normalizing flows is invariant with respect to these symmetries by design. Furthermore, we propose building blocks for flows which preserve symmetries which are usually found in physical/chemical many-body particle systems. Using benchmark systems motivated from molecular physics, we demonstrate that those symmetry preserving flows can provide better generalization capabilities and sampling efficiency.

1. Introduction

Many-body densities often have physical symmetries, but ignoring them can impair flow-based density estimation and sampling. The paper introduces equivariant flows and demonstrates improved generalization and sampling-related performance on particle systems.

  • Physical systems commonly exhibit symmetries including global rotations and permutations of identical elements.
  • Ignoring these symmetries in flow-based density estimation and enhanced sampling can lead to inferior results.
  • Equivariant flows provide symmetry-preserving generative models that satisfy the exact-likelihood requirements of Boltzmann generators.
  • Equivariant flows approximate highly symmetric particle-system densities and generalize beyond biased data, whereas non-equivariant normalizing flows cannot.
  • The framework includes a numerically tractable and efficient implementation using gradient flows derived from a simple mixture potential.
  • Considering target-distribution symmetries in density-estimation and sampling architectures can improve generalization and may be critical for successful learning.

2. Related Work

Normalizing flows model densities through invertible transformations and tractable likelihoods, while Boltzmann generators target known energy-based distributions for efficient sampling. Their exact likelihood enables reweighting generated samples toward the target density.

  • Normalizing Flows: Normalizing flows are diffeomorphisms that transform samples from a simple prior into generated samples and define densities through change of variables.
  • Normalizing Flows: Tractable flow densities support variational inference and approximate energy-based sampling, which can be made exact using importance sampling.
  • Boltzmann-Generating Flows: Boltzmann generators learn to sample a known target density proportional to exp(−u(x)) rather than estimate an unknown density solely from samples.
  • Boltzmann-Generating Flows: A Boltzmann generator combines an exact-likelihood model approximating the Boltzmann target with a statistical-mechanics reweighting algorithm.
  • Boltzmann-Generating Flows: Exact likelihood is required for reweighting, and flows used to generate asymptotically unbiased target samples are called Boltzmann-generating flows.
  • Boltzmann-Generating Flows: Energy-based training matches the energy corresponding to the generated density with the target energy function.

3. Invariant Densities via Equivariant Flows

The paper gives a sufficient condition for preserving density symmetries: pushing a G-invariant density through an H-equivariant diffeomorphism yields an H-invariant density. In particular, equivariant transformations of isotropic normal proposals preserve the relevant subgroup symmetry.

  • The framework assumes that prior and target densities over R^n share symmetries represented by a group acting on the space.
  • If a density is G-invariant and f is an H-equivariant diffeomorphism with G > H, the push-forward density is H-invariant.
  • For H < O(n), pushing an isotropic normal distribution through an H-equivariant diffeomorphism produces an H-invariant proposal density.

4. Constructing Equivariant Flows

The paper constructs equivariant continuous normalizing flows from symmetry-preserving vector fields, proving that their push-forward densities inherit the target symmetries. It also addresses numerical challenges in integration, training, and divergence computation, while contrasting the framework with Hamiltonian flows.

  • Theoretical framework: A H-equivariant vector field generates an H-equivariant bijection Fv,T, so continuous normalizing flows preserve the corresponding density symmetry.The flow is defined through characteristics of a time-dependent vector field and its integrated divergence.
  • Theoretical framework: Equivariant diffeomorphisms map G-invariant densities to H-invariant push-forward densities when H is a subgroup of G.This theorem supplies the sufficient criterion underlying the construction of symmetry-preserving flows.
  • Vector-field construction: Taking the gradient of an invariant potential produces an equivariant vector field, providing a generic construction for the required flow dynamics.For the many-body setting, the authors use an invariant potential under the relevant rotations and particle permutations.
  • Numerical implementation: Discrete integration does not guarantee invertibility, and adaptive solvers such as Dopri5 may require hundreds of vector-field evaluations for numerical accuracy.The authors identify numerical integration as one of three major implementation challenges.
  • Numerical implementation: DTO training offers similar results to OTD while providing faster iterations, more robust training, and faster convergence for the presented architecture.The comparison concerns adjoint-based optimize-then-discretize training versus fixed-grid discretize-then-optimize training.
  • Numerical implementation: Exact divergence computation avoids the infeasible O(n) autodifferentiation overhead and the high-variance Hutchinson estimator that can make importance reweighting ineffective.The proposed Gaussian radial-basis-function vector field permits a numerically exact divergence for general symmetries studied in the paper.
  • Relation to Hamiltonian flows: Augmenting a density with momentum variables creates an artificial Hamiltonian formulation whose finite-step conditional momentum distribution generally differs between source and target.Using an isotropic normal joint prior therefore requires a variational approximation of the target conditional momentum distribution.
  • Relation to Hamiltonian flows: Hamiltonian equivariant flows are a special case for linearly represented symmetries, but they handle broader spaces or nonlinear representations than the presented framework.Hamiltonian flows also avoid the numerical complications through symplectic integration and volume preservation.

5. Sampling of Coupled Particle Systems

The paper constructs equivariant continuous normalizing flows for symmetric many-body densities by using invariant potentials and exact divergence computation. Pairwise distance couplings provide an efficient implementation whose complexity can be controlled during training.

  • Symmetry-preserving flow design: The benchmark systems contain rotation, permutation, and center-of-mass symmetries, motivating flows whose vector fields respect these transformations.The target systems consist of interchangeable particles coupled through a potential energy and follow Boltzmann-type equilibrium distributions.
  • Symmetry-preserving flow design: Because simultaneous rotation and permutation symmetry prevents equivariant autoregressive or coupling layers, the method uses continuous residual flows with an equivariant vector field.The vector field is constructed from the gradient of an invariant potential.
  • Symmetry-preserving flow design: The invariant potential is formed from pairwise couplings over particle distances, producing per-particle updates from relative positions and distance-dependent forces.The pairwise construction uses r_ij(t), d_ij(t), and a coupling potential over distance and time.
  • Symmetry-preserving flow design: Distance and time features are expanded with Gaussian radial basis functions and mixed through a trainable shared weight matrix.This architecture combines distance-dependent and time-dependent components in the flow update.
  • Efficient computation: The gradient and divergence are computed exactly as one vectorized operation, avoiding the high-dimensional overhead of stochastic divergence estimation.The implementation is designed to remain numerically tractable for many-body particle systems.
  • Efficient computation: Weight decay and bandwidth control regulate the dynamics’ complexity while enabling fixed-grid integration with negligible integration error.The paper reports that a small amount of weight decay is sufficient for optimization with a fixed-grid solver.
  • Efficient computation: Compared with invariant neural-network potentials, the simple pairwise couplings are faster, use fewer parameters, and consistently outperform them on the studied target systems.Neural-network potentials would require automatic differentiation for gradients and repeated Laplacian computation during integration.

6. Benchmark Systems

The evaluation uses two symmetric coupled-particle systems: a frustrated double-well model and a 13-particle Lennard-Jones cluster. Their distinct energy landscapes provide benchmarks for sampling and structure-generation methods.

  • Benchmark Systems: The experiments study two systems in which rotation, permutation, and center-of-mass symmetries are all present.These systems are selected to evaluate symmetry-preserving flows under coupled many-body dynamics.
  • Double-well systems: The DW-2 and DW-4 systems contain 2 to 4 particles in two dimensions with pairwise double-well interactions.The double-well potential creates two low-energy modes separated by an energy barrier.
  • Double-well systems: Coupling multiple double-well particles produces frustrated systems with multiple metastable states.The design parameters and dimensionless temperature determine the resulting energy landscape.
  • Lennard-Jones system: The LJ-13 system uses 13 particles in three dimensions with a Lennard-Jones potential, representing solid-state models and rare-gas clusters.Its complex energy landscape contains energy minima that are difficult to find and sample between.
  • Lennard-Jones system: The Lennard-Jones cluster is used as a benchmark because its complex landscape challenges structure-generation methods.A small additional term prevents particle dissociation at finite sampling temperature.
  • Benchmark Systems: Figure 2 compares the energy contribution per distance for the double-well and Lennard-Jones potentials.The two panels correspond to the double-well and Lennard-Jones model systems.

7. Experiments

Experiments show that exact divergence computation makes equivariant flows computationally practical, while symmetry-aware models generalize better and discover more metastable states than non-equivariant alternatives.

  • 7.3. Statistical Efficiency for Density Estimation: Hutchinson divergence-estimation error grows with particle count, making importance reweighting impossible even for the simple DW-2 system.Brute-force divergence computation is also impractical because of its wall-clock overhead.
  • 7.3. Statistical Efficiency for Density Estimation: Exact closed-form traces provide fast divergence computation while retaining exact reweighting.The comparison covers brute-force automatic differentiation, Hutchinson estimation, and exact trace computation.
  • 7.3. Statistical Efficiency for Density Estimation: Regularized DTO and OTD training achieve similar test likelihoods, but OTD incurs a growing function-evaluation overhead while DTO remains flat.The experiment compares OTD with dopri5 against DTO using a fixed 20-step fourth-order Runge–Kutta grid.
  • 7.3. Statistical Efficiency for Density Estimation: Equivariant flows generalize to unseen trajectories in the low-data regime, whereas augmented non-equivariant flows perform worse or fail to fit the data.Without augmentation, non-equivariant flows show strong overfitting; the equivariant flow uses 620 parameters versus 5256 for DW-4 and 21671 for LJ-13.
  • 7.4. Equivariance in Boltzmann-Generating Flows: The equivariant flow overlaps the target distribution and supports asymptotically unbiased reweighting, unlike the tested non-equivariant flows.The unaugmented non-equivariant flow collapses to one mode, while the augmented version fails to produce accurate low-energy states.
  • 7.5. Discovery of Meta-Stable States: The equivariant flow discovers all five DW-4 minima from data concentrated near one minimum, while the non-equivariant model reproduces only its training minimum.For LJ-13, it outperforms short MCMC sampling, consistently finds the unseen global minimum, and performs nearly as well as a trajectory with 100 times more evaluations.

8. Discussion

The paper presents symmetry-preserving normalizing flows with structural advantages over black-box continuous normalizing flows and demonstrates their generalization on particle systems.

  • 8. Discussion: The framework incorporates density symmetries into normalizing-flow structure and demonstrates superior generalization on two physics-motivated particle systems.The proposed equivariant gradient field has an analytically computable divergence, explicit numerical-stability handling, and few parameters.
  • 8. Discussion: Equivariant gradient fields use a simple mixture potential to provide an analytically computable divergence and few parameters.These are identified as structural advantages over black-box CNFs.

A. Proofs and derivations

The paper establishes conditions under which equivariant transformations preserve symmetry-related densities, then constructs symmetry-preserving flow variants for particle systems. It also compares these architectures with alternative equivariant and non-equivariant flows in density-estimation experiments.

  • A. Proofs and derivations: If a density is invariant under an invertible linear transformation, that transformation has determinant ±1.This determinant property supports the change-of-variables argument used in the symmetry proof.
  • A. Proofs and derivations: An H-equivariant diffeomorphism preserves an invariant density when the target density is invariant under a larger symmetry group G containing H.The proof combines equivariance of the map with the determinant result for the representation of H.
  • A. Proofs and derivations: Center-of-mass subtraction is represented as a linear projection with rank (N −1)D onto the center-of-mass-free subspace.The projection matrix satisfies A^2 = A and A^T = A.
  • A. Proofs and derivations: For an isotropic Gaussian prior, projecting samples with A produces the corresponding density on the center-of-mass-free subspace.The projected covariance is AAT, and the densities agree on vectors already in the subspace.
  • A. Proofs and derivations: Equivariant flows map center-of-mass-free configurations to center-of-mass-free configurations, defining a flow on the projected subspace.This follows from equivariance with respect to the stated symmetries.
  • B.1. Comparison with other equivariant gradient flows: The proposed kernel flow is compared with simple-gradient and SchNet-inspired equivariant gradient flows, while non-equivariant coupling flows use rNVP transformations with center-of-mass augmentation.The kernel flow is reported as considerably faster during training and at least one order of magnitude faster during sampling; rNVP performance depends strongly on augmentation.

B.3. Energy based training for the DW-4 system

The DW-4 energy-based experiment trains models on MCMC data and compares their energies, reweighted distributions, likelihoods, and learned kernel dynamics. Equivariant models support low-energy sampling and correct reweighting, while non-equivariant alternatives show weaker or less reliable behavior.

  • B.3. Energy based training for the DW-4 system: 10,000 samples are evaluated by comparing their energies and reweighted energies with expectations from a long MCMC trajectory.Models are pretrained on 1,000 MCMC samples and then trained with a KL-loss weight increasing from λ = 0 to λ = 0.5.
  • B.3. Energy based training for the DW-4 system: All equivariant flows generalize quickly to unseen trajectories, although the kernel flow obtains the lowest train and test log-likelihoods.The comparison covers both DW-4 and LJ-13 systems as the training-set size increases.
  • B.3. Energy based training for the DW-4 system: Equivariant models produce low energies and allow correct reweighting for the DW-4 system.This includes the kernel flow and the gradient flow with SchNet.
  • B.3. Energy based training for the DW-4 system: The kernel flow’s learned dynamics can be visualized by plotting its kernel weights W after training.The visualization is provided for both DW-4 and LJ-13.
  • B.3. Energy based training for the DW-4 system: The non-equivariant coupling flow with data augmentation performs significantly worse than the equivariant flow on both test systems.For LJ-13, it cannot fit the augmented data and remains close to the prior.

C.7. Error bars

The reported uncertainty and variability are defined by the number of repeated runs or resampled estimates used for each experiment. Error bars represent one standard deviation, with evaluation protocols varying across figures.

  • C.7. Error bars: Error bars in all plots represent one standard deviation.This convention applies throughout the reported experiments.
  • C.7. Error bars: Figure 3 uses 1,000 estimations per particle count for errors in panel a, 100 reweighted bootstrap subsamples in panel b, and 100 estimations per method in panel c.The panels therefore use distinct replication procedures for error estimates and timing.
  • C.7. Error bars: Figure 4a reports three runs per method, whereas Figures 5, 8, and 9 report five runs per model, system, and training-set size.The number of independent runs is therefore figure-specific.
  • C.7. Error bars: All experiments were conducted on a GeForce GTX 1080 Ti with 12 GB RAM.This hardware specification contextualizes the reported computational measurements.
Loading 2006.02425v2…