Learning Equivariant Energy Based Models with Equivariant Stein Variational Gradient Descent
Authors: Priyank Jaini
Lars Holdijk
Max Welling
Venue: NeurIPS 2021 Poster (2021)
Keywords: Energy Based Models
Equivariance
Equivariant Sampling
Stein Variational Gradient Descent
Unsupervised Learning
Self-Supervised Learning

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.  

Learning Equivariant Energy Based Models with Equivariant Stein Variational Gradient Descent
Priyank Jaini" Bosch-Delds Lab University of Amsterdam
Lars Holdijk" Bosch-Delta Lab University of Amsterdam
Max Welling Bosch-Delta Lab University of Amsterdam
Abstract
We focus on the problem of efficient sampling and learning of probability densities by incorporating symmetries in probabilistic models. We first introduce Equivari- anf 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 density through equiv- ariant kernels which makes the resultant sampler efficient both in terms of sample complexity and the quality of generated samples. Subsequently, we define equiv- ariant energy based models to model invariant densities that are learned using contrastive divergence. By utilizing 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.
1 Introduction
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. 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 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 more sample efficient, produces 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 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 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-workh 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 (NcurIPS 2021).

---

Ts) 2 Preliminaries and Setup
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 be a group acting on Rd through representation R G + GL(d) where GL(d)) the general linear group on R such that Vg E G. g + Rg. Given a target random variable X C Rd with density T, we say that π is G-invariant if Vg E 9 and æ E R = xRR,z) = (xx). Additionally, function f(-) is G-equivariant if Vy 9 and x E R f(R = [(x)) We denote with (xx) the orbit of an element a: E defined as O(z) == (a': x = RR Vy E 9). We call aig the factorized density of G-invariant density π where T has support the set XIg where the elements of Xg are indexing the orbits i.c. if Φ, & E Xg then æ # Rz Vg E 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:
ii) Equivariant Learning: Given access to i.idd samples Iz...... In from G-invariant density x, 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-Eijndee 2010: Tabak and Turner, 2013; Rezende and Mohamed, 2015) that transforms a simple latent G-invariant density do to the target density through a series of G -eequivariant diffeomorphic transformations T (T T2.
i.e. π := T 9o- They achieved this by proving (cf. (Köhler et al., 2020, Theorem 1). (Rezende et al., 2019. Lemmal)) I)) that if qo is a G-invariant density in R F is proper sub-group of g i.e. F <<> G. and T is F-equivariant diffeomorphic transformation, then π == T # go 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 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 dynamical system through a time-dependent Lipschitz velocity field v Rd x R+ + Rd with the following system of ordinary differential equations(ODEss:
dar(() dt
(x)(),))) (0) =
This allows to define bijective function Tv, x: a0) + So V (x((), () di which leads push-forward density a. at each time-step t satisfying d oz = - (iv()( t)), which implies to the following important result:
de
Lemma 1 ((Köhler et al., 2020, Theorem 2)). Let V be F-equivariant vector-field Rd. Then, a0) 1 (xx(t), 1) di is F-equivariant Vl. E R4.. Furthermore, the push-forward == Tv.1 is F-imvariant Vt, if go is G-invariant and F < G.
the transformation Te, (=)
Lemma 1 conveniently provides a framework transform any G-invariant density 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 flow'.
(ii) Equivariant Sampling: In this paper. we are also interested in solving the inference problem i.c. we are interested in evaluating E,]f], the expectation of f when given access to a G-invariant density which typically involves generating samples [xx .... .n S ~ π. Intuitively. sampling from a G-invariant density can be reduced to sampling from its corresponding factorized distribution Tip- This is because any set of samples { ~ Tig can be used to get samples representing i by applying group actions from G to (ū: 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 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 MCMC methods for approximate posterior sampling Stein Variational Gradient Descent (SVGD) (Liu and Wang, 2016) which is a particle optimization
In Appendix A, we discuss way 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 {2; E X C RR are evolved following a dynamical system to approximate the target (posterior) density π (xx) x exp( -xx)) where E(-) is the energy function. This is achieved in a series of T discrete steps that transform the set of particles {- sampled from a base distribution go (e.g. Gaussian) at / using the map xf f==1 ~ ()(2))
{a0) in
T(x)
x'-1+ + e (xx=11 where e is the step size and V((.) is a vector field. vr()) is chosen such that maximally decreases the KI. divergence between the push-forwan density (x T # (æ) and the target (x).
If is restricted to the unit ball of an RKHS Hg with positive definite kernel k : R R + R. the direction of steepest descent that maximizes the negative gradient of the KL divergence is given by:
V 9,
(x) == arg max - V. (L = Exnu (trace (x))),
(2)
MENS
where Ar (x:)= V (() + V (ax) is the Stein operator. Thus, an iterative paradigm
can be easily implemented wherein set of particles {æ 22.... x.. do are transformed to approximate the target density π (T)) using the optimal update V (a:) x 1Er-ggAA, m)]-
Since A. (x) = V (x)))
we have that E [A, V (x)) = O for any Vp implying convergence when q = æ. Replacing the expectation in the update with a Monte Carlo sum over the current set of particles that represent q get:
at + at + evi"" (xa), where, T (x/)
Vas (()) wi))
M(af wi) Vay E(x3))
(3)
Y
j=1 repulsive force
attractive force
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 8 0, Equation (3) results in system of ordinary differential equations describing the evolution of particles (x). .Z..... ,=.)) according to dæz de (x))
Furthermore, shown in Wang et al. (2019), geometric information using pre-condiiigning 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):
x ++ <- x
K (x)), ax:) K K(xlj. a.)) Vas E(x());
(4)
T1
where K(x, is 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 lot of attention over the past few years 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.
3 Equivariant Stein Variational Gradient Descent
We begin this section by presenting the main result of this section by introducing equivariant Stein variational gradient descent (E-SVGD) by utilizing Lemma and Equations (3) and (4).
Proposition Let n be a G-invariant density and [ { at...... 2 ~ qo be a set of particles at t= 0 with go being F-invariant where F > G. Then, the iterative update given by Equation (3) is G-equivariant and the density get1 defined by it at time 1 is G-invariant if the positive definite kernel k(.,.) is G-invariant. The same holds for Equation (4) if K (:y-) is G-equivariant.
3

---

Proof. Since the initial distribution is F F-invariant, following Lemma 1 the update in Equation (3)
G-equivariant if is G-equivariant. If k(.,.) .) is G-invariant then V,,((- x) is G-equivariant. Furthermore, since the terms for V/ are G-equivariant if .)) is G-equivariant making the update in Equation (3)
exp( E(z)) is G-invariant, (aE(() is also G-equivariant. Thus, both
g-equivariant. The result follows similarly for Equation (4) when K(,,)) is G-equivariant.
Following Proposition 1, we call the updates in Equations (3) & (4) equivariant Stein variational
gradient descent when the kernel ((;,)) .) (and ((;))) respectively) is invariant (equivariant) and the initial set of particles ?=p,... I are sampled from an invariant density qo- Thus, all that is required
to sample from a G-invariant density using equivariant SVGD is to construct a positive definite kernel that G-equivariant. Let us next give a few examples for constructing in- and equivariant positive definite kernels.
Example (Invariant scalar kernel), Let G be a finite group acting on Rd with representation R such that Vg G.g Rg. Then,
kg(2,
ΣE
k(x,
xEO(x) x'E0(()))
is G-invariant where k(.,-) -) is some positive-definit kernel. While this provides 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 by either using Monte Carlo approximation or using transformation that performs computations in the factorized space XIg as we show in the next example.
Example (Continuous Symmetry Groups). Let r(x)) be 50(2)-invarian) (ef. Figure 3b for an example) where x E R iee. ((x) == (æ': ||=|| = [|). We can either construct an invariant kernel for sampling from nt using Monte Carlo approximation by sampling random rotations on unit sphere i.e.
kg (x, x')
Σ k(g,o,
Ir.g E G, V(. E [m] x [m]
Or alternately, consider the function &g R² R such that (x(xx) |x||- Then, Φg(x)) is
S0(2) invariant since (g(gx) 2g Vgg G. Thus, we can now use the following kernel
kg( x²) = k (z)) (2)))
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-Value Kernels, Reisert and Burkhardt (2007)). Examples and define an invariant scalar kernel. Following Reisert and Burkhardt (2007), we can also construct G-equivariant matrix-valued kernel for the generalized update as Equation (4) by defining:
K (x, x²) )==== = k(x RR dg
where Rg is a group representation and k(-,-) is scalar symmetric, G-invariant function. It is easy to check that K (ax, ) is equivariant the first argument and anti-equivariant the second argument, leading to equivariant K K(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 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 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 ((x) of æ. This is because for any point x' the repulsive and attractive force terms are the same in Equations (3) and (4) for points that are in the orbit (ax). This ability to capture long-range interactions by equivariant Stein variational gradient descent subsequently makes it more efficient 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:
(i Ga-Gaussaans (cf. Figure 3a and 3c): This example consists of four Gaussians invariant to C4 symmetry group. In this case, the group fac- torized distribution < i is Gaussian with the original Ca-invariant density obtained by rotat- ing TiC, through the set {0', 90" 180", 270"). }. In Figure 3a, the first column shows the sam- ples generated by equivariant SVGD, the second column is the projection of these samples the group factorized space XIC, and, the third column shows the samples obtained by rotating the original samples through the Ca-symmetry group. Figure 3c shows similar setup for vanilla SVGD.
-16
-22
-222
1000
Stop
Figure 1: Sample efficiency
Circles (cf. Figure 3b and 3d): This example com- prises of two concentric circles invariant to the S0(2) symmetry group. In this case, the group factorized space is a union of two dis- connected lines with length equal to the thickness of the circles. In Figure 3b, the first column shows the samples generated by equivari- ant 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.
Concentric
-2.8
-2.
-3.0
Log-Likellhood ood
-31
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 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 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.
-3.2
-3.3
-3.4
33.5
Invoriant
Regular
Figure 2: Robustness
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 to the ground-truth average log-likelioood. We run vanilla SVGD with up 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 its inability 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 C₁ -auassss 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
5

---

(a) C2 Gaussians Invariant SVGD sampling
(b) Two Circles Invariant SVGD sampling
(c) C4 Gaussians Vanilla SVGD Sampling
(d) Two Circles: Vanilla SVGD Sampling
Figure 3: Recommended to view in color. 3a (Left to Right) Original Samples from E-SVGD, samples projected on to the croupffactorized space and, samples obtained after applying group actions to the original samples. Yellow, Green and, Blue samples represent original samples rotated by 90º, 180 and, 270 respectively. 3b (Left to Right) Original Samples from E-SVGD and, samples projected on the group-factorie factorize space. 3e-3d: Same as 3a-3b but for vanilla SVGD.
to collapse to few local modes that are closest 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- lig-liiclshoods 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- nauu-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 al. (2020) similar to SGLD (Welling and Teh, 2011) or using an annealing strategy (D' Angelo and Fortuin, 2021).
Equivariant Joint Energy Model
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 learn invariant (unnormalized) densities æ given access to i..dd samples [x), I2. 2nn π.
Given a set of samples [x,, D2, ........ Rnn Rd, energy-based models (LeCun et al., 2006) learn an energy function Eg (x: Rd + IR that defines a probability distribution 7ee (x) exp( (o where Zo J exp (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, π is G-invariant, we can use the existing equivariant deep network architectures to parameterize Ea(-) 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- likelihoood of the data under the given model:
0 = arg min LM (0) Eg-z log (x))
(5)
However, evaluating Zo is intractable for most (useful) choices of Eo(-) which makes learning EBMs via maximum likelihood estimation problematic. Contrastive divergence (Hinton al., 2006) provides paradigm to learn EBMs using maximum likelihood estimation without needing to compute Ze by

---

approximating the gradient of V ML in Equation (5) as follows:
VeLmL ~ Ex ~ V VoLa((**)
- Ex ~Fp (olig(
(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 a= which are the data-points from the target distribution. Since, training an EBM using MLE requires sampling from the current model 5 #0), successful training of EBMs relies heavily on sampling strategies that lead to faster mixing. Fortuitously, since (.g()) in our present setting is g-equivariant, we propose to use our equivariant sampler for more efficient training" of the equivariant energy based model.
Additionally, following Grathwohl et al. (2019), we can extend equiv- ariant energy based models to equiv- ariant joint energy models. Let
Algorithm 1: Equivariant EBM training
Input: [xx a . ~ (a)) while not converged do
Generate samples from current eqNN model Ee
{(x1, 3n); (I22, 3/2)," Rd x [K] be a set of samples with ob- servations æ; and labels yı- Given a parametric function fe : Rd + R², a classifier uses the conditional distribu- tion R x exp( ( (v)) where (a [y] is the logit corresponding the y" class label. As shown by Grathwohl et al. (2019), these logits can be used to define the joint density to (z y) and marginal density #o (x)) follows:
(2n. yn)} S
1 L. ..... , Im = EquivariantSVGD((e)) >> Optimizee objective (0):
AO + > i-l [o - V VoLig ( (x;;;;
Update using AO and Adam optimizer
end
tee (1 y) = exp(fo(x)lv))
and, #e(x) Ey exp(fo(x)lv))
Ze
Zo
Hence, the energy function at point is given by = Ey exp( (x) with joint energy Eo (az, .) = fo(==]y)) In our setting, the joint distribution x ((rr)) y) is G-invariant in the first argument i.e. π (R,, y) = T (xx y). Vg E 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 fe to be G-equivariant to model a G-invariant density No (2, v). Furthermore, a G-equivariant fe also makes the marginal density (r)) and conditional density no (y)()) G-invariant in the input x. We call such an energy model where fe is equivariant to symmetry transformation to be an equivariant joint energy model.
We can train this model by maximizing the log- ligeehkkl of the joint distribution as follows:
((0)= LML (0) + Lsu (0) log ( log (y(la)
(8)
where (0) 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 (o(y(z)) through back prop-propagttiinn Conveniently, Equation (8) can also be used for semi-supervised learning with (s)( substituted with the appropriate supervised loss e.g. MSE for regression.
Let us end this section with an empirical example for learning mixture of C4 Gaussians (Figure 4) as shown in row two of the leftmost column of Figure 4. The innermost Gaaasssii defines the class conditional probability π (æly = 0) (row 3) and the outer C4-Gaussian defines <<0 = 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 presented in Appendix D.I.
"compared to using regular sampler with no encoded symmetries.
7

---

Experiments
In this section, we present empirical analysis of equivariant EBMs and E-SVGD through periments to (i) reconstruct potential function describing many-body particle system (DW- 4) trained using limited number of meta-stable states, (ii) model 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 space constraints, details about all the experi- ments as well as detailed figures are relegated Appendix E.
DW-:: In this many-body particles system, 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, show that given access to only single example of each metastable state configuration, an equiv- ariant EBM trained with E-SVGD can recover the
other states with similar energy as those training set. In Figure 7, the first column shows the metastable states present in the training set. The second column 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 distinct from the training set but mimic the energies of the cor- responding metastable states in the training set. (2020); we discuss this in detail in Appendix E. (2020) for model.
Figure 4: C4- Ca-Gaussian mixture model. Row 1: Decision Boundary. Row 2: Samples and energy of joint distribution π (xx y). Row 3: Samples and energy of conditional distribution π ( 0). Row 4: Samples and energy of conditional distri- bution π<zyy = 1). Left: Target distribution. Mid- dle: Non-equivariant EBM trained with vanilla SVGD. Right: E-EBM trained with E-SVGD.
Our setup is different from that of Köhler al. and also produce similar results as Köhler et al.
QM9: QM9 a molecular dataset containing over 145,000 molecules used for moleccular property prediction. However, use this for molecular structure generation of constitutional isomers of C₂H&O,, Similar 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 CsHoO,, molecules present in the QM9 dataset and used the trained energy model to generate novel samples that are isomers of C.H.O,. We show these novel generated molecules in Figure 5 wherein we used the relative distance between atoms a proxy for determining the covalent bonds. Our generated
a)
026
0310
10
22
Graly oo
Anger
Figure 5: (a) Molecules sampled from a EBM parameterized by E-GNN trained using E-SVGD. (b) Distribution of distance and angle between atom pairs and triplets.

---

Joint distribution
Conditional distribution
Accuracy
a
20.
Figure 6: Left & Center: Samples generated from joint and class-conditional distribution using equivariant EBM. Right: Plot of classification aceuracy vs. training iterations for equivariant and regular EBMs trained using vanilla SVGD and E-SVGD.
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 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.
(Details in Appendix E.3) In this experiment, 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 --ssymmeryy group. We train an
MNShionMMIST::
equivariant energy model where the energy function is Ca steerable
CNN (Weiler and Cesa, 2019) with both E-SVGD and vanilla SVGD.
Furthermore, we also compare to energy model with rotation symmetries and depict the performance in terms of classification accuracy the held out images of these three models function of the number of training iterations. The plot in Figure 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 condicondiii distribution.
Discussion and Conclusion
Figure 7: Col. Sam- ples from true potential en- ergy. Col. 2: Samples from EBM trained with SVGD. Col. 3: Samples from equivariant EBM trained with E-SVGD.
In this paper, we focused on incorporating inductive bias in the form of symmetry transformations using equivariant functions for sam- pling and learning invariant densities. We first proposed equivariant Stein variational gradient descent algorithm for sampling from in- variant 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 typer-parameters. An interesting future work might be develop strategies to either adapt or learn these syper-parameters while running the SVGD dynamics.
Subsequently, proposed equivariant energy based models wherein the energy function is pa- rameterized by an equivariant network. In our experiments, we leveraged the recent advances in geometric deep learning 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.

---

References
Bronstein, M. M., Bruna, J., Cohen, T., and Veličković, P. (2021). Geometric deep learning: Grids, groups, graphs, geodesics, and gauges. arXiv preprint arXiv:2104.13478.
Brooks, S., Gelman, A., Jones, G., and Meng, X.-I. (2011). Handbook of Markov Chain Monte Carlo. CRC press.
Cohen, T. and Welling, M. (2016). Group equivariant convolutional networks. In International conference on machine learning, pages 2990-2999, PMLR.
D' Angeelo F. and Fortuin, (2021). Annealed Stein Variational Gradient Descent. arXiv preprint arXiv:2101.09815
Dieleman, S., De Fauw, J., and Kavukcuoglu, K. (2016). Exploiting cyclic symmetry in convolutional neural networks. International conference on machine learning, pages 1889-1898. PMLR.
Du, Y. and Mordatch, I. (2019). Implicit generation and modeling with energy based models. NeurlPS.
Duncan, A., Nüsken, N., and Szpruch, L. (2019). On the geometry of Stein variational gradient descent. arXiv preprint arXiv: 191.00999.
Finn, C., Christiano, P., Abbeel, P., and Levine, S. (2016). A connection between generative adversarial networks, inverse reinforcement learning, and energy-based models. arXiv preprint arXiv:1611.0333
Grathwohl, W., Chen, R. T., Bettencourt, J., Sutskever, I,, and Duvenaud, D. (2018). Ffjord: Free-form continuous dynamics for scalable reversible generative models. arXiv preprint /7X1v:/8/0.01367.
Grathwohl, W., Wang, K.-C., Jacobsen, J.-H., Duvenaud, D., Norouzi, M., and Swersky, K. (2019). Your classifier is secretly energy based model and you should treat it like one. arXiv preprint arXiv:1912.03263.
Han, J., Ding. F., Liu, X., Torresani, L.. Peng. J., and Liu, Q. (2020). Stein Variational Inference for Discrete Distributions. In International Conference on Artificial Intelligence and Statistics, pages 4572-4572. PMLR.
Hinton, G., Osindero, S., Welling, M., and Teh, Y.W. (2006). Unsupervised discovery of nonlinear structure using contrastive backpropagation. Cognitive science, 30(4):725-731
Hoffmann, M., Sountsov, P., Dillon, J. V., Langmore, I., Tran, D., and Vasudevan, S. (2019). Neutraliz- ing Bad Geometry in Hamiltonian Monte Carlo using Neural Transport. ICML, 2019.
Jaini, Nielsen, D., and Welling, M. (2021). Sampling in Combinatorial Spaces with SurVAE Flow Augmented MCMC. In International Conference on Artificial Intelligence and Statistics, pages 3349-3357. PMLR.
Kingma, and Ba, J. (2014). Adam: A method for stochastic optimization. arXiv preprint 141:v:1412.6980.
Köhler, J., Klein, I.., and Noé, F. (2020). Equivariant flows: exact likelihood generative learning for symmetric densities. In International Conference on Machine Learning, pages 5361-5370. PMLR.
Kumar, R., Ozair, S., Goyal, A., Courville, A., and Bengio. Y. (2019). Maximum entropy generators for energy- models. arXiv preprint 101.0.990..8808..
LeCun, Y., Chopra, S., Hadsell, R., Ranzato, M., and Huang, F. J. (2006). A Tutorial on Energy Based Models. In MIT Press.
Liu, Q and Wang, D. (2016). Stein variational gradient descent: A general purpose bayesian inference algorithm. arXiv preprint 100::04008777
Liu, Q and Wang, (2017). LLearnnng deep energy models: Contrastive divergence vs. amortized
MLE. arXiv preprint arXiv:1707.0079
10

---

Neal, R. M. et al. (2011). MCMC using Hamiltonian Dynamics. Handbook of markov chain monte carlo. 2(11):2.
Nielsen, D., Jaini, P., Hoogeboom, E., Winther, O., and Welling, M. (2020). SurVAE Flows:
Surjections to Bridge the Gap between VAEs and Flows, arXiv preprint arXiv:2007.02731.
Nijkamp, E.,, Hill, M., Zhu, S.-C., and Wu, Y. (2019). Learning non-convergent non-persistent short-run MCMC toward energy-based model. arXiv preprint arXiv: 1904.09700
Parno, M. and Marzouk, Y. (2018). Transport Map Accelerated Markov Chain Monte Carlo. SIAM/ASA Journal on Uncertainty Quantification, 6(2):645-682.
Reisert, M. and Burkhardt, (2007). Learning Equivariant Functions with Matrix Valued Kernels. Journal of Machine Learning Research, 8(3).
Rezende, D. J. and Mohamed, S. (2015). Variational inference with normalizing flows. In ICML.
Rezende, J. Racanière, S., Higgins, L,, and Toth, P. (2019). Equivariant Hamiltonian Flows. arXiv preprint 1rXiv:1909.1311
Satorras, V. G., Hoogeboom, E., and Welling, M. (2021). E(n) equivariant graph neural networks. arXiv preprint arXiv: 102.09444.
Schulz, H., Müller, A., and Behnke, S. (2010). Investigating convergence of restricted boltzmann machine learning. In NIPS 2010 Workshop on Deep Learning and Unsupervised Feature Learning. volume 1, pages 6-1.
Tabak, E. G. and Turner, C. V. (2013). A family of nonparametric density estimation algorithms. Communications on Pure and Applied Mathematics, 66(2):145-164)
Tabak, E. G. and Handen-Eijnden, E. (2010). Density estimation by dual ascent of the log-likelihood. Communications Mathematical Sciences, 8(1):217-233.
Tieleman, T. (2008). Training restricted Boltzmann machines using approximations to the likelihood gradient. In Proceedings of the 25th international conference on Machine learning ICML '08, pages 1064-1071, 1, Helsinki, Finland. ACM Press.
Wang, D., Tang, Z., Bajaj, C., and Liu, Q. (2019). Stein variational gradient descent with matrix- valued kernels. Advances neural information processing systems, 32:7834.
Weiler, M. and Cesa, G. (2019), General B(2)-Eqqivariant Steerable CNNs. In Conference on Neural Information Processing Systems (NeurIPS).
Welling, M. and Teh, Y. W. (2011). Bayesian learning via stochastic gradient Langevin dynamics. In Proceedings of the 28th international conference on machine learning (ICML-I1), pages 68 688-688 Citeseer.
Xiao, H., Rasul, K., and Vollgraf, R. (2017). MNNIS:: Novel Image Dataset for Bench- marking Machine Learning Algorithms. arXiv:1708.0774:
Xie, J., Lu, Y., Gao, R., Zhu, S.-C., and Wu, Y. N. (2018). Cooperative training of descriptor and generator networks. IEEE transactions on pattern analysis and machine intelligence, 42:27-45
Zhang, J., Zhang, R., Carin, L., and Chen, C. (2020). Stochastic particle-optimizztion sampling and the non-asymptotic convergence theory. In International Conference Artificial Intelligence and Statistics, pages 1877-1887. PMLR.
Zhuo, J., Liu, C., Shi. J., Zhu, J., Chen, N., and Zhang, B. (2018). Message passing Stein variational gradient descent. In International Conference on Machine Learning, pages 6018-6027. PMLR.


Reviews:
[{"review_id": "apffxtSrfP", "reviewer": "Reviewer_FRDu", "summary": "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.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Overall, I think the paper provides practical and useful extensions for SVGD, but some practically and theoretically important points are still missing to publish this work.\n\n# Pros\n-  This paper provides a new SVGD using a $\\mathcal{G}$-invariant kernel function and its superior performance is supported in various tasks. \n\n# Cons\n-  The discussions on previous works should be presented in the main paper, otherwise, it is hard to follow for a non-expert why this work is important.\n-  My major concern is the theoretical and numerical properties of $mathcal{G}-$invariant kernel is unclear. I will comment on it below.\n- Another concern is the hyperparameter tuning in the kernel function is not clear.  In vanilla SVGD, we can adaptively tune the bandwidth of the RBF kernel using the median trick (although the authors used the constant bandwidth in numerical experiments).\n\n# Comments and Questions:\n- I am not sure that using the $\\mathcal{G}-$invariant kernel provides a valid gradient flow and it can converge successfully. I think the geometry of the proposed dynamics is significantly changed from the vanilla SVGD due to the $\\mathcal{G}-$invariant kernel,  so authors should discuss it, for example, considering the following points.\nIn previous works for SVGDs, to ensure the convergence, we often assume the translation-invariant kernel function, $k(x,y)=k(x-y)$. (for example, On the geometry of Stein variational gradient descent (A. Duncan, et al)). When using RBF kernel, this is satisfied but when we extend RBF kernel to $\\mathcal{G}-$invariant kernel, it is not satisfied. Another common assumption is the eigenvalues of the kernel function (SVGD as a kernelized Wasserstein gradient flow of the chi-squared divergence, (S Chewi, et al)). Assumptions used in (A Non-Asymptotic Analysis for Stein Variational Gradient Descent (A Korba et al)) is that the norm of the derivative of kernel function is upper bounded,  ($|\\nabla k(x,y)|\\leq B$)  and the larger the constant $B$ is, the slower convergence we obtain. Interestingly, this is contradicted with the numerical experiment in Figure 1 on page 5 in the paper. I think the norm of the derivative of the  $\\mathcal{G}-$invariant kernel function based on the RBF kernel is larger than the original RBF kernel.\n\n\n\n- Authors should describe what kernel function is used in numerical experiments in the main paper although it was written in Appendix currently. Moreover, as the authors mentioned in Section 6, tuning-hyperparameters in kernel function is very important for SVGD. Thus, the authors should mention how the bandwidths are tuned in numerical experiments.\n- Related to the bandwidth tuning, the authors described that bandwidths were fixed constant during the optimization in Appendix E. I think previous works including the original SVGD paper uses the median trick to tune bandwidths. As far as I know, using the median trick to adaptively tune the bandwidth during the optimization seems very important numerically. I strongly recommend running experiments not using constant bandwidth but uses the median trick for the bandwidth tuning for SVGD. However, I am not sure how to use the median trick for $\\mathcal{G}-$invariant kernel, especially for continuous groups since it has the integral term and there are infinitely many \"virtual\" particles.\n \n\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes, they are shown in the checklist.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3 hours"}, "extra_text": {}}, {"review_id": "GNRIbAVuotS", "reviewer": "Reviewer_7bET", "summary": "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.\n\nSection 2 provides theoretical background on equivariance and learning equivariant distributions.\n\nSection 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. \n\nSection 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. \n\nIn section 5 they demonstrate their approach on a number of datasets.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Describe the strengths of the work:\n\nThis paper develops a thorough theoretical background on equivariant learning and different approaches for learning equivariant distributions. It clearly explains its relation to prior work and develops detailed theory for equivariant SVGD.\n\nExplain the limitations of this work:\n\nThe paper does not sufficiently disentangle the effects of using an equivariant EBM and sampling from it using equivariant SVGD. While there are experiments in the appendix that demonstrate this on toy 2-dimensional data, there aren’t experiments which attempt to investigate this question on larger scale datasets. For example, it would be interesting to see if training with E-SVGD improves the performance of the energy function at out-of-distribution detection or calibration.\nIt would also be interesting to compare the run-times of the approaches with SVGD or E-SVGD. does SVGD require more steps then E-SVGD for the same accuracy/sample-quality? How much longer is an E-SVGD step compared to an SVGD step?\n\nClarity:\n\nOverall, the paper is written clearly. There are a few minor typos/fixes I suggest:\nIntroduce acronym Reproducing Kernel Hilbert Space (RKHS)\nline 40 just use the acronym SVGD\nline 65: implies to -> implies\nline 171: use of an equivariant kernel\n\nReproducibility:\n\nAn anonymized github link is provided for code, and thorough details of data processing, architectures, and hyperparameters is given in the appendix. \n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors have adequately addressed the limitations and potential negative societal impact of their work.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "9sOLr8zhXae", "reviewer": "Reviewer_FYs8", "summary": "The paper proposes an equivariant version of the Stein Variational Gradient Descent (SVGD) Algorithm for sampling from densities with symmetries. Using this sampling algorithm, the authors develop equivariant versions of energy-based models for learning invariant densities.", "questions": "", "limitations": "", "rating": 7, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "**Quality**: I checked the mathematical formulations and proofs in the paper, and to the best of my knowledge they seem to be correct. \n\n**Clarity**: The paper is in general clear and very well-written. Section 2 on Preliminaries and Setup can benefit from having a few more bolded subsections that distinguish between different methods (e.g. original SVGD vs. generalized form). Section 4 on the Equivariant Joint Energy Model can also be made more clear in the same way by adding subsections (e.g. original energy-based model vs. joint energy model).\n\n**Originality**: The equivariant extension of SVGD methods is novel, while it also seems to be relatively straightforward. Specific examples of applying equivariant energy-based models have been considered in concurrent work [1], though this paper is original in the sense that (i) it works on general energy-based models than a specific one, and (ii) it applies an equivariant sampling method within the energy model.\n\n**Significance**: The idea of introducing equivariance in energy-based models is a novel one, and the authors did a good job in ensuring equivariance either in the sampling step or for the network. At the same time, the theory part of the paper is more or less a straightforward extension (i.e. adding equivariance everywhere) from my understanding, so I will be more interested in the experimental results given in the paper.\n\nOn a high level, I think the experimental results are comprehensive in showing the benefits of equivariant SVGD vs the original one for learning symmetries. However, I come from a normalizing flow background and I am not very familiar with the experimental behavior of SVGD as a sampling method in general. I wonder if the authors can clarify some of the following observations about experiments:\n\n(1) In Figure 3(d), is there a good explanation for why the samples from vanilla SVGD all lie nearly exactly on the same orbit? While I do believe that equivariant SVGD may have samples that are more widely distributed across the width of the circle, it is a bit surprising to see vanilla SVGD converging to everything on the same orbit with literally no width in distribution. \n\n(2) In Figure 1, is there a good explanation for why all the vanilla SVGD experiments converge consistently to something different from the ground truth?\n\n**References**: \n\n[1] Wu, J., Shen, T., Lan, H., Bian, Y., & Huang, J. (2021). SE (3)-Equivariant Energy-based Models for End-to-End Protein Folding. bioRxiv.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors addressed the limitations and societal impact in Appendix F.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "10"}, "extra_text": {}}, {"review_id": "0Q4vfzFKcBZ", "reviewer": "Reviewer_WvCB", "summary": "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. ", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This paper builds off some recent work (such as by Kohler) on incorporating symmetries into probabilistic models. More specifically it is concerned with the problem of efficient sampling and learning of equivariant probability densities for this purpose. Towards this end, the paper makes a number of contributions. First. an equivariant version of the Stein variational gradient descent procedure (SVGD) is proposed. SVGD leverages some ideas from matrix-valued and equivariant kernels to construct a set of particles along an optimal gradient path. SVGD is then extended to invariant densities. Next, using recent work on equivariant deep learning, equivariant energy based models are proposed. These are trained by using contrastive divergence, with samples generated using SVGD. On a diverse set of experiments, the paper attempts to show the benefit of incorporating symmetry into such probabilistic models. \n\nThe paper is divided into four sections. The first gives a self-contained introduction, introducing the background, some of the recent work in the area and states the problem. The next introduces and formally specifies an equivariant sampler, followed by equivariant energy based models. The last section provides evidence for the benefit afforded by the approach (both in terms of the equivariant sampling, and equivariant EBMs).\n\nThe work of Kohler is used as a starting point. The problem is that given access to a set of i.i.d samples from a G-invariant density \\pi, we want to approximate \\pi. Prior work has approached this problem by means of an equivariant normalizing flow (which has the usual setup for normalizing flows but with G-invariance thrown in -- an initial, latent, G-invariant density is subject to a series of G-equivariant diffeomorphic transformations, so as to convert it into the target density). Efficient implementation of such a normalizing flow requires careful construction of an underlying dynamical system. While Kohler's work is referred to as a starting point and inspiration, it is not sufficient for sampling (at least in the way stated in the paper). The paper then turns to the Stein Variational Gradient descent procedure for this purpose, and considers the incorporation of equivariance in it, much like in the work of Kohler. The usual formulation of SVGD results in a system of ODEs which describes the evolution of particles as mentioned earlier. Furthermore, prior knowledge can be incorporated in the formulation of SVGD easily by considering kernels that reflect some geometric properties of the problem. Equivariant SVGD is achieved by simply using group equivariant matrix valued kernels. One of the nice things about equivariant SVGD, as compared to plain SVGD is that due to the fact that any particle can exert a force on any point in the neighborhood of any point in its orbit, it can model some long-range interactions better. The advantage of equivariant SVGD is demonstrated using some simple and standard(-ish) toy tasks, such as C4 gaussians and concentric circles. It is also demonstrated that invariant SVGD converges much more rapidly and with much fewer samples. The paper reports that vanilla SVGD does not converge even with 32 times the number of samples and 50 times more iterations. Next the paper extends energy based models to incorporate equivariance. They are trained with the usual CD procedure, with samples generated using equivariant SVGD. \n\nThe experiments comprise of a toy dataset where the task is to reconstruct the potential function of a many-body system trained using a few meta-stable states. Another experiment focuses on generating molecular structures (isomers of a particular molecule). Recently proposed E(n) equivariant GNNs are used as the underlying equivariant network, while equivariant SVGD is used for generating samples. Such experiments are certainly to be taken with a pinch of salt, however, whatever their limitations (and scale), we see that the equivariant EBM is able to recover correct 3D arrangement of bonds. No baseline comparisons are made, which is a weakness. Last is a familiar task on fashion-MNIST. \n\nIn summary: The paper presents a complete pipeline for incorporating symmetries into certain probabilistic models. For each of the steps, existing machinery (such as equivariant kernels, equivariance in EBMs) is cleverly used. Overall, the paper is well written, and makes a strong contribution. All the technical details seem correct and well fleshed out. While the experiments have significant room for improvement, I believe they are sufficient for corroborating the story of the paper. \n\nMinor comments:\n\n- Line 30: typo \"the energy function is equivariant neural network\" -> \"the energy function is a equivariant neural network\"?\n\n- Line 65: typo \"which implies to the following important result\" -> \"which implies the following important result\"\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Limitations are discussed in context of the unstable behaviour of contrastive divergence. ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3 hours"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper proposes an equivariant version of the Stein variational gradient descent algorithm for sampling from densities with symmetries. Using this sampling algorithm, the authors developed equivariant versions of energy-based models for learning invariant densities. The reviewers all agree that this is an interesting paper with nice contributions. I recommend acceptance as poster. 
Decision: Accept (Poster)