(Kernelized) Stein variational gradient descent

1 Stein operators

We start with the classic Stein’s identity which turns out to be a useful trick for quantifying how well we have approximated some density.

We care about a target density $p$ and another density $q$ (which will end up approximating it). $p$ needs to be differentiable for this to work. They are both densities over, by assumption $\mathcal{X}\subseteq {\mathbb{R}}^d$. We also introduce a family of ${\mathbb{R}}^d$ to ${\mathbb{R}}^d$ test functions $\mathcal{F}$. We require that $\underset{x\to \pm \mathrm{\infty }}{lim}p(xb)f(xb)=0$ for $\parallel b\parallel =1$, and some other stuff which we get to in a moment.

Next, we choose a Stein operator ${\mathcal{A}}_x:\mathcal{F}\to \mathcal{G}$. $\mathcal{F}$ and $\mathcal{G}$ are spaces of functions from ${\mathbb{R}}^d$ to $\mathbb{R}$. I gave them different names because it is not clear to me that they are necessarily the same space, but we can probably ignore that detail for now. We do not go into the details of the requirements of the spaces, but they should be smooth functions that are square-integrable (with respect to the target distribution $p$? or Lebesgue measure? something else?) whose derivatives also go to zero at infinity. For an operator ${\mathcal{A}}_x$ to be a Stein operator for a target distribution $p(x)$, it must satisfy the following key property:

Stein’s Operators: ${\mathcal{A}}_x$ is a Stein operator with respect to a suitable class of test functions $\mathcal{F}$, and target distribution $p(x)$ if the expectation of ${\mathcal{A}}_{x}f(x)$ under $p$ is zero, i.e. for all those $f$ in $\mathcal{F}$,

$$\begin{array}{}\text{(1)}& {\mathbb{E}}_{X\sim p}[{\mathcal{A}}_{X}f(X)]=0.\end{array}$$

For $\mathcal{F}$ which include a non-trivial linear subspace, we can see that ${\mathcal{A}}_x$ must be linear, because expectation is linear, and otherwise we could make linear changes to an $f$ and end up violating the equality.

A popular choice is to set the Stein Operator to $$\begin{array}{}\text{(2)}& {\mathcal{A}}_{x}f(x):=f(x){\mathrm{\nabla }}_x\cdot \log p(x)+{\mathrm{\nabla }}_x\cdot f(x).\end{array}$$

Anastasiou et al. (2023) call this the Langevin Stein Operator, and it seems that if you do not otherwise specify, this is the one you get. The Langevin Stein operator makes this into a score-based method — See C. Liu et al. (2019) for some deep theory about that.

Example time! Let us make this more concrete, by choosing a specific $p$ which is not trivial but not baffling either. I reckon a 2d Gaussian with standard deviation 1 and mean 0 will do the trick. Let us give it a correlation $\rho$, which we leave unspecified, to keep things spicy. This implies mean $\mu =(0,0)$ and covariance $\mathrm{\Sigma }=\left[\begin{array}{cc}1& \rho \\ \rho & 1\end{array}\right]$, and thence inverse covariance ${\mathrm{\Sigma }}^{-1}=\frac{1}{1-{\rho }^2}\left[\begin{array}{cc}1& -\rho \\ -\rho & 1\end{array}\right]$. The pdf for this distribution is $$\begin{array}{rl}p(\mathit{x})& =\frac{1}{2\pi \sqrt{|\mathrm{\Sigma }|}}\exp (-\frac{1}{2}(x-\mu {)}^{\mathrm{\top }}{\mathrm{\Sigma }}^{-1}(x-\mu ))\\ & =\frac{1}{2\pi \sqrt{1-{\rho }^2}}\exp (-\frac{1}{2}\left[\begin{array}{cc}{x}_1& x_2\end{array}\right]\frac{1}{1-{\rho }^2}\left[\begin{array}{cc}1& -\rho \\ -\rho & 1\end{array}\right]\left[\begin{array}{c}{x}_1\\ x_2\end{array}\right])\\ & =\frac{1}{2\pi \sqrt{1-{\rho }^2}}\exp (-\frac{1}{2(1-{\rho }^2)}(x_1^2-2\rho x_1{x}_2+x_2^2)).\end{array}$$

