Learning Equivariant Energy Based Models with Equivariant Stein Variational Gradient Descent

Abstract:
We focus on the problem of efficient sampling and learning of probability densities by incorporating symmetries in probabilistic models. We first introduce Equivariant Stein Variational Gradient Descent algorithm – an equivariant sampling method based on Stein’s identity for sampling from densities with symmetries. Equivariant SVGD explicitly incorporates symmetry information in a density through equivariant kernels which makes the resultant sampler efficient both in terms of sample complexity and the quality of generated samples. Subsequently, we define equivariant energy based models to model invariant densities that are learned using contrastive divergence. By utilizing our equivariant SVGD for training equivariant EBMs, we propose new ways of improving and scaling up training of energy based models. We apply these equivariant energy models for modelling joint densities in regression and classification tasks for image datasets, many-body particle systems and molecular structure generation.

heading: 1 Introduction
text: 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.
*Equal contribution. 35th Conference on Neural Information Processing Systems (NeurIPS 2021).

heading: 2 Preliminaries and Setup
text: 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 G be a group acting on Rd through a representation R : G → GL(d) where GL(d) is the general linear group on Rd, such that ∀g ∈ G, g → Rg . Given a target random variable X ⊆ Rd with density π, we say that π is G-invariant if ∀g ∈ G and x ∈ Rd, π(Rgx) = π(x). Additionally, a function f(·) is G-equivariant if ∀g ∈ G and x ∈ Rd, f(Rgx) = Rgf(x). We denote with O(x) the orbit of an element x ∈ X defined as O(x) := {x′ : x′ = Rgx,∀g ∈ G}. We call π|G the factorized density of a G-invariant density π where π|G has support on the set X|G where the elements of XG are indexing the orbits i.e. if x, x̃ ∈ XG then x 6= Rgx̃,∀g ∈ 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:
(i) Equivariant Learning: Given access to an i.i.d. samples *x1, . . . ,xn+ ∼ π from a G-invariant density π, we want to approximate π. Rezende et al. (2019) and Köhler et al. (2020) addressed this by learning an equivariant normalizing flow (Tabak and Vanden-Eijnden, 2010; Tabak and Turner, 2013; Rezende and Mohamed, 2015) that transforms a simple latent G-invariant density q0 to the target density π through a series of G−equivariant diffeomorphic transformations T = (T1,T2, · · · ,Tk) i.e. π := T#q0. They achieved this by proving (cf. (Köhler et al., 2020, Theorem 1), (Rezende et al., 2019, Lemma1)) that if q0 is a G-invariant density in Rd, F is a proper sub-group of G i.e. F < G, and T is an F -equivariant diffeomorphic transformation, then π := T#q0 is F -invariant. However, a major drawback of this formalism is that it requires T to not only be a F -equivariant diffeomorphism, but computation of the inverse and Jacobian must be cheap as well. This is problematic in practice.
Köhler et al. (2020) overcame this issue by using continuous normalizing flows (Grathwohl et al., 2018) that define a dynamical system through a time-dependent Lipschitz velocity field Ψ : Rd × R+ → Rd with the following system of ordinary differential equations(ODEs):
dx(t)
dt = Ψ(x(t), t), x(0) = z (1)
This allows to define a bijective function TΨ,t(z) := x(0) + ∫ t
0 Ψ(x(t), t) dt which leads to a push-forward density qt at each time-step t satisfying d log qtdt = −div ( Ψ(x(t), t) ) , which implies to the following important result: Lemma 1 ((Köhler et al., 2020, Theorem 2)). Let Ψ be an F-equivariant vector-field on Rd. Then, the transformation TΨ,t(z) := x(0) + ∫ t 0
Ψ(x(t), t) dt is F-equivariant ∀t ∈ R+. Furthermore, the push-forward qt := TΨ,t,#q0 is F-invariant ∀t, if q0 is G-invariant and F < G.
Lemma 1 conveniently provides a framework to transform any G-invariant density to an F -invariant density along a path in which each intermediate density is also F-invariant. However, equivariant normalizing flows cannot be used directly to generate samples when given access to an invariant density π since they require i.i.d. samples from π to train the flow1.
(ii) Equivariant Sampling: In this paper, we are also interested in solving the inference problem i.e. we are interested in evaluating Eπ[f ], the expectation of f when given access to a G-invariant density π which typically involves generating samples *x1,x2, · · · ,xn+ ∼ π. Intuitively, sampling from a G-invariant density can be reduced to sampling from its corresponding factorized distribution π|G . This is because any set of samples {x̃i}ni=1 ∼ π|G can be used to get samples representing π by applying group actions from G to {x̃i}ni=1. Indeed, sampling methods like Markov Chain Monte Carlo (MCMC) (Brooks et al., 2011) or Hybrid Monte Carlo (HMC) Neal et al. (2011) and their variants, in principle, can use this paradigm to sample from an invariant density π. However, MCMC methods for approximate posterior sampling are often slow and it still remains challenging to scale them up to big data settings. An alternate to MCMC methods for approximate posterior sampling is Stein Variational Gradient Descent (SVGD) (Liu and Wang, 2016) which is a particle optimization
1In Appendix A, we discuss a way to use equivariant normalizing flow for direct sampling given access to π.
variational inference method that combines the paradigms of sampling and variational inference for Bayesian inference problems.
In SVGD, a set of n particles {xi}ni=1 ∈ X ⊆ Rd are evolved following a dynamical system to approximate the target (posterior) density π(x) ∝ exp ( − E(x) ) where E(·) is the energy function.
This is achieved in a series of T discrete steps that transform the set of particles {x0i }ni=1 ∼ q0(x) sampled from a base distribution q0 (e.g. Gaussian) at t = 0 using the map xt = T(x) := xt−1 + ε · Ψ(xt−1) where ε is the step size and Ψ(·) is a vector field. Ψ(·) is chosen such that it maximally decreases the KL divergence between the push-forward density qt(x) = T#qt−1(x) and the target π(x).
If Ψ is restricted to the unit ball of an RKHSHdk with positive definite kernel k : Rd ×Rd → R, the direction of steepest descent that maximizes the negative gradient of the KL divergence is given by:
Ψ∗q,π(x) := arg max Ψ∈Hdk
−∇εKL ( q||π ) |ε→0 = Ex∼q[trace(AπΨ(x))], (2)
where AπΨ(x) = ∇x log π(x)Ψ(x)> +∇xΨ(x) is the Stein operator. Thus, an iterative paradigm can be easily implemented wherein a set of particles {x01,x02, · · · ,x0n} ∼ q0 are transformed to approximate the target density π(·) using the optimal update Ψ∗q,π(x) ∝ Ex′∼q[Aπk(x′,x)]. Since AπΨ(x) = ∇x[π(x)Ψ(x)]/π(x) we have that Ex∼π[AπΨ(x)] = 0 for any Ψ implying convergence when q = π. Replacing the expectation in the update with a Monte Carlo sum over the current set of particles that represent q we get:
xt+1i ← x t i + εΨ̃ ∗(xti), where, Ψ̃ ∗(xti) :=
1
n n∑ j=1 ( ∇xtjk(x t j ,xi)︸ ︷︷ ︸
repulsive force
− k(xtj ,xi) · ∇xtjE(x t j)︸ ︷︷ ︸
attractive force
) (3)
Stein variational gradient descent intuitively encourages diversity among particles by exploring different modes in the target distribution π through a combination of the second term in Equation (3) which attracts the particles to high density regions using the score function and the repulsive force (first term) which ensures the particles do not collapse together. In the continuous time limit, as ε→ 0, Equation (3) results in a system of ordinary differential equations describing the evolution of particles {x01,x02, · · · ,x0n} according to dxdt = Ψ̃ ∗(x).
Furthermore, as shown in Wang et al. (2019), geometric information using pre-conditioning matrices can be incorporated in Equation (3) by using matrix valued kernels (cf. Definition 2.3 (Reisert and Burkhardt, 2007)) leading to the following generalized form of SVGD (Wang et al., 2019):
xt+1i ← x t i +
ε
n n∑ j=1 ( ∇xtjK(x t j ,xi)−K(xtj ,xi) · ∇xtjE(x t j) ) , (4)
where K(x,x′) is a matrix valued kernel. Matrix-valued SVGD allows to flexibly incorporate preconditioning matrices yielding acceleration in the exploration of the given probability landscape.
SVGD has gained a lot of attention over the past few years as a flexible and scalable alternative to MCMC methods for approximate Bayesian posterior sampling. Further, it is more particle efficient since it generates diverse particles due to the deterministic repulsive force induced by kernels instead of Monte Carlo randomness. A natural question to ask is: Can we incorporate symmetry information into SVGD for more efficient sampling from invariant densities? We answer this in the affirmative in the next section by proposing equivariant Stein variational gradient descent algorithm for sampling from invariant densities.

