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.

References in corpus (16)

• Fashion-MNIST: a Novel Image Dataset for Benchmarking Machine Learning Algorithms
• MCMC using Hamiltonian dynamics
• Variational Inference with Normalizing Flows
• Geometric Deep Learning: Grids, Groups, Graphs, Geodesics, and Gauges
• Group Equivariant Convolutional Networks
• A Connection between Generative Adversarial Networks, Inverse Reinforcement Learning, and Energy-Based Models
• Exploiting Cyclic Symmetry in Convolutional Neural Networks
• E(n) Equivariant Graph Neural Networks
• General $E(2)$-Equivariant Steerable CNNs
• Learning Non-Convergent Non-Persistent Short-Run MCMC Toward Energy-Based Model
• NeuTra-lizing Bad Geometry in Hamiltonian Monte Carlo Using Neural Transport
• Maximum Entropy Generators for Energy-Based Models
• Message Passing Stein Variational Gradient Descent
• Equivariant Hamiltonian Flows
• On the geometry of Stein variational gradient descent
• Annealed Stein Variational Gradient Descent

Cited by in corpus (1)

• Explaining the Explainers in Graph Neural Networks: a Comparative Study