We can simplify the Langevin Stein operator for this choice of $p$, since $$\begin{array}{rl}{\mathrm{\nabla }}_x\log p(x)& ={\mathrm{\nabla }}_x\log (\frac{1}{2\pi \sqrt{1-{\rho }^2}}\exp (-\frac{1}{2(1-{\rho }^2)}(x_1^2-2\rho x_1{x}_2+x_2^2)))\\ & ={\mathrm{\nabla }}_x(-\frac{1}{2(1-{\rho }^2)}(x_1^2-2\rho x_1{x}_2+x_2^2))\\ & =-\frac{1}{2(1-{\rho }^2)}{\mathrm{\nabla }}_x(x_1^2-2\rho x_1{x}_2+x_2^2)\\ & =\frac{1}{1-{\rho }^2}\left[\begin{array}{c}{x}_1-\rho x_2\\ x_2-\rho x_1\end{array}\right]\end{array}$$

Equation 1 then comes out to

$$\begin{array}{rl}{\mathbb{E}}_{X\sim p}[{\mathcal{A}}_{x}f(X)]& ={\mathbb{E}}_{X\sim p}[f(x){\mathrm{\nabla }}_x\cdot \log p(x)+{\mathrm{\nabla }}_x\cdot f(x)]\\ & ={\mathbb{E}}_{X\sim p}[\frac{f(X)}{1-{\rho }^2}(X_1-\rho X_2+X_2-\rho X_1)+{\partial }_{X_1}f(X)+{\partial }_{X_2}f(X)]\\ & ={\mathbb{E}}_{X\sim p}[\frac{f(X)(1-\rho )}{1-{\rho }^2}(X_1+X_2)+{\partial }_{X_1}f(X)+{\partial }_{X_2}f(X)]\\ & ={\mathbb{E}}_{X\sim p}[f(X)\frac{X_1+X_2}{1+\rho }+{\partial }_{X_1}f(X)+{\partial }_{X_2}f(X)]\end{array}$$

We can choose some simple $\mathcal{F}$ for the purposes of intuition building, e.g. the linear set $$\mathcal{F}:=\{f(x)=a{x}_1+b{x}_2+c;a,b,c\in \mathbb{R}\}$ — we would normally use something a bit more interesting. The expectation of the Stein operator for our bivariate Gaussian for this function class is then

$$\begin{array}{rl}{\mathbb{E}}_{X\sim p}[{\mathcal{A}}_{x}f(X)]& ={\mathbb{E}}_{X\sim p}[f(X)\frac{X_1+X_2}{1+\rho }+{\partial }_{X_1}f(X)+{\partial }_{X_2}f(X)]\\ & ={\mathbb{E}}_{X\sim p}[(a{X}_1+b{X}_2)\frac{X_1+X_2}{1+\rho }+{\partial }_{X_1}(a{X}_1+b{X}_2)+{\partial }_{X_2}(a{X}_1+b{X}_2)]\\ & ={\mathbb{E}}_{X\sim p}[(a{X}_1+b{X}_2)\frac{X_1+X_2}{1+\rho }+a+b]\\ & ={\mathbb{E}}_{X\sim p}[\frac{(X_1+X_2)(a{X}_1+b{X}_2)}{1+\rho }+a+b]\end{array}$$

Phew! OK that worked. I itch to plot these functions; I think there are two qualities of interest; the first is the function $f$ itself, and the second is the Stein operator applied to $f$ weighted by the density $p$.

import jax.numpy as jnp
import numpy as np
from jax import grad, vmap
import plotly.graph_objects as go
import plotly.io as pio

pio.templates.default = "none"

# Plot params
n = 61

# Define the scalar function f
def f(x, a, b, c):
    return a * x[..., 0] + b * x[..., 1] + c

# Define the log-density of the generic probability density function p
def log_p(x, rho):
    return -0.5 * (
        x[..., 0]**2 + x[..., 1]**2
         - 2 * rho * x[..., 0] * x[..., 1]
    ) / (
        1 - rho**2
    ) - jnp.log(
        2 * jnp.pi * jnp.sqrt(1 - rho**2)
    )

# Define the Stein operator applied to some f and p
def A_x_f(x, f, log_p):
    grad_log_p = vmap(grad(log_p))(x)
    grad_f = vmap(grad(f))(x)
    return grad_log_p.sum(axis=-1) * f(x) + grad_f.sum(axis=-1)

# Fix specific values for rho and the parameters of f
rho = 0.3
a, b, c = 0.4, 0.25, -0.5
x1min, x1max = -3, 3
x2min, x2max = -3, 3

