[Paper Review] Equivariant Flows: Exact Likelihood Generative Learning for Symmetric Densities

Jonas Köhler, Leon Klein+1 more|arXiv (Cornell University)|Jun 3, 2020

Machine Learning in Materials Science45 references51 citations

TL;DR

The paper introduces symmetry-preserving (equivariant) normalizing flows for exact-likelihood Boltzmann-generating flows, improving sampling efficiency and generalization on symmetric many-body systems.

ABSTRACT

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 extit{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.

Motivation & Objective

• Motivate the need to incorporate physical symmetries into density estimation and sampling models.
• Propose a theoretical criterion showing when equivariant flows preserve target symmetries in density generation.
• Provide practical, tractable constructions of equivariant flows for many-body particle systems.
• Demonstrate improved generalization and sampling efficiency over non-equivariant baselines on benchmark physical systems.

Proposed method

• Formalize symmetry via group actions on R^n and prove that H-equivariant diffeomorphisms mapped from a G-invariant density yield an H-invariant push-forward.
• Construct equivariant flows using continuous normalizing flows with an H-equivariant vector field, enabling exact density change via divergence computed in closed form.
• Design an invariant potential-based gradient field v(x)=∇Φ(x) where Φ is H-invariant to ensure v is H-equivariant.
• Employ a simple mixture of Gaussian radial basis functions to implement the vector field with tractable, exact divergence computations.
• Utilize an exact-divergence approach to avoid Hutchinson-type estimators that scale poorly with particle count.
• Benchmark on symmetric particle systems (DW-2, DW-4, LJ-13) with full rotation, translation, and permutation symmetries.

Experimental results

Research questions

• RQ1 How can normalizing flows be constructed to respect the symmetries of the target density in many-body systems?
• RQ2 Do symmetry-preserving (equivariant) flows generalize better and sample more efficiently than non-equivariant flows on symmetric energies?
• RQ3 Can exact divergence be efficiently computed in equivariant CNFs to enable unbiased reweighting in Boltzmann-generating flows?
• RQ4 What is the impact of equivariance on discovering meta-stable states in symmetric potentials?
• RQ5 How do DTO vs OTD training regimes compare in the context of equivariant flows for these systems?

Key findings

• Equivariant flows guarantee H-invariance of the generated density when G>H and f is H-equivariant, ensuring symmetry is built into the model by design.
• An explicit, tractable implementation using a gradient field of an invariant potential yields exact divergence and stable, efficient training.
• Equivariant flows generalize well with limited data and outperform non-equivariant flows, especially when data augmentation is used for symmetry (per DW-4 and LJ-13).
• In Boltzmann-generating setups, equivariant flows achieve substantial overlap with the target distribution and enable asymptotically unbiased reweighting, unlike some non-equivariant configurations.
• Experiments show equivariant models discover more meta-stable states and better match energy landscapes than non-equivariant baselines.