The equivalence between Stein variational gradient descent
and black-box variational inference

The goal of Bayesian inference is to compute the posterior $P(x|z)$ over a variable of interest $x$. In principle, this posterior may be computed from the prior $P(x)$ and the likelihood $P(z|x)$ of observing data $z$, using the equation

We denote the posterior as $p(x)$ for convenience of notation. Unfortunately, the integral in the denominator is usually intractable, which motivates variational inference techniques, which approximate the true posterior $p(x)$ with an approximate posterior $q(x)$, often by minimizing the KL divergence $\mathrm{KL}(q(x)\,\|\,p(x))$. In this paper, we consider two popular variational inference techniques, black-box variational inference (Ranganath et al., 2014) and Stein variational gradient descent (Liu & Wang, 2016), and show that they are equivalent when viewed as instances of kernel gradient flow.

Stein variational gradient descent (Liu & Wang, 2016), or SVGD, is a technique for Bayesian inference that approximates the true posterior $p(x)$ with a set of particles $x_{1},\ldots,x_{n}$.

In the continuous-time limit of small step size, each particle undergoes the update rule

where $q_{t}$ denotes the empirical distribution of particles at time $t$:

and $k(x,y)$ is a user-specified kernel function, such as the RBF kernel $k(x,y)=e^{-||x-y||^{2}}$.

In the mean-field limit as $n\to\infty$ (Lu et al., 2019), an equivalent form of the dynamics (2) is obtained by an application of Stein’s identity (integration by parts on the second term):

Black-box variational inference (Ranganath et al., 2014), or BBVI, is another technique for Bayesian inference that approximates the true posterior $p(x)$ with an approximate posterior $q_{\phi}(x)$, where $q_{\phi}$ is a family of distributions parameterized by $\phi$. In BBVI, we maximize the evidence lower bound, or ELBO, objective

by gradient ascent on $\phi$. This procedure effectively minimizes the KL divergence between $q_{\phi}(x)$ and the true posterior $p(x)=P(x|z)$, since the KL divergence and the ELBO objective differ by only the evidence $P(z)$, which is constant w.r.t. $\phi$:

Our claim is:

To see this, we observe that the evolution of the parameters $\phi$ under gradient ascent is governed by

Next, we specialize to the case where the family of approximate posteriors is parameterized via the reparameterization trick (Kingma & Welling, 2014). That is, suppose that there exists a fixed distribution $\omega$ and a parameterized function $f_{\phi}$ such that the following two sampling methods result in the same distribution over $x$:

As an example, the family of normal distributions $\mathcal{N}(\mu,\sigma)$ may be reparameterized as

In this setting, Roeder et al. (2017) and Feng et al. (2017) noted that

Now, we consider the dynamics of a sample $x=f_{\phi}(\varepsilon)$ under the parameter dynamics (7). By the chain rule, we have that

Let us introduce the neural tangent kernel of Jacot et al. (2018)

and define

making the additional assumption that $\varepsilon\mapsto f_{\phi}(\varepsilon)$ is injective. Note that if $x\in\mathbb{R}^{n}$, then $\Theta_{\phi}(\varepsilon,w)$ and $k_{\phi}(x,y)$ are both $n$-by-$n$ matrices that depend on $\phi$. Then, substituting (7) and (10) into (11), we find that the samples satisfy

Comparing (16) with the SVGD dynamics (4), we find an exact correspondence between SVGD and BBVI, where in BBVI, the kernel is given by (13) and defined by the neural tangent kernel.

In the previous section, we found that SVGD and BBVI both correspond to particle dynamics of the form

One peculiar feature of the BBVI dynamics is that the kernel $k_{\phi}$ depends on the current parameter $\phi$, rather than being constant as the approximate posterior $q_{\phi}$ changes, as in the SVGD case.

In fact, we argue that this feature of BBVI is quite natural:

To make this claim, let us first review Euclidean and Riemannian gradient flows. In Euclidean space, following the negative gradient of a function $J:\mathbb{R}^{n}\to\mathbb{R}$ according to

can lead to a minimizer of $J$. Analogously, on a Riemannian manifold $M$, following the negative Riemannian gradient of a function $J:M\to\mathbb{R}$ according to

can lead to a minimizer of $J$. Here, $G$ is a positive-definite matrix-valued function called the Riemannian metric, which defines the local geometry at $x$ and perturbs the Euclidean gradient $\nabla J$ pointwise. Note that in the case that $G(x)$ is the identity matrix for all $x$, Riemannian gradient flow reduces to the Euclidean gradient flow.

Next, we review Wasserstein gradient flows, which generalize gradient flows to the space of probability distributions (Ambrosio et al., 2008). Here, we consider the set of all probability distributions over a particular space formally as an “infinite-dimensional” manifold $\mathcal{P}$, and we consider a function $J:\mathcal{P}\to\mathbb{R}$. In variational inference, the most relevant such function is the KL divergence $J(q):=\mathrm{KL}(q\,\|\,p)$, where we are interested in finding an approximate posterior that minimizes $J$. Analogous to before, a minimizer of $J$ may be obtained by following the analogue of a gradient; the trajectory of the distribution $q$ turns out to take the form of the PDE

Here, $\nabla\Psi_{q}$ serves as the correct analogue of the gradient of $J$ evaluated at $q$, and it turns out that $\Psi_{q}(x)=\log q(x)-\log p(x)$ for the variational inference case $J(q):=\mathrm{KL}(q\,\|\,p)$. This function $\Psi_{q}$ is known variously as the functional derivative, first variation, or von Mises influence function.

