Learning Equivariant Energy Based Models
with Equivariant Stein Variational Gradient Descent

Many real-world observations comprise symmetries and admit probabilistic models that are invariant to such symmetry transformations. Naturally, overlooking these inductive biases while encoding such domains will lead to models with inferior performance capabilities. In this paper, we focus on the problem of efficient sampling and learning of equivariant probability densities by incorporating symmetries in probabilistic models.

We accomplish this by first proposing equivariant Stein variational descent algorithm in §3 for sampling from invariant densities. Stein Variational Gradient Descent (SVGD) is a kernel-based inference method that constructs a set of particles iteratively along an optimal gradient path in an RKHS to approximate and sample from a target distribution. We extend SVGD for invariant densities by considering equivariant kernel functions that evolve the set of particles such that the density at each time-step is invariant to the same symmetry transformations as encoded in the kernel. We demonstrate that equivariant SVGD is more sample efficient, produces a more diverse set of samples, and is more robust compared to regular SVGD when sampling from invariant densities.

Subsequently, in §4, we build equivariant Energy Based Models EBMs for learning invariant densities given access to i.i.d. data by leveraging the tremendous recent advances in geometric deep learning where the energy function is equivariant neural network. We train these equivariant EBMs through contrastive divergence by generating samples using equivariant SVGD. We show that incorporating the symmetries present in the data into the energy model as well as the sampler provides an efficient learning paradigm to train equivariant EBMs that generalize well beyond training data.

We empirically demonstrate the performance of equivariant EBMs using equivariant SVGD in §5. We consider real-world applications comprising of problems from many-body particle systems, molecular structure generation and, classification and generation for image datasets.

In this section we set-up our main problem, introduce key definitions and notations and formulate an approach to incorporate symmetries in particle variational inference optimization methods through Stein variational gradient descent. Along the way, we also discuss directly related work and relegate a detailed discussion on previous work to Appendix B.

Let $\mathcal{G}$ be a group acting on $\mathds{R}^{d}$ through a representation $\mathrm{R}:\mathcal{G}\to\mathrm{GL}(d)$ where $\mathrm{GL}(d)$ is the general linear group on $\mathds{R}^{d}$, such that $\forall g\in\mathcal{G}$, $g\to\mathrm{R}_{g}$. Given a target random variable $\mathsf{X}\subseteq\mathds{R}^{d}$ with density $\pi$, we say that $\pi$ is $\mathcal{G}$-invariant if $\forall g\in\mathcal{G}$ and $\boldsymbol{x}\in\mathds{R}^{d}$, $\pi(\mathrm{R}_{g}\boldsymbol{x})=\pi(\boldsymbol{x})$. Additionally, a function $f(\cdot)$ is $\mathcal{G}$-equivariant if $\forall g\in\mathcal{G}$ and $\boldsymbol{x}\in\mathds{R}^{d}$, $f(\mathrm{R}_{g}\boldsymbol{x})=\mathrm{R}_{g}f(\boldsymbol{x})$. We denote with $\mathcal{O}(\boldsymbol{x})$ the orbit of an element $\boldsymbol{x}\in\mathsf{X}$ defined as $\mathcal{O}(\boldsymbol{x}):=\{\boldsymbol{x}^{\prime}:\boldsymbol{x}^{\prime}=\mathrm{R}_{g}\boldsymbol{x},\forall g\in\mathcal{G}\}$. We call $\pi_{|\mathcal{G}}$ the factorized density of a $\mathcal{G}$-invariant density $\pi$ where $\pi_{|\mathcal{G}}$ has support on the set $\mathsf{X}_{|\mathcal{G}}$ where the elements of $\mathsf{X}_{\mathcal{G}}$ are indexing the orbits i.e. if $\boldsymbol{x},\tilde{\boldsymbol{x}}\in\mathsf{X}_{\mathcal{G}}$ then $\boldsymbol{x}\neq\mathrm{R}_{g}\tilde{\boldsymbol{x}},\forall g\in\mathcal{G}$. In this paper, we are interested to incorporate inductive biases given by symmetry groups to develop efficient sampling and learning paradigms for generative modelling. Precisely, we consider the following problems:

We begin this section by presenting the main result of this section by introducing equivariant Stein variational gradient descent (E-SVGD) by utilizing Lemma 1 and Equations (3) and (4).

