Venue: NeurIPS 2021 Conference

Reviews:

final_score: rating: 7: Good paper, accept
review_content: 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. 
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. 

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

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

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

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

Minor comments:

- Line 30: typo "the energy function is equivariant neural network" -> "the energy function is a equivariant neural network"?

- Line 65: typo "which implies to the following important result" -> "which implies the following important result"

limitations_and_societal_impact: Limitations are discussed in context of the unstable behaviour of contrastive divergence.
review_title: null

final_score: rating: 7: Good paper, accept
review_content: 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.
main_review: Describe the strengths of the work:

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

Explain the limitations of this work:

The 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.
It 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?

Clarity:

Overall, the paper is written clearly. There are a few minor typos/fixes I suggest:
Introduce acronym Reproducing Kernel Hilbert Space (RKHS)
line 40 just use the acronym SVGD
line 65: implies to -> implies
line 171: use of an equivariant kernel

Reproducibility:

An anonymized github link is provided for code, and thorough details of data processing, architectures, and hyperparameters is given in the appendix. 

limitations_and_societal_impact: The authors have adequately addressed the limitations and potential negative societal impact of their work.
review_title: null

final_score: rating: 7: Good paper, accept
review_content: 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.
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. 

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

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

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

On 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:

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

(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?

**References**: 

[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.
limitations_and_societal_impact: The authors addressed the limitations and societal impact in Appendix F.
review_title: null

final_score: rating: 6: Marginally above the acceptance threshold
review_content: 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.
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.

# Pros
-  This paper provides a new SVGD using a $\mathcal{G}$-invariant kernel function and its superior performance is supported in various tasks. 

# Cons
-  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.
-  My major concern is the theoretical and numerical properties of $mathcal{G}-$invariant kernel is unclear. I will comment on it below.
- 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).

# Comments and Questions:
- 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.
In 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.



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



limitations_and_societal_impact: Yes, they are shown in the checklist.
review_title: null

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)