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

Introduction

Generative learning using exact-likelihood methods based on invertible transformations has had remarkable success in domains like accurately representing images , audio , 3D point cloud data as well as applications in physics .

Recently, Boltzmann Generators (BG) have been introduced for sampling Boltzmann type distributions $p(x)\propto\exp(-u(x))$ of high-dimensional multi-body problems, such as valid conformations of proteins. In contrast to typical generative learning problems, the target density $p(x)$ is specified by definition of the multi-body energy function $u(x)$ and the difficulty lies is learning to sample it efficiently. BGs do that in two steps: (1) An exact-likelihood method that generates samples from a density $q_{X,\theta}(x)$ that approximates the Boltzmann density $p(x)$. (2) An algorithm to reweigh the generated density to the target density $p(x)$.

In , it has been demonstrated that such BGs can be trained to efficiently sample prototypical many-body systems such as proteins in implicit solvent and condensed matter systems. In order to make further progress, it is essential to develop exact-likelihood generative models that respect the symmetries of $u(x)$, for example invariance of the energy with respect to global rotation or permutations of identical particles.

In this work, we derive a sufficient criterion to design symmetry invariant BGs and design a simple instance that is symmetric with respect to translations, rotations and permutations. By comparing this model with the non-symmetric approach based on RealNVP layers as used in , we show that encoding symmetries in the model is critical for generalization.

Boltzmann-generating flows

In order to reweigh the generated density $q_{X,\theta}(x)$ to the target density $p_{X}(x)\propto\exp(-u(x))$ (step 2 above), BGs require an exact likelihood generative model. This can be achieved by transforming a simple prior density $q_{Z}(z)$, e.g. a multivariate Normal distribution, via an invertible function $f_{\theta}$ . Sampling from $q_{X,\theta}(x)$ is achieved by sampling $z\sim q_{Z}(z)$ and transforming to $f_{\theta}(z)\sim q_{X,\theta}(x)$. Due to the invertibility of $f_{\theta}$, the probability density of any generated point can be computed with the change of variables equation

Two obvious options are available in order to train the generator to match $q_{X,\theta}\approx p_{X}$ :

ML-training: If some data $\left\{x_{n}\right\}_{n=1\ldots N}$ is given that at least represents one or a few high-probability modes of $p(x)$, we can maximize the likelihood under the model, as is typically done when training flow-based models for images.

KL-training We minimize the reverse Kullback-Leibler divergence $KL(q_{X,\theta}\|p_{X})$. This approach is also known as energy-based training where the energy corresponding to the generated density is matched with $u(x)$.

Equivariant Flows and Boltzmann Generators

Symmetries may be discussed in terms of a group $G$ acting on a vector space $V$. A representation $\rho$ of $G$ is a map $\rho\colon G\rightarrow GL(V)$ satisfying $\rho(g)\rho(h)=\rho(gh)$. Examples are permutations and rotations in 3D which can be represented by permutation and rotation matrices.

Our first result is the following theorem, which gives a sufficient criterion, for when a transformed density resulting from a normalizing flow is symmetric

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.

2 Equivariant dynamics yield equivariant continuous normalizing flows

A recently proposed approach to flow based modeling describes the bijection implicitly via solving the Cauchy problem

for a fixed time horizon $[t_{0},t_{1}]$, where $g_{\theta}$ can be a freely chosen / learned dynamics function. The solution of this initial value problem is $x:=y(t_{1})=y(t_{0})+\int_{t_{0}}^{t_{1}}dtg_{\theta}(y(t),t)$, which (under mild constraints) gives a bijection $z\rightarrow x$

As shown in , the change of density can be computed at any integration time, using the continuous change of variable rule, given by the differential equation

Integrating eqs. (1) & (2), e.g. using a black-box solver, yields an exact-likelihood generative model.

For such continuous normalizing flows (CNF), a statement similar to theorem 1 can be made. Our second result is:

Let $f$ be a $G$-equivariant dynamics function. If $q_{Z}(z)$ is a $G$-invariant prior density, $z$ is sampled from $q_{Z}(z)$ and $x$ is obtained from $z$ by solving (1). Then $x$ is distributed according to a $G$-invariant density.