Following Proposition 1, we call the updates in Equations (3) & (4) equivariant Stein variational gradient descent when the kernel $k(\cdot,\cdot)$ (and $\boldsymbol{K}(\cdot,\cdot)$ respectively) is invariant (equivariant) and the initial set of particles $\lbag\boldsymbol{x}^{0}_{1},\cdots,\boldsymbol{x}^{0}_{n}\rbag$ are sampled from an invariant density $q_{0}$. Thus, all that is required to sample from a $\mathcal{G}$-invariant density $\pi$ using equivariant SVGD is to construct a positive definite kernel that is $\mathcal{G}$-equivariant. Let us next give a few examples for constructing in- and equivariant positive definite kernels.

Example 1 is only restricted to finite groups and does not directly apply to continuous symmetry groups. We can construct kernels for continuous groups following Example 1 by either using a Monte Carlo approximation or using a transformation that performs computations in the factorized space $\mathsf{X}_{|\mathcal{G}}$ as we show in the next example.

Examples (1) and (2) are both invariant scalar kernels. Let us next give an example of an equivariant matrix valued kernel for matrix valued SVGD (cf. Equation 4).

In Section 3, we developed equivariant Stein variational gradient descent algorithm for sampling from invariant densities. In this section, we leverage the recent tremendous advances in deep geometric learning (Cohen and Welling,, 2016; Dieleman et al.,, 2016; Bronstein et al.,, 2021) to propose equivariant energy based models that are trained contrastively using our proposed equivariant Stein variational gradient descent algorithm to learn invariant (unnormalized) densities $\pi$ given access to i.i.d. samples $\lbag\boldsymbol{x}_{1},\boldsymbol{x}_{2},\cdots,\boldsymbol{x}_{n}\rbag\sim\pi$.

Given a set of samples $\lbag\boldsymbol{x}_{1},\boldsymbol{x}_{2},\cdots,\boldsymbol{x}_{n}\rbag\subseteq\mathds{R}^{d}$, energy-based models (LeCun et al.,, 2006) learn an energy function $\mathrm{E}_{\boldsymbol{\theta}}(\boldsymbol{x}):\mathds{R}^{d}\to\mathds{R}$ that defines a probability distribution $\tilde{\pi}_{\boldsymbol{\theta}}(\boldsymbol{x})=\nicefrac{{\mathsf{exp}\big{(}-\mathrm{E}_{\boldsymbol{\theta}}(\boldsymbol{x})\big{)}}}{{Z_{\boldsymbol{\theta}}}}$, where $Z_{\boldsymbol{\theta}}=\int\mathsf{exp}\big{(}-\mathrm{E}_{\boldsymbol{\theta}}(\boldsymbol{x})\big{)}\mathop{}\!\mathrm{d}\boldsymbol{x}$ is the partition function. Unlike other popular tractable density models like normalizing flows, EBMs are less restrictive in the parameterization of the functional form of $\tilde{\pi}_{\boldsymbol{\theta}}(\cdot)$ since the energy function does not need to integrate to one, it can be parameterized using any nonlinear function. Conveniently, if $\pi$ is $\mathcal{G}$-invariant, we can use the existing equivariant deep network architectures to parameterize $\mathrm{E}_{\theta}(\cdot)$ to encode the symmetries into the energy network. Such an equivariant energy network defines an equivariant energy based model. EBMs are usually trained by maximizing the log-likelihood of the data under the given model:

However, evaluating $Z_{\boldsymbol{\theta}}$ is intractable for most (useful) choices of $\mathrm{E}_{\boldsymbol{\theta}}(\cdot)$ which makes learning EBMs via maximum likelihood estimation problematic. Contrastive divergence (Hinton et al.,, 2006) provides a paradigm to learn EBMs using maximum likelihood estimation without needing to compute $Z_{\boldsymbol{\theta}}$ by approximating the gradient of $\nabla_{\boldsymbol{\theta}}\mathcal{L}_{\mathsf{ML}}(\boldsymbol{\theta})$ in Equation 5 as follows:

Intuitively, the gradient in Equation 6 drives the model such that it assigns higher energy to the negative samples $\boldsymbol{x}^{-}$ sampled from the current model and decreases the energy of the positive samples $\boldsymbol{x}^{+}$ which are the data-points from the target distribution. Since, training an EBM using MLE requires sampling from the current model $\tilde{\pi}(\boldsymbol{\theta})$, successful training of EBMs relies heavily on sampling strategies that lead to faster mixing. Fortuitously, since $\mathrm{E}_{\boldsymbol{\theta}}(\cdot)$ in our present setting is $\mathcal{G}$-equivariant, we propose to use our equivariant sampler for more efficient training²²2compared to using a regular sampler with no encoded symmetries. of the equivariant energy based model.

