Learning Equivariant Energy Based Models with Equivariant Stein Variational Gradient Descent

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 . 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 ) . 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 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 can not 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 .

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.  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.