f_specific = lambda x: f(x, a, b, c)
log_p_specific = lambda x: log_p(x, rho)

# Create a grid of points
x1, x2 = np.meshgrid(
    np.linspace(x1min, x1max , n, endpoint=True),
    np.linspace(x2min, x2max, n, endpoint=True)
)
x = np.stack([x1, x2], axis=-1).reshape(-1, 2)

# Compute the function f at each point in x
f_x = f_specific(x).reshape(x1.shape)
p_x = np.exp(log_p_specific(x)).reshape(x1.shape)

# Compute the Stein operator for f at each point in x
A_x_f_x = A_x_f(x, f_specific, log_p_specific).reshape(x1.shape)
p_A_x_f_x = A_x_f_x * p_x

# Determine the z range with a margin
z_min = np.min(p_A_x_f_x) - 0.1
z_max = np.max(p_A_x_f_x) + 0.1

# Create the 3D surface plot
fig = go.Figure()

# Add the surface plot for the Stein operator coloured by the density p_x
fig.add_trace(
    go.Surface(
        z=p_A_x_f_x,
        x=x1,
        y=x2,
        surfacecolor=p_x,
        colorscale='Viridis',
        showscale=False,  # Remove the colour bar
        opacity=0.9,  # slightly transparent
        name='<i>p A<sub>x</sub> f</i>'  # Add name for legend
    )
)

# Add the contour plot for f on the same axes, with a different colour scheme and semi-transparent
fig.add_trace(
    go.Surface(
        z=f_x,
        x=x1,
        y=x2,
        colorscale='Cividis',
        showscale=False,
        opacity=0.5,  # make this semi-transparent
        name='f',  # Add name for legend
        contours={
            "z": {
                "show": True,
                "start": np.min(f_x),
                "end": np.max(f_x),
                "size": (np.max(f_x) - np.min(f_x)) / 10,
                "color":"white",
            }
        }
    )
)

# Set the layout with an initial camera view closer to the z=0 plane
fig.update_layout(
    title='<i>p A<sub>x</sub> f</i>&nbsp;and&nbsp;<i>f</i>',
    scene=dict(
        xaxis=dict(title='x<sub>1</sub>'),
        yaxis=dict(title='x<sub>2</sub>'),
        zaxis=dict(
            # title='p A<sub>x</sub> f, f',
            range=[z_min, z_max]),
        camera=dict(
            eye=dict(x=1.25, y=-1.25, z=0.5)  # Lower down closer to the z=0 plane
        )
    ),
    width=800,
    height=800,
    font=dict(family="Alegreya, serif"),
    paper_bgcolor='rgba(0,0,0,0)',
    plot_bgcolor='rgba(0,0,0,0)',
    ## legends don’t work on 3d contours
    # showlegend=True,  # Show legend
    # legend=dict(
    #     x=0.02,  # Position the legend on the left
    #     y=0.98,
    #     bgcolor='rgba(255,255,255,0.7)',  # Semi-transparent background for better visibility
    #     bordercolor='Black',
    #     borderwidth=1
    # )
)
# Show the plot
fig.show()

Did that help us? Well, kinda. It is not really clear to me that I should trust that the second figure should actually integrate to 0. Did it?

p_A_x_f_x.sum().item()*(x1max-x1min)*(x2max-x2min)/(p_A_x_f_x.size)

0.01856179405239057

Hm, not convincingly exactly 0, but not so far off that we cannot persuade ourselves that it is simply a truncation problem.

2 Stein discrepancy

We make Equation 1 into a quantity that depends on two, potentially-different densities by taking the expectation over a different density $q$ than the one that generated the operator ${\mathcal{A}}_x,$ and seeing if that does something useful:

$$\begin{array}{}\text{(3)}& {\mathbb{E}}_{x\sim q}[{\mathcal{A}}_{x}f(X)]=0\end{array}$$

Spoiler: it turns out that this does do something useful.

