Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Jonas Köhler, Leon Klein, Frank Noé

TL;DR

The paper tackles efficient sampling of Boltzmann-type distributions for high-dimensional multi-body systems by enforcing energy symmetries directly in flow-based generators. It introduces a theoretical framework showing that using a $G$-invariant prior with a $G$-equivariant transform yields $G$-invariant densities, and extends this to continuous normalizing flows with equivariant dynamics. An explicit equivariant flow is demonstrated for particle systems with translation, rotation, and permutation invariances, improving generalization over non-equivariant baselines. Experiments on a toy multi-body system show that equivariant Boltzmann Generators can generalize unseen trajectories and uncover unseen metastable states, highlighting the practical impact of symmetry-aware generative modeling in physics-inspired sampling tasks.

Abstract

Flows are exact-likelihood generative neural networks that transform samples from a simple prior distribution to the samples of the probability distribution of interest. Boltzmann Generators (BG) combine flows and statistical mechanics to sample equilibrium states of strongly interacting many-body systems such as proteins with 1000 atoms. In order to scale and generalize these results, it is essential that the natural symmetries of the probability density - in physics defined by the invariances of the energy function - are built into the flow. Here we develop theoretical tools for constructing such equivariant flows and demonstrate that a BG that is equivariant with respect to rotations and particle permutations can generalize to sampling nontrivially new configurations where a nonequivariant BG cannot.

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

TL;DR

The paper tackles efficient sampling of Boltzmann-type distributions for high-dimensional multi-body systems by enforcing energy symmetries directly in flow-based generators. It introduces a theoretical framework showing that using a -invariant prior with a -equivariant transform yields -invariant densities, and extends this to continuous normalizing flows with equivariant dynamics. An explicit equivariant flow is demonstrated for particle systems with translation, rotation, and permutation invariances, improving generalization over non-equivariant baselines. Experiments on a toy multi-body system show that equivariant Boltzmann Generators can generalize unseen trajectories and uncover unseen metastable states, highlighting the practical impact of symmetry-aware generative modeling in physics-inspired sampling tasks.

Abstract

Flows are exact-likelihood generative neural networks that transform samples from a simple prior distribution to the samples of the probability distribution of interest. Boltzmann Generators (BG) combine flows and statistical mechanics to sample equilibrium states of strongly interacting many-body systems such as proteins with 1000 atoms. In order to scale and generalize these results, it is essential that the natural symmetries of the probability density - in physics defined by the invariances of the energy function - are built into the flow. Here we develop theoretical tools for constructing such equivariant flows and demonstrate that a BG that is equivariant with respect to rotations and particle permutations can generalize to sampling nontrivially new configurations where a nonequivariant BG cannot.

Paper Structure

This paper contains 11 sections, 2 theorems, 6 equations, 1 figure, 1 table.

Table of Contents

  1. Introduction
  2. Boltzmann-generating flows
  3. Equivariant Flows and Boltzmann Generators
  4. Equivariant normalizing flows yield symmetric densities
  5. Equivariant dynamics yield equivariant continuous normalizing flows
  6. Modeling symmetric multi-body systems with normalizing flows
  7. Equivariant flows for multi-body systems
  8. Experiments
  9. First experiment: generalization on unseen trajectories
  10. Second experiment: discovery of unseen metastable states
  11. Discussion

Key Result

Theorem 1

Let $q_Z(z)$ be a $G$-invariant prior density and $f_{\theta}$ be a $G$-equivariant bijection. Let $q_{X, \theta}(x)$ be the density of $x = f(z)$ for $z \sim q_{Z}(z)$. Then $q_{X, \theta}$ is $G$-invariant.

Figures (1)

  • Figure 1: Comparison: sampling equivariant and non-equivariant Boltzmann Generators. a, c) Energies for samples taken from trained nBG and eqBG respectively. b, d) Marginal distribution of the first particle position $x_1$ for states sampled from trained nBG and eqBG respectively. e) Potential energy of the system illustrated for $N=2$. f-h)Top row: sampled states from training data, trained nBG and trained eqBG respectively together with corresponding energies. Bottom row: corresponding configurations after energy minimization.

Theorems & Definitions (2)

  • Theorem 1
  • Theorem 2