Additionally, following Grathwohl et al., (2019), we can extend equivariant energy based models to equivariant joint energy models. Let $\{(\boldsymbol{x}_{1},y_{1}),(\boldsymbol{x}_{2},y_{2}),\cdots,(\boldsymbol{x}_{n},y_{n})\}\subseteq\mathds{R}^{d}\times[K]$ be a set of samples with observations $\boldsymbol{x}_{i}$ and labels $y_{i}$. Given a parametric function $f_{\boldsymbol{\theta}}:\mathds{R}^{d}\to\mathds{R}^{k}$, a classifier uses the conditional distribution $\tilde{\pi}_{\boldsymbol{\theta}}(y|\boldsymbol{x})\propto\mathsf{exp}(f_{\boldsymbol{\theta}}(\boldsymbol{x})[y])$ where $f_{\boldsymbol{\theta}}(\boldsymbol{x})[y]$ is the logit corresponding to the $y^{\text{th}}$ class label. As shown by Grathwohl et al., (2019), these logits can be used to define the joint density $\tilde{\pi}_{\boldsymbol{\theta}}(\boldsymbol{x},y)$ and marginal density $\tilde{\pi}_{\boldsymbol{\theta}}(\boldsymbol{x})$ as follows:

Hence, the energy function at a point $\boldsymbol{x}$ is given by $\mathrm{E}_{\boldsymbol{\theta}}=-\log\sum_{y}\mathsf{exp}(f_{\boldsymbol{\theta}}(\boldsymbol{x})[y])$ with joint energy $\mathrm{E}_{\boldsymbol{\theta}}(\boldsymbol{x},y)=-f_{\boldsymbol{\theta}}(\boldsymbol{x})[y]$. In our setting, the joint distribution $\pi(\boldsymbol{x},y)$ is $\mathcal{G}$-invariant in the first argument i.e. $\pi(\mathrm{R}_{g}\boldsymbol{x},y)=\pi(\boldsymbol{x},y),\forall g\in\mathcal{G}$. An example of such a setting is any image data-set where the class label does not change if the image is rotated by an angle. Using Equation 7, it suffices for the function $f_{\boldsymbol{\theta}}$ to be $\mathcal{G}$-equivariant to model a $\mathcal{G}$-invariant density $\tilde{\pi}_{\boldsymbol{\theta}}(\boldsymbol{x},y)$. Furthermore, a $\mathcal{G}$-equivariant $f_{\boldsymbol{\theta}}$ also makes the marginal density $\tilde{\pi}_{\boldsymbol{\theta}}(\boldsymbol{x})$ and conditional density $\tilde{\pi}_{\boldsymbol{\theta}}(y|\boldsymbol{x})$ $\mathcal{G}$-invariant in the input $\boldsymbol{x}$. We call such an energy model where $f_{\boldsymbol{\theta}}$ is equivariant to a symmetry transformation to be an equivariant joint energy model.

We can train this model by maximizing the log-likelihood of the joint distribution as follows:

where $\mathcal{L}_{\mathsf{SL}}(\boldsymbol{\theta})$ is the supervised loss which is the cross-entropy loss in the case of classification. Thus, an equivariant joint energy model can now be trained by applying the gradient estimator in Equation 6 for $\log\tilde{\pi}_{\boldsymbol{\theta}}(\boldsymbol{x})$ and evaluating the gradient of $\log\tilde{\pi}_{\boldsymbol{\theta}}(y|\boldsymbol{x})$ through back-propagation. Conveniently, Equation 8 can also be used for semi-supervised learning with $\mathcal{L}_{\mathsf{SL}}((\boldsymbol{\theta}))$ substituted with the appropriate supervised loss e.g. MSE for regression.

Let us end this section with an empirical example for learning a mixture of $C_{4}$-Gaussians (Figure 4) as shown in row two of the leftmost column of Figure 4. The innermost $C_{4}$-Gaussian defines the class conditional probability $\pi(\boldsymbol{x}|y=0)$ (row 3) and the outer $C_{4}$-Gaussian defines $\pi(\boldsymbol{x}|y=1)$ (row 4). We learn a non-equivariant joint EBM using vanilla SVGD (cf. Figure 4 center column) and an equivariant joint EBM using equivariant SVGD (cf. Figure 4 right column) keeping the number of iterations and particles the same for training. In Figure 4, we plot the decision boundaries learned by the model in the top row. The star marked samples in the figure are the samples generated by the underlying model. We plot the joint distribution and the class conditional distributions in row two-four respectively. The figure abundantly demonstrates the superior performance of an equivariant joint energy model trained using equivariant SVGD over its non-equivariant counterpart. A more detailed figure with comparisons to an equivariant joint energy model trained using vanilla SVGD is presented in Section D.1.