$$\begin{array}{rl}{\mathbb{E}}_{x\sim q}[{\mathcal{A}}_{x}f(X)]& ={\mathbb{E}}_{x\sim q}[{\mathcal{A}}_{x}f(X)]-\stackrel{=0}{\overbrace{{\mathbb{E}}_{x\sim q}[{\mathcal{A}}_{q}f(X)]}}\\ & ={\mathbb{E}}_{x\sim q}[f(x)\cdot {\mathrm{\nabla }}_x\log p(x)+{\mathrm{\nabla }}_x\cdot f(x)\\ & -f(x)\cdot {\mathrm{\nabla }}_x\log p(x)-{\mathrm{\nabla }}_x\cdot f(x)]\\ & ={\mathbb{E}}_{x\sim q}[f(x)\cdot {\mathrm{\nabla }}_x\log p(x)-f(x)\cdot {\mathrm{\nabla }}_x\log q(x)]\\ & ={\mathbb{E}}_{x\sim q}[f(x){\delta }_{p,q}(x)]\end{array}$$ where ${\delta }_{p,q}(x):={\mathrm{\nabla }}_x\log p(x)-{\mathrm{\nabla }}_x\log q(x)$ is the difference in score function between $p$ and $q$.

By choosing a $f$ from some sufficiently rich $\mathcal{F}$ we can make this non-zero unless $p=q$ a.e., so this equation tells us something about how distinct are two densities $p$ and $q$, in this slightly weird but credible-seeming sense where we care about the difference in their score functions. i.e. this is some kind of score matching method.

This looks neat. How can we calculate it in practice? Obstacle: we have not specified $f$. We could fix some $f$ and use it to measure how different are $p$ and $q$ in some sense. Or we could choose some stochastic process which generates some random $f$s and estimate it over many $f$s, I guess? I assume that has been done.

The Stein Discrepancy takes a strong approach to controlling those $f$s: We control the supremum of that difference over all $f$ in some function class $\mathcal{F}$, so that we know that this difference $p$ and $q$ is not too bad for any $f$, since if we have found this Stein discrepancy, we have found how bad it is over the worst $f$:

$$\sqrt{S(q,p)}=\underset{f\in \mathcal{F}}{sup}|{\mathbb{E}}_{x\sim q}[\mathrm{trace}({\mathcal{A}}_{x}f(X))]|$$

Notice we snuck in a trace there as well to make it a scalar? This ended up being the most confusing thing for me; how many dimensions even is anything in this equation?

Let us consider how we might find this ‘worst’ $f$ which gives us this most powerful guarantee of the difference between $p$ and $q$. There are a few steps.

First, we use the linearity of that Stein operator ${\mathcal{A}}_x$, mentioned earlier. Suppose that $f$ can be represented as a finite linear combination $f(x)=\sum _i{w}_i{f}_i(x)$ of a set of basis functions $f_i(x)$ for some coefficients $w_i$ s.t. $\parallel w\parallel \le 1$. Then we can define the ‘violation of Stein-ness’ by $${\mathbb{E}}_q\left[{\mathcal{A}}_{x}f\right]={\mathbb{E}}_q[{\mathcal{A}}_x\sum _i{w}_i{f}_i(x)]=\sum _i{w}_i{\beta }_i,$$ where $${\beta }_i={\mathbb{E}}_{x\sim q}[{\mathcal{A}}_x{f}_i(x)].$$

This only works for univariate densities, so far. To make the discrepancy be a scalar even for multivariate problems (in the sense of densities over multidimensional spaces) we define the violation as $${\mathbb{E}}_{X\sim p}\left[\mathrm{trace}({\mathcal{A}}_q\mathit{f}(x))\right]={\mathbb{E}}_{X\sim p}[{({\mathit{s}}_q(x)-{\mathit{s}}_p(x))}^{\mathrm{\top }}\mathit{f}(x)]$$

The optimal (i.e. greatest _mis_match) coefficients $w_i$ would then b $$\underset{w}{max}\sum _i{w}_i{\beta }_i,\text{ s.t. }\parallel w\parallel \le 1$$

OK, so this is notionally an optimisation problem we can solve, choosing the $w_i$ values to be as terrible as possible, and then seeing how bad the most-terrible values are.

The distances arising from these are apparently integral probability metric according to (Anastasiou et al. 2023).

However, it looks like a nested optimisation problem, which can be tedious. Can we do better?

3 Kernelized Stein Discrepancy

When we see a challenge of this kind — where we wish we could use ‘more tricks’ in our function space — it typically suggests that the trick we are looking for might be the kernel trick. This entail choosing the tricky function class $\mathcal{F}$ to be a reproducing kernel Hilbert space (“RKHS”) and seeing what that does to the problem, which we call i kernelising. Frequently that makes life easier. Spoiler: it helps here too.