heading: 3 Equivariant Stein Variational Gradient Descent
text: 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). Proposition 1. Let π be a G-invariant density and *x01,x02, · · · ,x0n+ ∼ q0 be a set of particles at t = 0 with q0 being F-invariant where F > G. Then, the iterative update given by Equation (3) is G-equivariant and the density qt+1 defined by it at time t+ 1 is G-invariant if the positive definite kernel k(·, ·) is G-invariant. The same holds for Equation (4) ifK(·, ·) is G-equivariant.
Proof. Since the initial distribution q0 is F -invariant, following Lemma 1 the update in Equation (3) is G-equivariant if Ψ is G-equivariant. If k(·, ·) is G-invariant then ∇xk(·,x) is G-equivariant. Furthermore, since π = exp ( − E(x) ) is G-invariant, ∇xE(x) is also G-equivariant. Thus, both the terms for Ψ are G-equivariant if k(·, ·) is G-equivariant making the update in Equation (3) G-equivariant. The result follows similarly for Equation (4) whenK(·, ·) is G-equivariant.
Following Proposition 1, we call the updates in Equations (3) & (4) equivariant Stein variational gradient descent when the kernel k(·, ·) (andK(·, ·) respectively) is invariant (equivariant) and the initial set of particles *x01, · · · ,x0n+ are sampled from an invariant density q0. Thus, all that is required to sample from a G-invariant density π using equivariant SVGD is to construct a positive definite kernel that is G-equivariant. Let us next give a few examples for constructing in- and equivariant positive definite kernels. Example 1 (Invariant scalar kernel). Let G be a finite group acting onRd with representation R such that ∀g ∈ G, g → Rg . Then,
kG(x,x ′) = ∑ x∈O(x) ∑ x′∈O(x′) k(x,x′)
is G-invariant where k(·, ·) is some positive-definite kernel. While this provides a general method to construct invariant kernels for finite groups, the double summation can be computationally expensive. In practice, usually simple kernels like RBF kernel (for rotation symmetries) or uniform kernel suffice as more practical alternatives.
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 X|G as we show in the next example.
Example 2 (Continuous Symmetry Groups). Let π(x) be SO(2)-invariant (cf. Figure 3b for an example) where x ∈ R2 i.e. O(x) := {x′ : ‖x‖ = ‖x′‖}. We can either construct an invariant kernel for sampling from π using a Monte Carlo approximation by sampling random rotations on a unit sphere i.e.
kG(x,x ′) = n∑ i,j=1 k(gjx, gix ′), gi, gj ∈ G, ∀(i, j) ∈ [n]× [n]
Or alternately, we can consider the function ΦG : R2 → R such that ΦG(x) = ‖x‖. Then, ΦG(x) is SO(2) invariant since ΦG(gx) = ΦG(x),∀g ∈ G. Thus, we can now use the following kernel
kG(x,x ′) = k ( ΦG(x),ΦG(x ′) )
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)). Example 3 (Equivariant Matrix-Valued Kernels, Reisert and Burkhardt (2007)). Examples 1 and 2 define an invariant scalar kernel. Following Reisert and Burkhardt (2007), we can also construct a G-equivariant matrix-valued kernel for the generalized update as in Equation (4) by defining:
K(x,x′) = ∫ G k(x, gx′)Rg dg
where Rg is a group representation and k(·, ·) is a scalar symmetric, G-invariant function. It is easy to check thatK(x,x′) is equivariant in the first argument and anti-equivariant in the second argument, leading to an equivariantK(x,x′) (cf. Proposition 2.2 Reisert and Burkhardt (2007)).
Advantages of Equivariant Sampling: As we discussed briefly in Section 2, SVGD works by evolving a set of particles using a dynamical system through a combination of attractive and repulsive forces among the particles that are governed by the inter-particle distance. Thus, a particle exerts these forces in a restricted neighbourhood around it. Equivariant SVGD, on the other hand, is able to model long-range interactions among particles due to the use of equivariant kernel. Intuitively, any point x is able to exert these forces on any other point x′ in equivariant SVGD if x′ is in the
neighbourhood of any point in the orbit O(x) of x. This is because for any point x′ the repulsive and attractive force terms are the same in Equations (3) and (4) for all points that are in the orbit O(x). This ability to capture long-range interactions by equivariant Stein variational gradient descent subsequently makes it more efficient in sample complexity and running time with better sample quality, and makes it more robust to different initial configurations of the particles compared to vanilla SVGD. We illustrate these next with the help of the following examples:
0 1000 2000 3000 4000 5000
4.0
3.5
3.0
2.5
2.0
1.5
1.0
Lo g-
lik el
ih oo
d
(i) C4-Gaussians (cf. Figure 3a and 3c): This example consists of four Gaussians invariant to C4 symmetry group. In this case, the group factorized distribution π|C4 is Gaussian with the original C4-invariant density obtained by rotating π|C4 through the set {0◦, 90◦, 180◦, 270◦}. In Figure 3a, the first column shows the samples generated by equivariant SVGD, the second column is the projection of these samples on the group factorized space X|C4 and, the third column shows the samples obtained by rotating the original samples through the C4-symmetry group. Figure 3c shows a similar setup for vanilla SVGD.
(ii) Concentric Circles (cf. Figure 3b and 3d): This example comprises of two concentric circles invariant to the SO(2) symmetry group. In this case, the group factorized space is a union of two disconnected lines with length equal to the thickness of the circles. In Figure 3b, the first column shows the samples generated by equivariant SVGD and, the second column is the projection of these samples on the group factorized space X|SO(2). Figure 3d shows a similar setup for vanilla SVGD.
We keep the experimental setup i.e. number of particles and number iterations exactly the same for both vanilla SVGD and equivariant SVGD. For both the examples, it may seem that the original samples from the vanilla SVGD capture the target distribution better than the equivariant counterpart (first column for Figs. 3a-3d). However, projecting the samples onto the factorized space (second column for the aforementioned figures) shows that equivariant SVGD more faithfully captures the target density compared to vanilla SVGD. Furthermore, due to its ability to model long-range interactions, we see for both examples that in the projected space of the invariant sampler, the samples are not close together whereas for vanilla
SVGD most samples end up in a configuration where they reside in the same orbit. This phenomena is most evident for the concentric circles example where samples from vanilla SVGD reside on the high density region throughout the two circles resulting in all the samples being positioned on top of each other in the factorized space demonstrating its inability to capture the distribution. On the other hand, invariant SVGD prevents any sample from residing on the same orbit of another sample due to long-range repulsive force from the equivariant kernel allowing it to sample more faithfully from the invariant densities.
Secondly, we study the effect of increasing the number of particles used for vanilla SVGD for the two concentric circles example. In Figure 1, we plot the average log-likelihoods of the particles from vanilla SVGD and particles from invariant SVGD as a function of number of iterations and compare it to the ground-truth average log-likelihood. We run vanilla SVGD with up to 32 times more particles than invariant SVGD. As evident from the plot, invariant SVGD converges to the final configuration within the first 100 iterations with average log-likelihood closely matching the ground truth. Vanilla SVGD, on the other hand, is unable to converge to the ground truth with even 32 times more samples and 5000 iterations due to its inability to interact with particles at longer distances.
Finally, we study the effect of different configurations of the initial particles on the performance of vanilla and invariant SVGD in Figure 2 for the C4-Gaussian example. As shown by Zhuo et al. (2018); Zhang et al. (2020) and D’Angelo and Fortuin (2021), the particles in vanilla SVGD have a tendency
to collapse to a few local modes that are closest to the initial distribution of the particles. We test the robustness of invariant SVGD to particles with initial distributions localized to different regions in the space. We plot the average log-likelihoods of the converged samples for both invariant and vanilla SVGD for all random initializations in Figure 2 and compare this to the ground truth average log-likelihood. The plot illustrates that equivariant SVGD is more robust to the initial distribution of particles than vanilla SVGD. Nevertheless, if the group-factorized space is multi-modal, equivariant SVGD might exhibit a tendency to favour one of modes. However, this can be easily alleviated by either adding some noise to the SVGD update as proposed by Zhang et al. (2020) similar to SGLD (Welling and Teh, 2011) or using an annealing strategy (D’Angelo and Fortuin, 2021).