Now, we review the recent perspective that SVGD can be interpreted as a generalized Wasserstein gradient flow under the Stein geometry (Liu, 2017; Liu et al., 2019; Duncan et al., 2019). We follow the presentation of Duncan et al. (2019) and refer to it for a rigorous treatment. To set the stage, we take a non-parametric view of the SVGD update (2), in which the dependence on $\phi$ is interpreted as dependence on the distribution $q$ itself:

Substituting $\Psi_{q}(y)=\log q(y)-\log p(y)$ and the linear operator $T_{q}$ defined by

we have

Under these dynamics, the probability distribution $q$ evolves according to the PDE

Comparing (30) with (26), we see that (30) defines a modified gradient flow in which the gradient $\nabla\Psi_{q}$ is perturbed by the operator $T_{q}$.

We now advocate for generalizing Wasserstein gradient flow in the same way that Riemannian gradient flow generalizes Euclidean gradient flow. The operator $T_{q}$ perturbs the gradient in a way analogous to how the Riemannian metric perturbs the Euclidean gradient in (25), so the operator $T_{q}$ thereby defines an analogue of a Riemannian metric on $\mathcal{P}$. However, there is no fundamental reason that $T_{q}$ must have the restrictive form prescribed by (28). Indeed, because $T_{q}$ is analogous to the Riemannian metric $G(x)$, it is natural to let the kernel, whose action defines the operator $T_{q}$, depend on the current value of $q$. It is also natural to allow the kernel to output a matrix rather than a scalar so that $T_{q}$ may mix all components of $\varphi$. Duncan et al. (2019) in fact speculate on these possibilities (Remarks 17 and 1).

With these considerations in mind, we propose replacing (28) with

where the kernel $k_{q}$ now depends on $q$ and outputs a matrix. This defines a gradient flow by (30) that we will refer to as kernel gradient flow.¹¹1To further the analogy between Euclidean and Riemannian gradient flow and Wasserstein and kernel gradient flow, note that just as setting the Riemannian metric to identity matrix for all $x$ reduces Riemannian to Euclidean gradient flow, setting $T_{q}$ to the identity operator for all $q$ reduces kernel to Wasserstein gradient flow. The special “Euclidean” Riemannian metric obtained this way is the central object of the Otto calculus (Otto, 2001; Ambrosio et al., 2008).

Once $T_{q}$ has the form (31), BBVI may naturally be regarded as an instance of kernel gradient flow, in which the kernel $k_{q}$ is the neural tangent kernel which depends on the current distribution $q$. More abstractly, we see that the neural tangent kernel defines a Riemannian metric on the space of probability distributions. We summarize the perspective that this framework gives on variational inference:

We now argue that the kernel gradient flow perspective we have developed describes not only SVGD and BBVI, but also describes the training dynamics of generative adversarial networks (Goodfellow et al., 2014).

Generative adversarial networks, or GANs, are a technique for learning a generator distribution $q_{\phi}$ that mimics an empirical data distribution $p_{\mathrm{data}}$. The generator distribution $q_{\phi}$ is defined implicitly as the distribution obtained by sampling from a fixed distribution $\omega$, often a standard normal, and running the sample through a neural network $f_{\phi}$ called the generator. The learning process is facilitated by another neural network $D_{\theta}$ called the discriminator that takes a sample and outputs a real number, and is trained to distinguish between a real sample from $p_{\mathrm{data}}$ and a fake sample from $q_{\phi}$. The generator and discriminator are trained simultaneously until the discriminator is unable to distinguish between real and fake samples, at which point the generator distribution $q_{\phi}$ hopefully mimics the data distribution $p_{\mathrm{data}}$.

For many GAN variants, the rule to update the generator parameters $\phi$ can be expressed in the continuous-time limit as

or by the chain rule,

The discriminator parameters $\theta$ are updated simultaneously to minimize a separate discriminator loss $L_{\text{disc}}(\theta,\phi)$, but it is common for theoretical purposes to assume that the discriminator achieves optimality at every training step. Denoting this optimal discriminator as $-\Psi_{\phi}$ (i.e. setting $\Psi_{\phi}:=-D_{\theta^{*}}$ for $\theta^{*}:=\operatorname*{arg\,min}_{\theta}L_{\text{disc}}(\theta,\phi)$), we have

This matches the BBVI update rule (10) with $\log q_{\phi}(y)-\log p(y)$ replaced by the discriminator $\Psi_{\phi}(y)$. Hence, analogous to (16), the generated points $x$ satisfy

where here $k_{\phi}$ is defined as in (13) by the neural tangent kernel of the generator $f_{\phi}$. Finally, it was observed that the optimal discriminator $\Psi_{q}$ of the minimax GAN equals the functional derivative of the Jensen–Shannon divergence (Chu et al., 2019, 2020); hence we conclude:

Similarly, non-saturating and Wasserstein GAN updates generate kernel gradient flows on the directed divergence $J(q):=\mathrm{KL}(\frac{1}{2}p_{\text{data}}+\frac{1}{2}q\,\|\,p_{\text{data}})$ and Wasserstein-1 distance $J(q):=W_{1}(p_{\text{data}},q)$ respectively.

We have cast SVGD and BBVI, as well as the dynamics of GANs, into the same theoretical framework of kernel gradient flow, thus identifying an area ripe for further study.