In this section, we present empirical analysis of equivariant EBMs and E-SVGD through experiments to (i) reconstruct potential function describing a many-body particle system (DW-4) trained using limited number of meta-stable states, (ii) model a generative distribution of molecules (QM9) and generate novel samples and (iii) hybrid (generative & discriminative) model invariant to rotations for FashionMNIST trained using dataset with no rotations. Due to space constraints, details about all the experiments as well as detailed figures are relegated to Appendix E.

DW-4: In this many-body particles system, a double-well potential describes the configuration of four particles that is invariant to rotations, translations and, permutation of the particles. This system comprises five distinct metastable states which are characterized as the mimina in the potential function. In our experiment, we show that given access to only a single example of each metastable state configuration, an equivariant EBM trained with E-SVGD can recover other states with similar energy as those of in the training set. In Figure 7, the first column shows the metastable states present in the training set. The second column are the states recovered by an EBM trained with vanilla SVGD which results in configurations that exactly copy the training set. The third column shows configurations generated by the equivariant model which are distinct from the training set but mimic the energies of the corresponding metastable states in the training set. Our setup is different from that of Köhler et al., (2020); we discuss this in detail in Section E.1 and also produce similar results as Köhler et al., (2020) for our model.

QM9: QM9 is a molecular dataset containing over 145,000 molecules used for moleccular property prediction. However, we use this for molecular structure generation of constitutional isomers of \chC5H8O1. Similar to DW-4, the molecules here are invariant to rotations, translations and, permutations of the same atoms. We encode these symmetries using E-GNN (Satorras et al.,, 2021), an equivariant graph neural network, to represent the energy. We trained our model via E-SVGD using \chC5H8O1 molecules present in the QM9 dataset and used the trained energy model to generate novel samples that are isomers of \chC5H8O1. We show these novel generated molecules in Figure 5 wherein we used the relative distance between atoms as a proxy for determining the covalent bonds. Our generated molecules demonstrate the correct 3D arrangement of bonds while containing complex atom structures like aromatic rings. This is further supported by the plots comparing the radial distribution functions of the two most common heavy atom pairs to quantify our model fit to QM9 (Figure 5). While, the generated molecules have a larger distributional spread, the range of values and modes – for both angles and distances– resemble the true distribution. We provide more details in Section E.2.

FashionMNIST: (Details in Section E.3) In this experiment, we take the FashionMNIST dataset with training set consisting regular images whereas the test set is processed to contain images that are randomly rotated using the $C_{4}$-symmetry group. We train an equivariant energy model where the energy function is a $C_{4}$ steerable CNN (Weiler and Cesa,, 2019) with both E-SVGD and vanilla SVGD. Furthermore, we also compare to an energy model with no rotation symmetries and depict the performance in terms of classification accuracy on the held out images of these three models as a function of the number of training iterations. The plot in Figure 6 shows, albeit unsurprisingly, that an equivariant energy model performs better than a regular model. Furthermore, the results also illustrate that an equivariant model trained with E-SVGD converges faster than when trained with vanilla SVGD highlighting the benefit of using E-SVGD for training equivariant EBMs. Furthermore, in Figure 6, we show samples generated by E-SVGD using the trained equivariant EBM from the joint and the class-conditional distribution.

In this paper, we focused on incorporating inductive bias in the form of symmetry transformations using equivariant functions for sampling and learning invariant densities. We first proposed equivariant Stein variational gradient descent algorithm for sampling from invariant densities by using equivariant kernels which affords many benefits in terms of efficiency due to its ability to model long-range interactions between particles. However, a major limitation of Stein variational gradient descent algorithm in general is its sensitivity to the kernel hyper-parameters. An interesting future work might be to develop strategies to either adapt or learn these hyper-parameters while running the SVGD dynamics.

Subsequently, we proposed equivariant energy based models wherein the energy function is parameterized by an equivariant network. In our experiments, we leveraged the recent advances in geometric deep learning to model EBMs using steerable CNNs (Weiler and Cesa,, 2019) for images, equivariant graph networks (Satorras et al.,, 2021) for representing molecules, and group equivariant networks (Cohen and Welling,, 2016) for many-body particle systems. We used equivariant SVGD to train these equivariant energy based models for modelling invariant densities and demonstrated that incorporating symmetries in the energy model as well as the sampler leads to efficient training. However, as discussed in previous works (Grathwohl et al.,, 2019), training EBMs using contrastive divergence and short sampling chains is often unstable and challenging. These issues remain with equivariant samplers and have to be addressed to be able to train large-scale energy based models.