heading: 4 Equivariant Joint Energy Model
text: 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 π given access to i.i.d. samples *x1,x2, · · · ,xn+ ∼ π. Given a set of samples *x1,x2, · · · ,xn+ ⊆ Rd, energy-based models (LeCun et al., 2006) learn an energy function Eθ(x) : Rd → R that defines a probability distribution π̃θ(x) = exp ( −Eθ(x) ) /Zθ,
where Zθ = ∫ exp ( − Eθ(x) ) dx 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 π̃θ(·) since the energy function does not need to integrate to one, it can be parameterized using any nonlinear function. Conveniently, if π is G-invariant, we can use the existing equivariant deep network architectures to parameterize Eθ(·) 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:
θ∗ := arg min θ
LML(θ) = Ex∼π [ − log π̃θ(x) ] (5)
However, evaluating Zθ is intractable for most (useful) choices of Eθ(·) 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θ by
approximating the gradient of∇θLML(θ) in Equation (5) as follows: ∇θLML(θ) ≈ Ex+∼π [ ∇θEθ(x+) ] − Ex−∼π̃θ [ ∇θEθ(x−) ] (6)
Intuitively, the gradient in Equation (6) drives the model such that it assigns higher energy to the negative samples x− sampled from the current model and decreases the energy of the positive samples x+ which are the data-points from the target distribution. Since, training an EBM using MLE requires sampling from the current model π̃(θ), successful training of EBMs relies heavily on sampling strategies that lead to faster mixing. Fortuitously, since Eθ(·) in our present setting is G-equivariant, we propose to use our equivariant sampler for more efficient training2 of the equivariant energy based model.
Algorithm 1: Equivariant EBM training Input: *x+1 ,x+2 , · · · ,x+m+ ∼ π(x) while not converged do
. Generate samples from current eqNN model Eθ *x−1 ,x−2 , · · · ,x−m+ = EquivariantSVGD(Eθ) ; . Optimize objective LML(θ): ∆θ ← ∑m i=1∇θEθ(x + i )−∇θEθ(x − i ) ;
. Update θ using ∆θ and Adam optimizer end Additionally, following Grathwohl et al. (2019), we can extend equivariant energy based models to equivariant joint energy models. Let {(x1, y1), (x2, y2), · · · , (xn, yn)} ⊆ Rd × [K] be a set of samples with observations xi and labels yi. Given a parametric function fθ : Rd → Rk, a classifier uses the conditional distribution π̃θ(y|x) ∝ exp(fθ(x)[y]) where fθ(x)[y] is the logit corresponding to the yth class label. As shown by Grathwohl et al. (2019), these logits can be used to define the joint density π̃θ(x, y) and marginal density π̃θ(x) as follows:
π̃θ(x, y) = exp ( fθ(x)[y] ) Zθ , and, π̃θ(x) = ∑ y exp ( fθ(x)[y] ) Zθ
(7)
Hence, the energy function at a point x is given by Eθ = − log ∑ y exp(fθ(x)[y]) with joint energy Eθ(x, y) = −fθ(x)[y]. In our setting, the joint distribution π(x, y) is G-invariant in the first argument i.e. π(Rgx, y) = π(x, y),∀g ∈ 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θ to be G-equivariant to model a G-invariant density π̃θ(x, y). Furthermore, a G-equivariant fθ also makes the marginal density π̃θ(x) and conditional density π̃θ(y|x) G-invariant in the input x. We call such an energy model where fθ 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:
L(θ) : = LML(θ) + LSL(θ) = log π̃θ(x) + log π̃θ(y|x) (8)
where LSL(θ) 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 π̃θ(x) and evaluating the gradient of log π̃θ(y|x) through back-propagation. Conveniently, Equation (8) can also be used for semi-supervised learning with LSL((θ)) 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 C4-Gaussians (Figure 4) as shown in row two of the leftmost column of Figure 4. The innermost C4-Gaussian defines the class conditional probability π(x|y = 0) (row 3) and the outer C4-Gaussian defines π(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 Appendix D.1.
2compared to using a regular sampler with no encoded symmetries.
5 Experiments
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 (DW4) 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 cor-
responding metastable states in the training set. Our setup is different from that of Köhler et al. (2020); we discuss this in detail in Appendix 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 C5H8O1. 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 C5H8O1 molecules present in the QM9 dataset and used the trained energy model to generate novel samples that are isomers of C5H8O1. 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 Appendix E.2.
FashionMNIST: (Details in Appendix 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 C4-symmetry group. We train an equivariant energy model where the energy function is aC4 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 ESVGD 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.

heading: 6 Discussion and Conclusion
text: 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.

Review summaries:
This paper worked on efficient sampling from distributions that have special structures regarding symmetries. Focusing on the group invariance of the target distribution, the authors developed a new SVGD algorithm, of which kernel function incorporates the group invariance. Various numerical experiments show the superior efficiency of the new SVGD compared to naive SVGD.

This paper studies the incorporation of symmetries (of densities) into probabilistic models for efficient sampling and learning. It first proposes a method dubbed equivariant Stein variational gradient descent (SVGD), which uses work on equivariant, matrix-valued kernels to construct a set of particles along an optimal gradient path in order to sample from and approximate a target distribution -- which due to the exploitation of symmetry is able to handle long-range interactions. Next, recent work in equivariant and geometric deep learning is leveraged to build equivariant energy based models -- where the energy function is modeled by an equivariant network. These equivariant EBMs are trained using contrastive divergence while also using equivariant SVGD for sampling. A set of experiments, on toy data, molecular structure generation, and a vision task demonstrate the efficacy and sample efficiency of the methods proposed, and thus the utility of incorporating symmetry. 

This paper aims to develop EBMs which respect certain symmetries in the underlying data distribution. This is accomplished by first developing an equivariant sampling method, equivariant SVGD, and then applying that to EBMs.  Section 2 provides theoretical background on equivariance and learning equivariant distributions.  Section 3 provides a thorough investigation comparing vanilla SVGD to equivariant SVGD. Figure 3 is immensely helpful in explaining the visual intuition of the method. It explains that because equivariant SVGD uses an equivariant kernel, it is better able to capture long-range interactions of particles. Particles can interact from a long distance if they are along the same orbit that SVGD will be invariant to. It is then demonstrated that this leads to far more efficiency in terms of sample complexity and log-likelihood.   Section 4 introduces equivariant EBMs. First, this requires using recent advances in geometric deep learning to develop high-capacity equivariant neural-network based energy functions. Since the energy is equivariant, they propose using the equivariant SVGD sampler to improve mixing of sampling, which would allow for more efficient learning from CD. They additionally propose training joint EBMs following Grathwohl et al. (2019). This formulation allows them to incorporate labels, and train the joint model by factoring the loss into a CD term and a CE term. This loss also allows for learning on unlabelled data, for example in semi-supervised learning. They demonstrate their approach on a simple 2-dimensional toy example, and show they model can generate better conditional samples and a better energy.   In section 5 they demonstrate their approach on a number of datasets.