So, how kernelized Stein discrepancy works is as follows: $\mathcal{H}$ is the RKHS with associated kernel $k:{\mathbb{R}}^d\times {\mathbb{R}}^d\to \mathbb{R}$. We require that $k(x,x^{\prime }):{\mathbb{R}}^d\times {\mathbb{R}}^d\to \mathbb{R}$ be positive definite kernel. The RKHS $\mathcal{H}$ with kernel $k$ includes functions of form $f(x)=\sum _i{w}_{i}k(x,x_i)$, equipped with RKHS inner product $⟨f,g{⟩}_{\mathcal{H}}=\sum _{ij}{w}_i{v}_{j}k(x_i,x_j)$ for $g=\sum _j{v}_{j}k(x,x_j)$ and RKHS norm $\parallel f{\parallel }_{\mathcal{H}}^2=\sum _{ij}{w}_i{w}_{j}k(x_i,x_j)$.

Now we have some extra structure: $$\begin{array}{rlrl}f(x)& =⟨f(\cdot ),k(x,\cdot ){⟩}_{\mathcal{H}}& & \text{reproducing property}\\ {\mathrm{\nabla }}_{x}f(x)& ={⟨f(\cdot ),{\mathrm{\nabla }}_{x}k(x,\cdot )⟩}_{\mathcal{H}}& & \text{gradient property}\end{array}$$

This is one of two kernels that we invoke; in fact this one gets “Steinalized” using our operator ${\mathcal{A}}_x$ to give us another kernel:

$$\begin{array}{rl}{k}_{\mathcal{A}}(x,x^{\prime })& :=\mathrm{trace}[{\mathcal{A}}_x{\mathcal{A}}_{x^{\prime }}k(x,x^{\prime })].\end{array}$$

If we plug in Equation 2, we get a special form,

$$\begin{array}{rl}{k}_{\mathcal{A}}(x,x^{\prime })=& {\mathrm{\nabla }}_x\cdot {\mathrm{\nabla }}_{x^{\prime }}k(x,x^{\prime })\\ & +{\mathrm{\nabla }}_{x}k(x,x^{\prime })\cdot {\mathrm{\nabla }}_{x^{\prime }}\log p(x^{\prime })\\ & +{\mathrm{\nabla }}_{x^{\prime }}k(x,x^{\prime })\cdot {\mathrm{\nabla }}_x\log p(x)\\ & +k(x,x^{\prime })({\mathrm{\nabla }}_x\log p(x))\cdot ({\mathrm{\nabla }}_{x^{\prime }}\log p(x^{\prime }))\end{array}$$ Woof! look at those score functions everywhere!

Moreover, $$\begin{array}{rl}{\mathbb{E}}_{x\sim q}[\mathrm{trace}{\mathcal{A}}_x\mathit{f}(x)]& =\sum _{i=1}^d{⟨f_i(\cdot ),{\mathbb{E}}_{x\sim q}[{\mathcal{A}}_x{k}_i(\cdot ,x)]⟩}_{\mathcal{H}}\\ & ={⟨\mathit{f}(\cdot ),{\mathbb{E}}_{x\sim q}[{\mathcal{A}}_{x}k(\cdot ,x)]⟩}_{{\mathcal{H}}^d}\end{array}$$

We take $\mathcal{F}$ to be the unit ball in that RKHS, i.e. $\mathcal{F}:=\{\mathit{f};\parallel \mathit{f}{\parallel }_{{\mathcal{H}}^d}\le 1\}$.

Then we can write

$$\sqrt{S(q,p)}=\underset{f\in \mathcal{H},\parallel f{\parallel }_{{\mathcal{H}}^d}\le 1}{sup}\left\{{\mathbb{E}}_{x\sim q}[\mathrm{trace}{\mathcal{A}}_{x}f(x)]\right\}.$$ i.e. it is just the same, but we have restricted the function class to be an RKHS.

Define $${\beta }_{q,p}(\cdot )={\mathbb{E}}_{x^{\prime }\sim q}{\mathcal{A}}_{x}k(\cdot ,x^{\prime }).$$

Finding that supremum is then equivalent to solving $$\underset{f}{sup}{⟨f,{\beta }_{q,p}⟩}_{\mathcal{H}},\text{ s.t. }\parallel f{\parallel }_{\mathcal{H}}\le 1.$$