3 Modeling symmetric multi-body systems with normalizing flows

Permutation invariance: Swapping the labels of any two interchangeable particles $(x_{1},\ldots,x_{K})\rightarrow(x_{\sigma(1)},\ldots,x_{\sigma(K)})$.

Rotation invariance: Any 2D/3D rotation of the system $Rx=(Rx_{1},\ldots,Rx_{K})$.

Translation invariance: Any 2D/3D translation $x+v=(x_{1}+v,\ldots,x_{K}+v)$.

A natural prior distribution $q_{Z}(z)$ being invariant to all these symmetries and thus satisfying the conditions of thm. 1 can be obtained by sampling ${z^{\prime}=(z^{\prime}_{1},\ldots,z^{\prime}_{N})\sim\mathcal{N}(0,I_{D\times D})}$ i.i.d. and subtract its center of mass $\mu(x):=\sum_{i=1}^{K}x_{i}$ yielding $z=z^{\prime}-\mu(z^{\prime})$. For a mean-free configuration $z$, we can evaluate its likelihood by $\log q_{Z}(z)=-1/2\sum_{i}^{K}\|z_{i}\|^{2}+const.$, which is invariant with respect to the symmetries 1.-3. above.

4 Equivariant flows for multi-body systems

We give a simple example of an equivariant flow which obeys the symmetries 1.-3. above. For this, we employ thms. 1 & 2 to design equivariant dynamics functions with tractable exact traces.

Each particle is updated using a vector field in radial direction

which can be computed exactly in an efficient way using one backpropagation pass through the graph.

It is simple to see that this vector field is equivariant with respect to symmetries 1.-3.. Further this field does not alter the center of mass of a particle system, thus integrating mean-free states along it will always result in mean-free states.

Experiments

Using a simple 2D toy system, we show empirically, that a Boltzmann Generator using an equivariant flow encoding the symmetries of the system’s energy function explicitly as given in the last section (eqBG) can generalize beyond the training data whereas a non-equivariant approach (furthermore abbreviated as nBG) based on RealNVP transformations cannot.

Our system ($N=4,D=2$) is given by the pairwise potential

which produces two distinct low energy modes separated by a high energy barrier (see Figure 1 e). By coupling multiple particles with such double-well interactions we can create a frustrated system with multiple metastable states.

In a first experiment, we compare how both models generalize to unseen trajectories. For this we use standard MCMC sampling using the classic Metropolis-Hasting algorithm to produce short out-of-equilibrium trajectories of the system. We split a trajectory using the first 50% of samples for initializing both, the nBG and the eqBG, with ML training. After that we fine-tune both models on the energy using a weighted sum of ML and KL training. Both models are able to sample states with acceptable energies (Figure 1 a, c). However, evaluating both models on the second 50% of the trajectory, we clearly see that the nBG is not generalizing at all, whereas the eqBG is able to achieve similar likelihoods on train and test set simultaneously (Figure 1 b, d, Table 1).

2 Second experiment: discovery of unseen metastable states

In a second experiment, we evaluate to which extend both models are able to discover new metastable states, which have not been observed in the data set. Here we (1) create a training set, by perturbing a single minimum state with a tiny amount of Gaussian noise (Figure 1 f), (2) initialize both, the nBG and eqBG, with ML training on the training set, (3) fine-tune both models using a weighted sum ML and KL training on the energy, (4) generate samples from the converged models (Figure 1 g,h, top row) and (5) locally minimize the energies $u(x)$ of the generated structures $x$ using a non-momentum second-order optimizer (Figure 1 g, h, bottom row). As can be seen, the nBG will only reproduce the unique state from which the training set was constructed. However, the eqBG will produce a variety of local minima, which did not appear in the training set.

Discussion

We gave a sufficient criterion, by which equivariant flows can be constructed and implemented an instance for symmetries frequently appearing in multi-particle systems. Using a simple example we have demonstrated that encoding symmetries of the target density/energy into the invertible function is critical, if we expect generative flows such as BGs to generalize beyond seen trajectories. Scaling such equivariant flows and BGs to large real-world systems remains a major challenge for future research.

References