From this we get $\varphi (x)={\varphi }_{q,p}^{\ast }(x)/{\Vert {\varphi }_{q,p}^{\ast }\Vert }_{{\mathcal{H}}^d}$, where $${\varphi }_{q,p}^{\ast }(\cdot )={\mathbb{E}}_{x\sim q}[{\mathcal{A}}_{x}k(x,\cdot )],\text{ for which we have }\mathbb{S}(q,p)={\Vert {\varphi }_{q,p}^{\ast }\Vert }_{{\mathcal{H}}^d}^2$$

We maximize this, I assert, if we set $f={\beta }_{q,p}/{\Vert {\beta }_{q,p}\Vert }_{\mathcal{H}}$, normalising it to be on the unit ball (question: why can it not be on the interior?) at the point that maximises the expectation. Thus $$\begin{array}{rl}S(q,p)& ={\Vert {\beta }_{q,p}\Vert }_{{\mathcal{H}}^d}^2\\ & ={\mathbb{E}}_{x,x^{\prime }\sim q}\left[{\kappa }_p(x,x^{\prime })\right]\end{array}$$ where $${\kappa }_p(x,x^{\prime }):={\mathcal{A}}_x^x{\mathcal{A}}_x^{x^{\prime }}k(x,x^{\prime }).$$ Here we defined ${\mathcal{A}}_x^x$ and ${\mathcal{A}}_x^{x^{\prime }}$ represents the Stein operator w.r.t. variable $x$ and $x^{\prime }$, respectively. ${\kappa }_p(x,x^{\prime })$ is the “Steinalized” kernel obtained by applying Stein operator on $k(x,x^{\prime })$ twice.

$$S(p,q)={\mathbb{E}}_{x,x^{\prime }\sim p}[{\mathit{\delta }}_{q,p}(x{)}^{\mathrm{\top }}k(x,x^{\prime }){\mathit{\delta }}_{q,p}\left(x^{\prime }\right)],$$ where ${\mathit{\delta }}_{q,p}(x)=s_q(x)-s_p(x)$ is the score difference between $p$ and $q$, and $x,x^{\prime }$ are i.i.d. draws from $p(x)$.

It is a mess to write out in full though.

4 Stein Variational Gradient Descent

The next bit comes from Q. Liu and Wang (2019). It turns out that we can use this Stein trick to sample from some interesting distributions, by using the Stein discrepancy as a loss function. Interestingly, this works on posterior distributions in particular.

We manufacture an empirical $q$ by using a set of particles $\{x_i{\}}_{i=1}^n$.

The gradient descent here is not SGD where we assimilate gradient steps by looking at examples; it is rather a gradient descent in parameter space which converges towards a good approximation of the posterior.

A worked example will sort this out.

import jax
import jax.numpy as jnp
from jax import grad, jit, vmap
from jax.scipy.stats import norm
import plotly.graph_objects as go
import plotly.io as pio

pio.templates.default = "none"

# Define the target distribution (standard normal in this case)
def log_p(x, rho):
    return -0.5 * (
        x[..., 0]**2 + x[..., 1]**2
         - 2 * rho * x[..., 0] * x[..., 1]
    ) / (
        1 - rho**2
    ) - jnp.log(
        2 * jnp.pi * jnp.sqrt(1 - rho**2)
    )
# Define the RBF kernel
def rbf_kernel(x, y, h):
    return jnp.exp(-jnp.sum((x - y)**2) / (2 * h**2))

# Compute the Stein kernel
@jit
def stein_kernel(x, x_prime, h, rho):
    k = rbf_kernel(x, x_prime, h)
    grad_k_x = grad(rbf_kernel, argnums=0)(x, x_prime, h)
    grad_k_x_prime = grad(rbf_kernel, argnums=0)(x_prime, x, h)

    grad_log_p_x = grad(log_p)(x, rho)
    grad_log_p_x_prime = grad(log_p)(x_prime, rho)

    term1 = jnp.dot(grad_k_x, grad_k_x_prime)
    term2 = jnp.dot(grad_k_x, grad_log_p_x_prime)
    term3 = jnp.dot(grad_k_x_prime, grad_log_p_x)
    term4 = k * jnp.dot(grad_log_p_x, grad_log_p_x_prime)

    return term1 + term2 + term3 + term4

# Define the SVGD update
@jit
def svgd_update(particles, h, rho, lr=0.001):
    n_particles = particles.shape[0]
    updates = jnp.zeros_like(particles)

    def particle_update(i, updates):
        x_i = particles[i]
        grad_log_p_i = grad(log_p)(x_i, rho)

        def kernel_and_grad(j):
            x_j = particles[j]
            k_stein = stein_kernel(x_i, x_j, h, rho)
            return k_stein

        k_stein_values = vmap(kernel_and_grad)(jnp.arange(n_particles))

        phi = jnp.mean(k_stein_values, axis=0)
        updates = updates.at[i].set(phi)
        return updates

    updates = jax.lax.fori_loop(0, n_particles, particle_update, updates)

    return particles + lr * updates

# Initialize particles
key = jax.random.PRNGKey(0)
n_particles = 20
particles = jax.random.normal(key, (n_particles, 2))  # 2D particles

# Set kernel bandwidth
h = 0.1
rho = 0.8  # Correlation parameter

# Set the number of iterations
n_iterations = 1000

# Store particle locations at different stages
initial_particles = particles.copy()
mid_particles = None
final_particles = None

# Run SVGD for a few iterations
for i in range(n_iterations):
    particles = svgd_update(particles, h, rho)

    # Detect NaN values
    if jnp.isnan(particles).any():
        num_nan_particles = jnp.isnan(particles).any(axis=1).sum()
        raise ValueError(f"Detected {num_nan_particles} NaN values at iteration {i + 1}. Diagnostic Info: particles shape {particles.shape}, step number {i + 1}")

    if i == n_iterations // 2:
        mid_particles = particles.copy()

final_particles = particles.copy()

# Create a mesh grid for plotting the density
# Create a mesh grid for plotting the density
x = np.linspace(-3, 3, 100)
y = np.linspace(-3, 3, 100)
X, Y = np.meshgrid(x, y)
XY = np.stack([X.ravel(), Y.ravel()], axis=-1)

# Compute the log-density for the mesh grid
rho = 0.5  # Example correlation value
Z = np.exp(jax.vmap(lambda xy: log_p(xy, rho))(XY).reshape(X.shape))

fig = go.Figure()

# Add the density heatmap
fig.add_trace(go.Contour(
    x=x,
    y=y,
    z=Z,
    colorscale='Viridis',
    opacity=0.5
))

# Add the scatter plot of initial particles
fig.add_trace(go.Scatter(
    x=initial_particles[:, 0],
    y=initial_particles[:, 1],
    mode='markers',
    marker=dict(size=5, color='red', opacity=0.8),
    name='Initial Particles'
))

# Add the scatter plot of middle particles
fig.add_trace(go.Scatter(
    x=mid_particles[:, 0],
    y=mid_particles[:, 1],
    mode='markers',
    marker=dict(size=5, color='green', opacity=0.8),
    name='Middle Particles'
))

# Add the scatter plot of final particles
fig.add_trace(go.Scatter(
    x=final_particles[:, 0],
    y=final_particles[:, 1],
    mode='markers',
    marker=dict(size=5, color='blue', opacity=0.8),
    name='Final Particles'
))

fig.update_layout(
    title='SVGD Particles at Different Stages with Target Density',
    xaxis_title='x1',
    yaxis_title='x2',
    font=dict(family="Alegreya, serif"),
    template=pio.templates.default,
    legend=dict(
        x=0,
        y=1,
        traceorder='normal',
        font=dict(size=12),
        bgcolor='rgba(255, 255, 255, 0.5)',
        bordercolor='Black',
        borderwidth=1
    )
)

fig.show()

5 For mixtures

Mixtures in general are helpful in variational inference (Ranganath, Tran, and Blei 2016). See ELBO-within-Stein (Rønning et al. 2021), Nonlinear Stein (D. Wang and Liu 2019), Stein Mixtures (Nalisnick and Smyth 2017)…

6 Stochastic variants

(Li et al. 2020; Zhang et al. 2020)

7 As moment matching

See Q. Liu and Wang (2018).

8 By message passing

Define a kernel over factors and now the Stein updates become local messages. Discovered simultaneously in 2018 by D. Wang, Zeng, and Liu (2018) and Zhuo et al. (2018).

To read: Zhou and Qiu (2023).

This turns Stein VGD into a particle message passing algorithm.

9 Incoming

• Qiang Liu did a lot of the groundwork and produced some helpful resources

  • Stein Variational Gradient Descent
  • UT Statistical Learning & AI Group
  • Stein’s Method for Practical Machine Learning

• The Stein Gradient | Sanyam Kapoor

10 References