Stein Variational Gradient Descent with Matrix-Valued Kernels

Abstract

Stein variational gradient descent (SVGD) is a particle-based inference algorithm that leverages gradient information for efficient approximate inference. In this work, we enhance SVGD by leveraging preconditioning matrices, such as the Hessian and Fisher information matrix, to incorporate geometric information into SVGD updates. We achieve this by presenting a generalization of SVGD that replaces the scalar-valued kernels in vanilla SVGD with more general matrix-valued kernels. This yields a significant extension of SVGD, and more importantly, allows us to flexibly incorporate various preconditioning matrices to accelerate the exploration in the probability landscape. Empirical results show that our method outperforms vanilla SVGD and a variety of baseline approaches over a range of real-world Bayesian inference tasks.

1. Introduction

Approximate inference of intractable distributions is a central task in probabilistic learning and statistics. An efficient approximation inference algorithm must perform both efficient optimization to explore the high probability regions of the distributions of interest, and reliable uncertainty quantification for evaluating the variation of the given distributions. Stein variational gradient descent (SVGD) (Liu & Wang, 2016) is a deterministic sampling algorithm that achieves both desiderata by optimizing the samples using a procedure similar to gradient-based optimization, while achieving reliable uncertainty estimation using an interacting repulsive mechanism. SVGD has been shown to provide a fast and flexible alternative to traditional methods such as Markov chain Monte Carlo (MCMC) (e.g., Neal et al., 2011; Hoffman & Gelman, 2014) and parametric variational inference (VI) (e.g., Wainwright et al., 2008; Blei et al., 2017) in various challenging applications (e.g., Pu et al., 2017; Wang & Liu, 2016; Kim et al., 2018; Haarnoja et al., 2017).

On the other hand, standard SVGD only uses the first order gradient information, and can not leverage the advantage of the second order methods, such as Newton’s method and natural gradient, to achieve better performance on challenging problems with complex loss landscapes or domains. Unfortunately, due to the special form of SVGD, it is not straightforward to derive second order extensions of SVGD by simply extending similar ideas from optimization. While this problem has been recently considered (e.g., Detommaso et al., 2018; Liu & Zhu, 2018; Chen et al., 2019), the presented solutions either require heuristic approximations (Detommaso et al., 2018), or lead to complex algorithmic procedures that are difficult to implement in practice (Liu & Zhu, 2018).

Our solution to this problem is through a key generalization of SVGD that replaces the original scalar-valued positive definite kernels in SVGD with a class of more general matrix-valued positive definite kernels. Our generalization includes all previous variants of SVGD (e.g., Wang et al., 2018; Han & Liu, 2018) as special cases. More significantly, it allows us to easily incorporate various structured preconditioning matrices into SVGD updates, including both Hessian and Fisher information matrices, as part of the generalized matrix-valued positive definite kernels. We develop theoretical results that shed insight on optimal design of the matrix-valued kernels, and also propose simple and fast practical procedures. We empirically evaluate both Newton and Fisher based extensions of SVGD on various practical benchmarks, including Bayesian neural regression and sentence classification, on which our methods show significant improvement over vanilla SVGD and other baseline approaches.

Notation and Preliminary

For notation, we use bold lower-case letters (e.g., x) for vectors in ${ℝ}^d$, and bold upper-case letters (e.g., Q) for matrices. A symmetric function k: ${ℝ}^d$ × ${ℝ}^d$ → ℝ is called a positive definite kernel if ${\sum }_{ij}{c}_{i}k\left({\mathit{x}}_i,{\mathit{x}}_j\right)c_j\ge 0$ for any $\left\{c_i\right\}\subset ℝ$ and $\left\{{\mathit{x}}_i\right\}\subset {ℝ}^d$. Every positive definite kernel k(x, x^′) is associated with a reproducing kernel Hilbert space (RKHS) ${\mathcal{H}}_k$, which consists of the closure of functions of form

$$f(\mathit{x})=\sum _i{c}_{i}k\left(\mathit{x},{\mathit{x}}_i\right),\forall \left\{c_i\right\}\subset ℝ,\left\{{\mathit{x}}_i\right\}\subset {ℝ}^d,$$ (1)

for which the inner product and norm are defined by ${\langle f,g\rangle }_{{\mathcal{H}}_k}={\sum }_{ij}{c}_i{s}_{j}k\left({\mathit{x}}_i,{\mathit{x}}_j\right)$, $\Vert f{\Vert }_{{\mathcal{H}}_k}^2={\sum }_{ij}{c}_i{c}_{j}k\left({\mathit{x}}_i,{\mathit{x}}_j\right)$, where we assume $g(\mathit{x})={\sum }_i{s}_{i}k\left(\mathit{x},{\mathit{x}}_i\right)$. Denote by ${\mathcal{H}}_k^d:={\mathcal{H}}_k\times \dots \times {\mathcal{H}}_k$ the vector-valued RKHS consisting of ${ℝ}^d$ vector-valued functions $\mathit{\varphi }={\left[{\varphi }^1,\dots ,{\varphi }^d\right]}^{\top }$ with each ${\varphi }^{\mathcal{l}}\in {\mathcal{H}}_k$. See e.g., Berlinet & Thomas-Agnan (2011) for more rigorous treatment. For notation convenience, we do not distinguish distributions on ${ℝ}^d$ and and their density functions.

2. Stein Variational Gradient Descent (SVGD)

We introduce the basic derivation of Stein variational gradient descent (SVGD), which provides a foundation for our new generalization. See Liu & Wang (2016, 2018); Liu (2017) for more details.

Let p(x) be a positive and continuously differentiable probability density function on ${ℝ}^d$. Our goal is to find a set of points (a.k.a. particles) ${\left\{{\mathit{x}}_i\right\}}_{i=1}^n\subset {ℝ}^d$ to approximate p, such that the empirical distribution $q(\mathit{x})={\sum }_i\delta \left(\mathit{x}-{\mathit{x}}_i\right)/n$ of the particles weakly converges to p when n is large. Here $\delta (\cdot )$ denotes the Dirac delta function.

SVGD achieves this by starting from a set of initial particles, and iteratively updating them with a deterministic transformation of form

$${\mathit{x}}_i\leftarrow {\mathit{x}}_i+ϵ{\varphi }_k^{*}\left({\mathit{x}}_i\right),\forall i=1,\cdots ,n,{\mathit{\varphi }}_k^{*}=\underset{\mathit{\varphi }\in {\mathcal{B}}_k}{\arg \max }\left\{-{\frac{\text{d}}{\text{d}ϵ}\mathrm{KL}\left(q_{[ϵ\varphi ]}\Vert p\right)|}_{ϵ=0}\right\},$$ (2)

where $ϵ$ is a small step size, ${\mathit{\varphi }}_k^{*}:{ℝ}^d\to {ℝ}^d$ is an optimal transform function chosen to maximize the decreasing rate of the KL divergence between the distribution of particles and the target p, and $q_{[ϵ\varphi ]}$ denotes the distribution of the updated particles ${\mathit{x}}^{\prime }=\mathit{x}+ϵ\mathit{\varphi }(\mathit{x})$ as x ∼ q, and ${\mathcal{B}}_k$ is the unit ball of RKHS ${\mathcal{H}}_k^d:={\mathcal{H}}_k\times \dots \times {\mathcal{H}}_k$ associated with a positive definite kernel k(x, x^′), that is,

$${\mathcal{B}}_k=\left\{\mathit{\varphi }\in {\mathcal{H}}_k^d:\Vert \mathit{\varphi }{\Vert }_{{\mathcal{H}}_k^d}\le 1\right\}.$$ (3)

Liu & Wang (2016) showed that the objective in (2) can be expressed as a linear functional of Φ,

$$-{\frac{\text{d}}{\text{d}ϵ}\text{KL}\left(q_{[ϵ\mathit{\varphi }]}\Vert p\right)|}_{ϵ=0}={\mathbb{E}}_{\mathit{x}~q}\left[{\mathcal{P}}^{\top }\mathit{\varphi }(\mathit{x})\right],{\mathcal{P}}^{\top }\mathit{\varphi }(\mathit{x})={\nabla }_{\mathit{x}}\log p{(\mathit{x})}^{\top }\mathit{\varphi }(\mathit{x})+{\nabla }_{\mathit{x}}^{\top }\mathit{\varphi }(\mathit{x}),$$ (4)

where $\mathcal{P}$ is a differential operator called Stein operator; here we formally view $\mathcal{P}$ and the derivative operator ${\nabla }_{\mathit{x}}\text{as}{ℝ}^d$ column vectors, hence ${\mathcal{P}}^{\top }\mathit{\varphi }$ and ${\nabla }_{\mathit{x}}^{\top }\mathit{\varphi }$ are viewed as inner products, e.g., ${\nabla }_{\mathit{x}}^{\top }\mathit{\varphi }={\sum }_{\mathcal{l}=1}^d{\nabla }_{x^{\mathcal{l}}}{\varphi }^{\mathcal{l}}$, with $x^{\mathcal{l}}$ and ${\varphi }^{\mathcal{l}}$ being the $\mathcal{l}$-th coordinate of vector x and ϕ, respectively. With (4), it is shown in Liu & Wang (2016) that the solution of (2) is

$${\mathit{\varphi }}_k^{*}(\cdot )\propto {\mathbb{E}}_{\mathit{x}~q}[\mathcal{P}k(\mathit{x},\cdot )]={\mathbb{E}}_{\mathit{x}~q}\left[{\nabla }_{\mathit{x}}\log p(\mathit{x})k(\mathit{x},\cdot )+{\nabla }_{\mathit{x}}k(\mathit{x},\cdot )\right].$$ (5)

Such ${\mathit{\varphi }}_k^{*}$ provides the best update direction for the particles within RKHS ${\mathcal{H}}_k^d$. By taking q to be the empirical measure of the particles, i.e., $q(\mathit{x})={\sum }_{i=1}^n\delta \left(\mathit{x}-{\mathit{x}}_i\right)/n$ and repeatedly applying this update on the particles, we obtain the SVGD algorithm using equations (2) and (5).

3. SVGD with Matrix-valued Kernels

Our goal is to extend SVGD to allow efficient incorporation of precondition information for better optimization. We achieve this by providing a generalization of SVGD that leverages more general matrix-valued kernels, to flexibly incorporate preconditioning information.

The key idea is to observe that the standard SVGD searches for the optimal ϕ in RKHS ${\mathcal{H}}_k^d={\mathcal{H}}_k\times \cdots \times {\mathcal{H}}_k$, a product of d copies of RKHS of scalar-valued functions, which does not allow us to encode potential correlations between different coordinates of $\mathit{\varphi }$. This limitation can be addressed by replacing ${\mathcal{H}}_k^d$ with a more general RKHS of vector-valued functions (called vector-valued RKHS), which uses more flexible matrix-valued positive definite kernels to specify rich correlation structures between different coordinates. In this section, we first introduce the background of vector-valued RKHS with matrix-valued kernels in Section 3.1, and then propose and discuss our generalization of SVGD using matrix-valued kernels in Section 3.2-3.3.

3.1. Vector-Valued RKHS with Matrix-Valued Kernels

We now introduce the background of matrix-valued positive definite kernels, which provides a most general framework for specifying vector-valued RKHS. We focus on the intuition and key ideas in our introduction, and refer the readers to Alvarez et al. (2012); Carmeli et al. (2006) for mathematical treatment.

Recall that a standard real-valued RKHS ${\mathcal{H}}_k$ consists of the closure of the linear span of its kernel function $k(\cdot ,\mathit{x})$ as shown in (1). Vector-valued RKHS can be defined in a similar way, but consist of the linear span of a matrix-valued kernel function:

$$\mathit{f}(\mathit{x})=\sum _i\mathit{K}\left(\mathit{x},{\mathit{x}}_i\right){\mathit{c}}_i,$$ (6)

for any $\{{\mathit{c}}_i\}\subset {ℝ}^d$, where $\mathit{K}:{ℝ}^d\times {ℝ}^d\to {ℝ}^{d\times d}$ is now a matrix-valued kernel function, and ci are vector-valued weights. Similar to the scalar case, we can define an inner product structure ${\langle \mathit{f},\mathit{g}\rangle }_{{\mathcal{H}}_K}=\sum _{ij}{\mathit{c}}_i^{\top }\mathit{K}\left({\mathit{x}}_i,{\mathit{x}}_j\right){\mathit{s}}_j$, where we assume $\mathit{g}={\sum }_i\mathit{K}\left(\mathit{x},{\mathit{x}}_i\right){\mathit{s}}_i$ and hence a norm $\Vert \mathit{f}{\Vert }_{{\mathcal{H}}_k}^2={\sum }_{ij}{\mathit{c}}_i^{\top }K\left({\mathit{x}}_i,{\mathit{x}}_j\right){\mathit{c}}_j$. In order to make the inner product and norm well defined, the matrix-value kernel K is required to be symmetric in that $\mathit{K}\left(\mathit{x},{\mathit{x}}^{\prime }\right)=\mathit{K}{\left({\mathit{x}}^{\prime },\mathit{x}\right)}^{\top }$, and positive definite in that ${\sum }_{ij}{\mathit{c}}_i^{\top }\mathit{K}({\mathit{x}}_i,{\mathit{x}}_j){\mathit{c}}_j\ge 0$, $\left\{{\mathit{x}}_i\right\}\subset {ℝ}^d$ for any $\left\{{\mathit{c}}_i\right\}\subset {ℝ}^d$.

Mathematically, one can show that the closure of the set of functions in (6), equipped with the inner product defined above, defines a RKHS that we denote by ${\mathcal{H}}_{\mathit{K}}$. It is “reproducing” because it has the following reproducing property that generalizes the version for scalar-valued RKHS: for any $\mathit{f}\in {\mathcal{H}}_{\mathit{K}}$ and any $\mathit{c}\in {ℝ}^d$, we have

$$\mathit{f}{(\mathit{x})}^{\top }\mathit{c}={\langle \mathit{f}(\cdot ),\mathit{K}(\cdot ,\mathit{x})\mathit{c}\rangle }_{{\mathcal{H}}_{\mathit{K}}},$$ (7)

where it is necessary to introduce c because the result of the inner product of two functions must be a scalar. A simple example of matrix kernel is $\mathit{K}\left(\mathit{x},{\mathit{x}}^{\prime }\right)=k\left(\mathit{x},{\mathit{x}}^{\prime }\right)\mathit{I}$, where I is the d × d identity matrix. It is related RKHS is ${\mathcal{H}}_{\mathit{K}}={\mathcal{H}}_k\times \cdots \times {\mathcal{H}}_k={\mathcal{H}}_k^d$, as used in the original SVGD.

3.2. SVGD with Matrix-Valued Kernels

It is now natural to leverage matrix-valued kernels to obtain a generalization of SVGD (see Algorithm 1). The idea is simple: we now optimize $\mathit{\varphi }$ in the unit ball of a general vector-valued RKHS ${\mathcal{H}}_{\mathit{K}}$ with a matrix valued kernel $\mathit{K}\left(\mathit{x},{\mathit{x}}^{\prime }\right)$.

$${\mathit{\varphi }}_{\mathit{K}}^{*}=\underset{\mathit{\varphi }\in {\mathcal{H}}_{\mathit{K}}}{\arg \max }\left\{{\mathbb{E}}_{\mathit{x}~q}\left[{\mathcal{P}}^{\top }\mathit{\varphi }(\mathit{x})\right],\mathit{\text{s.t.}}\Vert \mathit{\varphi }{\Vert }_{{\mathcal{H}}_{\mathit{K}}}\le 1\right\}.$$ (8)

This yields a simple closed form solution similar to (5).

Theorem 1.

Let $\mathit{K}\left(\mathit{x},{\mathit{x}}^{\prime }\right)$ be a matrix-valued positive definite kernel that is continuously differentiable on $\mathit{x}and{\mathit{x}}^{\prime }$, the optimal ${\mathit{\varphi }}^{*}$ in (8) is

$${\mathit{\varphi }}_{\mathit{K}}^{*}(\cdot )\propto {\mathbb{E}}_{\mathit{x}~q}[\mathit{K}(\cdot ,\mathit{x})\mathcal{P}]={\mathbb{E}}_{\mathit{x}~q}[\mathit{K}(\cdot ,\mathit{x}){\nabla }_{\mathit{x}}\log p(\mathit{x})+\mathit{K}(\cdot ,\mathit{x}){\nabla }_{\mathit{x}}],$$ (9)

Algorithm 1.

Stein Variational Gradient Descent with Matrix-valued Kernels (Matrix SVGD)

Input: A (possibly unnormalized) differentiable density function p(x) in ${ℝ}^d$. A matrix-valued positive definite kernel $\mathit{K}\left(\mathit{x},{\mathit{x}}^{\prime }\right)$. Step size $ϵ$.

Goal: Find a set of particles ${\left\{{\mathit{x}}_i\right\}}_{i=1}^n$ to represent the distribution p.

Initialize a set of particles ${\left\{{\mathit{x}}_i\right\}}_{i=1}^n$, e.g., by drawing from some simple distribution.

repeat

${\mathit{x}}_i\leftarrow {\mathit{x}}_i+\frac{ϵ}{n}\sum _{j=1}^n[\mathit{K}({\mathit{x}}_i,{\mathit{x}}_j){\nabla }_{{\mathit{x}}_j}\log p({\mathit{x}}_j)+\mathit{K}({\mathit{x}}_i,{\mathit{x}}_j){\nabla }_{{\mathit{x}}_j}],$

where $\mathit{K}(\cdot ,\mathit{x}){\nabla }_{\mathit{x}}$ is formally defined as the product of matrix $\mathit{K}(\cdot ,\mathit{x})$ and vector ${\nabla }_{\mathit{x}}$. The $\mathcal{l}$-th element of $\mathit{K}(\cdot ,\mathit{x}){\nabla }_{\mathit{x}}\text{is}{\left(\mathit{K}(\cdot ,\mathit{x}){\nabla }_{\mathit{x}}\right)}_{\mathcal{l}}=\sum _{m=1}^d{\nabla }_{x^m}{\mathit{K}}_{\mathcal{l},m}(\cdot ,\mathit{x})$, see also (10).

until Convergence

where the Stein operator $\mathcal{P}$ and derivative operator ∇x are again formally viewed as ${ℝ}^d$-valued column vectors, and $\mathit{K}(\cdot ,\mathit{x})\mathcal{P}$ and $\mathit{K}(\cdot ,\mathit{x}){\nabla }_{\mathit{x}}$ are interpreted by the matrix multiplication rule. Therefore, $\mathit{K}(\cdot ,\mathit{x})\mathcal{P}$ is a ${ℝ}^d$-valued column vector, whose $\mathcal{l}$-th element is defined by

$${(\mathit{K}(\cdot ,\mathit{x})\mathcal{P})}_{\mathcal{l}}=\sum _{m=1}^d({\mathit{K}}_{\mathcal{l},m}(\cdot ,\mathit{x}){\nabla }_{x^m}\log p(\mathit{x})+{\nabla }_{x^m}{\mathit{K}}_{\mathcal{l},m}(\cdot ,\mathit{x})),$$ (10)

where $K_{\mathcal{l},m}\left(\mathit{x},{\mathit{x}}^{\prime }\right)$ denotes the $(\mathcal{l},m)$-element of matrix $\mathit{K}\left(\mathit{x},{\mathrm{mathvariant=\text{'}bold-italic\text{'}x}}^{\prime }\right)$ and x^m the m-th element of x.

Similar to the case of standard SVGD, recursively applying the optimal transform ${\mathit{\varphi }}_{\mathit{K}}^{*}$ on the particles yields a general SVGD algorithm shown in Algorithm 1, which we call matrix SVGD.

Parallel to vanilla SVGD, the gradient of matrix SVGD in (9) consists of two parts that account for optimization and diversity, respectively: the first part is a weighted average of gradient ${\nabla }_{\mathit{x}}\log p(\mathit{x})$ multiplied by a matrix-value kernel $\mathit{K}(\cdot ,\mathit{x})$ the other part consists of the gradient of the matrix-valued kernel K, which, like standard SVGD, serves as a repulsive force to keep the particles away from each other to reflect the uncertainty captured in distribution p.

Matrix SVGD includes various previous variants of SVGD as special cases. The vanilla SVGD corresponds to the case when $\mathit{K}\left(\mathit{x},{\mathit{x}}^{\prime }\right)=k\left(\mathit{x},{\mathit{x}}^{\prime }\right)\mathit{I}$, with I as the d × d identity matrix; the gradient-free SVGD of Han & Liu (2018) can be treated as the case when $\mathit{K}\left(\mathit{x},{\mathit{x}}^{\prime }\right)=k\left(\mathit{x},{\mathit{x}}^{\prime }\right)w(\mathit{x})w\left({\mathit{x}}^{\prime }\right)\mathit{I}$, where w(x) is an importance weight function; the graphical SVGD of Wang et al. (2018); Zhuo et al. (2018) corresponds to a diagonal matrix-valued kernel: $\mathit{K}\left(\mathit{x},{\mathit{x}}^{\prime }\right)=\mathrm{diag}\left[{\left\{k_{\mathcal{l}}\left(\mathit{x},{\mathit{x}}^{\prime }\right)\right\}}_{\mathcal{l}=1}^d\right]$, where each $k_{\mathcal{l}}\left(\mathit{x},{\mathit{x}}^{\prime }\right)$ is a “local” scalar-valued kernel function related to the $\mathcal{l}$-th coordinate $x^{\mathcal{l}}$ of vector $\mathit{x}$.

3.3. Matrix-Valued Kernels and Change of Variables

It is well known that preconditioned gradient descent can be interpreted as applying standard gradient descent on a reparameterization of the variables. For example, let $\mathit{y}={\mathit{Q}}^{1/2}\mathit{x}$, where Q is a positive definite matrix, then log $p(\mathit{x})=\log p\left({\mathit{Q}}^{-1/2}\mathit{y}\right)$. Applying gradient descent on y and transform it back to the updates on x yields a preconditioned gradient update $\mathit{x}\leftarrow \mathit{x}+ϵ{\mathit{Q}}^{-1}{\nabla }_{\mathit{x}}\log p(\mathit{x})$.

We now extend this idea to SVGD, for which matrix-valued kernels show up naturally as a consequence of change of variables. This justifies the use of matrix-valued kernels and provides guidance on the practical choice of matrix-valued kernels. We start with a basic result of how matrix-valued kernels change under change of variables (see Paulsen & Raghupathi (2016)).

Lemma 2.

Assume ${\mathcal{H}}_0$ is an RKHS with a matrix kernel ${\mathit{K}}_0:{ℝ}^d\times {ℝ}^d\to {ℝ}^{d\times d}$. Let $\mathcal{H}$ be the set of functions formed by

$$\mathit{\varphi }(\mathit{x})=M(\mathit{x}){\mathit{\varphi }}_0(t(\mathit{x})),\forall {\mathit{\varphi }}_0\in {\mathcal{H}}_0,$$

where $\mathit{M}:{ℝ}^d\to {ℝ}^{d\times d}$ is a fixed matrix-valued function and we assume M (x) is an invertible matrix for all x, and $\mathit{t}:{ℝ}^d\to {ℝ}^d$ is a fixed continuously differentiable one-to-one transform on ${ℝ}^d$.

For $\forall \mathit{\varphi },{\mathit{\varphi }}^{\prime }\in \mathcal{H}$, we can identity an unique ${\mathit{\varphi }}_0,{\mathit{\varphi }}_0^{\prime }\in {\mathcal{H}}_0$ such that $\mathit{\varphi }(\mathit{x})=\mathit{M}(\mathit{x}){\mathit{\varphi }}_0(t(\mathit{x}))$ and ${\mathit{\varphi }}^{\prime }(\mathit{x})=\mathit{M}(\mathit{x}){\mathit{\varphi }}_0^{\prime }(\mathit{t}(\mathit{x}))$. Define the inner product on $\mathcal{H}$ via ${\langle \mathit{\varphi },{\mathit{\varphi }}^{\prime }\rangle }_{\mathcal{H}}={\langle {\mathit{\varphi }}_0,{\mathit{\varphi }}_0^{\prime }\rangle }_{{\mathcal{H}}_0}$, then $\mathcal{H}$ is also a vector-valued RKHS, whose matrix-valued kernel is

$$\mathit{K}(\mathit{x},{\mathit{x}}^{\prime })=\mathit{M}(\mathit{x}){\mathit{K}}_0(\mathit{t}(\mathit{x}),\mathit{t}({\mathit{x}}^{\prime }))\mathit{M}{({\mathit{x}}^{\prime })}^{\top }.$$

We now present a key result, which characterizes the change of kernels when we apply invertible variable transforms on the SVGD trajectory.

Theorem 3.

i) Let p and q be two distributions and p₀, q₀ the distribution of ${\mathit{x}}_0=\mathit{t}(\mathit{x})$ when x is drawn from p, q, respectively, where t is a continuous differentiable one-to-one map on ${ℝ}^d$. Assume p is a continuous differentiable density with Stein operator $\mathcal{P}$, and ${\mathcal{P}}_0$ the Stein operator of p₀. We have

$${\mathbb{E}}_{\mathit{x}~q_0}\left[{\mathcal{P}}_0^{\top }{\mathit{\varphi }}_0(\mathit{x})\right]={\mathbb{E}}_{\mathit{x}~q}\left[{\mathcal{P}}^{\top }\mathit{\varphi }(\mathit{x})\right],with\mathit{\varphi }(\mathit{x}):=\nabla t{(\mathit{x})}^{-1}{\mathit{\varphi }}_0(t(\mathit{x})),$$ (11)

where ∇t is the Jacobian matrix of t.

ii) Therefore, in the asymptotics of infinitesimal step size $(ϵ\to 0^{+})$, running SVGD with kernel ${\mathit{K}}_0$ on p₀ is equivalent to running SVGD on p with kernel

$$\mathit{K}\left(\mathit{x},{\mathit{x}}^{\prime }\right)=\nabla \mathit{t}{(\mathit{x})}^{-1}{\mathit{K}}_0\left(\mathit{t}(\mathit{x}),\mathit{t}\left({\mathit{x}}^{\prime }\right)\right)\nabla \mathit{t}{\left({\mathit{x}}^{\prime }\right)}^{-\top },$$

in the sense that the trajectory of these two SVGD can be mapped to each other by the one-to-one map t (and its inverse).

3.4. Practical Choice of Matrix-Valued Kernels

Theorem 3 suggests a conceptual procedure for constructing proper matrix kernels to incorporate desirable preconditioning information: one can construct a one-to-one map t so that the distribution $p_0$ of ${\mathit{x}}_0=\mathit{t}(\mathit{x})$ is an easy-to-sample distribution with a simpler kernel, which can be a standard scalar-valued kernel or a simple diagonal kernel. Practical choices of t often involve rotating x with either Hessian matrix or Fisher information, allowing us to incorporating these information into SVGD. In the sequel, we first illustrate this idea for simple Gaussian cases and then discuss practical approaches for non-Gaussian cases.

Constant Preconditioning Matrices

Consider the simple case when p is multivariate Gaussian, e.g., $\log p(\mathit{x})=-\frac{1}{2}{\mathit{x}}^{\top }\mathit{Q}\mathit{x}+const$, where Q is a positive definite matrix. In this case, the distribution p₀ of the transformed variable $\mathit{t}(\mathit{x})={\mathit{Q}}^{1/2}\mathit{x}$ is the standard Gaussian distribution that can be better approximated with a simpler kernel ${\mathit{K}}_0\left(\mathit{x},{\mathit{x}}^{\prime }\right)$, which can be chosen to be the standard RBF kernel suggested in Liu & Wang (2016), the graphical kernel suggested in Wang et al. (2018), or the linear kernels suggested in Liu & Wang (2018). Theorem 3 then suggests to use a kernel of form

$${\mathit{K}}_{\mathit{Q}}\left(\mathit{x},{\mathit{x}}^{\prime }\right):={\mathit{Q}}^{-1/2}{\mathit{K}}_0\left({\mathit{Q}}^{1/2}\mathit{x},{\mathit{Q}}^{1/2}{\mathit{x}}^{\prime }\right){\mathit{Q}}^{-1/2},$$ (12)

in which Q is applied on both the input x and the output side. As an example, taking ${\mathit{K}}_{\text{0}}$ to be the scalar-valued Gaussian RBF kernel gives

$${\mathit{K}}_{\mathit{Q}}\left(\mathit{x},{\mathit{x}}^{\prime }\right)={\mathit{Q}}^{-1}\exp \left(-\frac{1}{2h}{\Vert \mathit{x}-{\mathit{x}}^{\prime }\Vert }_{\mathit{Q}}^2\right),$$ (13)

where ${\Vert \mathit{x}-{\mathit{x}}^{\prime }\Vert }_{\mathit{Q}}^2:={\left(\mathit{x}-{\mathit{x}}^{\prime }\right)}^{\top }\mathit{Q}\left(\mathit{x}-{\mathit{x}}^{\prime }\right)$ and h is a bandwidth parameter. Define ${\mathit{K}}_{0,\mathit{Q}}\left(\mathit{x},{\mathit{x}}^{\prime }\right):={\mathit{K}}_0\left({\mathit{Q}}^{1/2}\mathit{x},{\mathit{Q}}^{1/2}{\mathit{x}}^{\prime }\right)$. One can show that the SVGD direction of the kernel in (12) equals

$${\mathit{\varphi }}_{{\mathit{K}}_{\mathit{Q}}}^{*}(\cdot )={\mathit{Q}}^{-1}{\mathbb{E}}_{\mathit{x}~q}\left[\nabla \log p(\mathit{x}){\mathit{K}}_{0,\mathit{Q}}(\cdot ,\mathit{x})+{\mathit{K}}_{0,\mathit{Q}}(\cdot ,\mathit{x}){\nabla }_{\mathit{x}}\right]={\mathit{Q}}^{-1}{\mathit{\varphi }}_{{\mathit{K}}_{0,\mathit{Q}}}^{*}(\cdot ),$$ (14)

which is a linear transform of the SVGD direction of kernel ${\mathit{K}}_{0,\mathit{Q}}\left(\mathit{x},{\mathit{x}}^{\prime }\right)$ with matrix Q⁻¹.

In practice, when p is non-Gaussian, we can construct Q by taking averaging over the particles. For example, denote by $\mathit{H}(\mathit{x})=-{\nabla }_{\mathit{x}}^2\log p(\mathit{x})$ the negative Hessian matrix at x, we can construct Q by

$$\mathit{Q}=\sum _{i=1}^n\mathit{H}\left({\mathit{x}}_i\right)/n,$$ (15)

where ${\left\{{\mathit{x}}_i\right\}}_{i=1}^n$ are the particles from the previous iteration. We may replace H with the Fisher information matrix to obtain a natural gradient like variant of SVGD.

Point-wise Preconditioning

A constant preconditioning matrix can not reflect different curvature or geometric information at different points. A simple heuristic to address this limitation is to replace the constant matrix Q with a point-wise matrix function Q(x); this motivates a kernel of form

$$\mathit{K}\left(\mathit{x},{\mathit{x}}^{\prime }\right)={\mathit{Q}}^{-1/2}(\mathit{x}){\mathit{K}}_0\left({\mathit{Q}}^{1/2}(\mathit{x})\mathit{x},{\mathit{Q}}^{1/2}\left({\mathit{x}}^{\prime }\right){\mathit{x}}^{\prime }\right){\mathit{Q}}^{-1/2}\left({\mathit{x}}^{\prime }\right).$$

Unfortunately, this choice may yield expensive computation and difficult implementation in practice, because the SVGD update involves taking the gradient of the kernel $\mathit{K}\left(\mathit{x},{\mathit{x}}^{\prime }\right)$, which would need to differentiate through matrix valued function Q(x). When Q(x) equals the Hessian matrix, for example, it involves taking the third order derivative of log p(x), yielding an unnatural algorithm.

Mixture Preconditioning

We instead propose a more practical approach to achieve point-wise preconditioning with a much simpler algorithm. The idea is to use a weighted combination of several constant preconditioning matrices. This involves leveraging a set of anchor points ${\left\{{\mathit{z}}_{\mathcal{l}}\right\}}_{\mathcal{l}=1}^m\subset {ℝ}^d$, each of which is associated with a preconditioning matrix ${\mathit{Q}}_{\mathcal{l}}=\mathit{Q}\left({\mathit{z}}_{\mathcal{l}}\right)$ (e.g., their Hessian or Fisher information matrices). In practice, the anchor points ${\left\{{\mathit{z}}_{\mathcal{l}}\right\}}_{\mathcal{l}=1}^m$ can be conveniently set to be the same as the particles ${\left\{{\mathit{x}}_i\right\}}_{i=1}^n$. We then construct a kernel by

$$\mathit{K}\left(\mathit{x},{\mathit{x}}^{\prime }\right)=\sum _{\mathcal{l}=1}^m{\mathit{K}}_{{\mathit{Q}}_{\mathcal{l}}}\left(\mathit{x},{\mathit{x}}^{\prime }\right)w_{\mathcal{l}}(\mathit{x})w_{\mathcal{l}}\left({\mathit{x}}^{\prime }\right),$$ (16)

where ${\mathit{K}}_{{\mathit{Q}}_{\mathcal{l}}}\left(\mathit{x},{\mathit{x}}^{\prime }\right)$ is defined in (12) or (13), and $w_{\mathcal{l}}(\mathit{x})$ is a positive scalar-valued function that decides the contribution of kernel ${\mathit{K}}_{{\mathit{Q}}_{\mathcal{l}}}$ at point x. Here $w_{\mathcal{l}}(\mathit{x})$ should be viewed as a mixture probability, and hence should satisfy ${\sum }_{\mathcal{l}=1}^m{w}_{\mathcal{l}}(\mathit{x})=1$ for all x. In our empirical studies, we take $w_{\mathcal{l}}(\mathit{x})$ as the Gaussian mixture probability from the anchor points:

$$w_{\mathcal{l}}(\mathit{x})=\frac{\mathcal{N}\left(\mathit{x};{\mathit{z}}_{\mathcal{l}},{\mathit{Q}}_{\mathcal{l}}^{-1}\right)}{\sum _{{\mathcal{l}}^{\prime }=1}^m\mathcal{N}\left(\mathit{x};{\mathit{z}}_{{\mathcal{l}}^{\prime }},{\mathit{Q}}_{{\mathcal{l}}^{\prime }}^{-1}\right)},\mathcal{N}\left(\mathit{x};{\mathit{z}}_{\mathcal{l}},{\mathit{Q}}_{\mathcal{l}}^{-1}\right):=\frac{1}{{\mathit{Z}}_{\mathcal{l}}}\exp \left(-\frac{1}{2}{\Vert \mathit{x}-{\mathit{z}}_{\mathcal{l}}\Vert }_{{\mathit{Q}}_{\mathcal{l}}}^2\right),$$ (17)

where $Z_{\mathcal{l}}={(2\pi )}^{d/2}\det {\left(Q_{\mathcal{l}}\right)}^{-1/2}$. In this way, each point x is mostly influenced by the anchor point closest to it, allowing us to apply different preconditioning for different points. Importantly, the SVGD update direction related to the kernel in (16) has a simple and easy-to-implement form:

$${\mathit{\varphi }}_{\mathit{K}}^{*}(\cdot )=\sum _{\mathcal{l}=1}^m{w}_{\mathcal{l}}(\cdot ){\mathbb{E}}_{\mathit{x}~q}\left[\left(w_{\mathcal{l}}(\mathit{x}){\mathit{K}}_{{\mathit{Q}}_{\mathcal{l}}}(\cdot ,\mathit{x})\right)\mathcal{P}\right]=\sum _{\mathcal{l}=1}^m{w}_{\mathcal{l}}(\cdot ){\mathit{\varphi }}_{w_{\mathcal{l}}{\mathit{K}}_{{\mathit{Q}}_{\mathcal{l}}}}^{*}(\cdot ),$$ (18)

which is a weighted sum of a number of SVGD directions with constant preconditioning matrices (but with an asymmetric kernel $w_{\mathcal{l}}(\mathit{x}){\mathit{K}}_{{\mathit{Q}}_{\mathcal{l}}}(\cdot ,\mathit{x}))$.

A Remark on Stein Variational Newton (SVN)

Detommaso et al. (2018) provided a Newton-like variation of SVGD. It is motivated by an intractable functional Newton framework, and arrives a practical algorithm using a series of approximation steps. The update of SVN is

$${\mathit{x}}_i\leftarrow {\mathit{x}}_i+ϵ{\tilde{\mathit{H}}}_i^{-1}{\mathit{\varphi }}_k^{*}\left({\mathit{x}}_i\right),\forall i=1,\dots ,n,$$ (19)

where ${\mathit{\varphi }}_k^{*}(\cdot )$ is the standard SVGD gradient, and ${\tilde{\mathit{H}}}_i$ is a Hessian like matrix associated with particle x_i, defined by

$${\tilde{\mathit{H}}}_i={\mathbb{E}}_{\mathit{x}~q}\left[\mathit{H}(\mathit{x})k{(\mathit{x},{\mathit{x}}_i)}^2+{({\nabla }_{{\mathit{x}}_i}k(\mathit{x},{\mathit{x}}_i))}^{\otimes 2}\right],$$

where $\mathit{H}(\mathit{x})=-{\nabla }_{\mathit{x}}^2\log p(\mathit{x})$ and ${\mathit{w}}^{\otimes 2}:=\mathit{w}{\mathit{w}}^{\top }$. Due to the approximation introduced in the derivation of SVN, it does not correspond to a standard functional gradient flow like SVGD (unless ${\tilde{\mathit{H}}}_i=\mathit{Q}$ for all i, in which case it reduces to using a constant preconditioning matrix on SVGD like (14)). SVN can be heuristically viewed as a “hard” variant of (18), which assigns each particle with its own preconditioning matrix with probability one, but the mathematical form do not match precisely. On the other hand, it is useful to note that the set of fixed points of SVN update (19) is the identical to that of the standard SVGD update with ${\mathit{\varphi }}_k^{*}(\cdot )$, once all ${\tilde{\mathit{H}}}_i$ are positive definite matrices. This is because at the fixed points of (19), we have ${\tilde{\mathit{H}}}_i^{-1}{\mathit{\varphi }}_k^{*}\left({\mathit{x}}_i\right)=0\text{for}\forall i=1,\dots ,n$, which is equivalent to ${\mathit{\varphi }}_k^{*}\left({\mathit{x}}_i\right)=0$, $\forall i$ when all the ${\tilde{\mathit{H}}}_i$, $\forall i$ are positive definite. Therefore, SVN can be justified as an alternative fixed point iteration method to achieve the same set of fixed points as the standard SVGD.

4. Experiments

We demonstrate the effectiveness of our matrix SVGD on various practical tasks. We start with a toy example and then proceed to more challenging tasks that involve logistic regression, neural networks and recurrent neural networks. For our method, we take the preconditioning matrices to be either Hessian or Fisher information matrices, depending on the application. For large scale Fisher matrices in (recurrent) neural networks, we leverage the Kronecker-factored (KFAC) approximation by Martens & Grosse (2015); Martens et al. (2018) to enable efficient computation. We use RBF kernel for vanilla SVGD. The kernel ${\mathit{K}}_0\left(\mathit{x},{\mathit{x}}^{\prime }\right)$ in our matrix SVGD (see (12) and (13)) is also taken to be Gaussian RBF. Following Liu & Wang (2016), we choose the bandwidth of the Gaussian RBF kernels using the standard median trick and use Adagrad (Duchi et al., 2011) for stepsize. Our code is available at https://github.com/dilinwang820/matrix_ssvgd.

The algorithms we test are summarized here:

Vanilla SVGD, using the code by Liu & Wang (2016);

Matrix-SVGD (average), using the constant preconditioning matrix kernel in (13), with Q to be either the average of the Hessian matrices or Fisher matrices of the particles (e.g., (15));

Matrix-SVGD (mixture), using the mixture preconditioning matrix kernel in (16), where we pick the anchor points to be particles themselves, that is, ${\left\{{\mathit{z}}_{\mathcal{l}}\right\}}_{\mathcal{l}=1}^m={\left\{{\mathit{x}}_i\right\}}_{i=1}^n$;

Stein variational Newton (SVN), based on the implementation of Detommaso et al. (2018);

Preconditioned Stochastic Langevin Dynamics (pSGLD), which is a variant of SGLD (Li et al., 2016), using a diagonal approximation of Fisher information as the preconditioned matrix.

4.1. Two-Dimensional Toy Examples

Settings

We start with illustrating our method using a Gaussian mixture toy model (Figure 1), with exact Hessian matrices for preconditioning. For fair comparison, we search the best learning rate for all algorithms exhaustively. We use 50 particles for all the cases. We use the same initialization for all methods with the same random seeds.

Figure 1:

Figure (a)-(e) show the particles obtained by various methods at the 30-th iteration. Figure (f) plots the log MMD (Gretton et al., 2012) vs. training iteration starting from the 10-th iteration. We use the standard RBF kernel for evaluating MMD.

Results

Figure 1 show the results for 2D toy examples. Appendix B shows more visualization and results on more examples. We can see that methods with Hessian information generally converge faster than vanilla SVGD, and Matrix-SVGD (mixture) yields the best performance.

4.2. Bayesian Logistic Regression

Settings

We consider the Bayesian logistic regression model for binary classification. Let D = ${\left\{\left({\mathit{x}}_j,y_j\right)\right\}}_{j=1}^N$ be a dataset with feature vector x_j and binary label $y_j\in \{0,1\}$. The distribution of interest is

$$p(\mathit{\theta }|\mathit{D})\propto p(\mathit{D}|\mathit{\theta })p(\mathit{\theta })\text{with}p(\mathit{D}|\mathit{\theta })=\prod _{j=1}^N\left[y_j\sigma \left({\mathit{\theta }}^{\top }{\mathit{x}}_j\right)+\left(1-y_j\right)\sigma \left(-{\mathit{\theta }}^{\top }{\mathit{x}}_j\right)\right],$$

where σ(z) := 1/(1 + exp(−z)), and $p_0(\mathit{\theta })$ is the prior distribution, which we set to be standard normal $\mathcal{N}(\mathit{\theta };\mathbf{0},\mathit{I})$. The goal is to approximate the posterior distribution $p(\mathit{\theta }|\mathit{D})$ with a set of particles ${\left\{{\mathit{\theta }}_i\right\}}_{i=1}^n$, and then use it to predict the class labels for testing data points. We compare our methods with preconditioned stochastic gradient Langevin dynamics (pSGLD) (Li et al., 2016). Because pSGLD is a sequential algorithm, for fair comparison, we obtain the samples of pSGLD by running n parallel chains of pSGLD for estimation. The preconditioning matrix in both pSGLD and matrix SVGD is taken to be the Fisher information matrix.

We consider the binary Covtype² dataset with 581, 012 data points and 54 features. We partition the data into 70% for training, 10% for validation and 20% for testing. We use Adagrad optimizer with a mini-batch size of 256. We choose the best learning rate from [0.001, 0.005, 0.01, 0.05, 0.1, 0.5, 1.0] for each method on the validation set. For all the experiments and algorithms, we use n = 20 particles. Results are average over 20 random trials.

Results

Figure 2 (a) and (b) show the test accuracy and test log-likelihood of different algorithms. We can see that both Matrix-SVGD (average) and Matrix-SVGD (mixture) converge much faster than both vanilla SVGD and pSGLD, reaching an accuracy of 0.75 in less than 500 iterations.

Figure 2:

(a)-(b) Results of Bayesian Logistic regression on the Covtype dataset. (c)-(d) Average test RMSE and log-likelihood vs. training batches on the Protein dataset for Bayesian neural regression.

4.3. Neural Network Regression

Settings

We apply our matrix SVGD on Bayesian neural network regression on UCI datasets. For all experiments, we use a two-layer neural network with 50 hidden units with ReLU activation functions. We assign isotropic Gaussian priors $\mathcal{N}(0,\alpha \text{I})$ to the neural network weights. We set α = 0.5. All datasets³ are randomly partitioned into 90% for training and 10% for testing. All results are averaged over 20 random trials, except for Protein and Year, on which 5 random trials are performed. We use n = 10 particles for all methods. We use Adam optimizer with a mini-batch size of 100; for large dataset such as Year, we set the mini-batch size to be 1000. We use the Fisher information matrix with Kronecker-factored (KFAC) approximation for preconditioning.

Results

Table 1 shows the performance in terms of the test RMSE and the test log-likelihood. We can see that both Matrix-SVGD (average) and Matrix-SVGD (mixture), which use second-order information, achieve better performance than vanilla SVGD. Matrix-SVGD (mixture) yields the best performance for both test RMSE and test log-likelihood in most cases. Figure 2 (c)-(d) show that both variants of Matrix-SVGD converge much faster than vanilla SVGD and pSGLD on the Protein dataset.

Table 1:

Average test RMSE and log-likelihood in test data for UCI regression benchmarks.

	Test RMSE				Test Log-Likelihood
Dataset	pSGLD	Vanilla SVGD	Matrix-SVGD (average)	Matrix-SVGD (mixture)	pSGLD	Vanilla SVGD	Matrix-SVGD (average)	Matrix-SVGD (mixture)
Boston	2.699±0.155	2.785±0.169	2.898±0.184	2.717±0.166	−2.847±0.182	−2.706±0.158	−2.669±0.141	−2.861±0.207
Concrete	5.053±0.124	5.027±0.116	4.869±0.124	4.721±0.111	−3.206±0.056	−3.064±0.034	−3.150±0.054	−3.207±0.071
Energy	0.985±0.024	0.889±0.024	0.795±0.025	0.868±0.025	−1.395±0.029	−1.315±0.020	−1.135±0.026	−1.249±0.036
Kin8nm	0.091±0.001	0.093±0.001	0.092±0.001	0.090±0.001	0.973±0.010	0.964±0.012	0.956±0.011	0.975±0.011
Naval	0.002±0.000	0.004±0.000	0.001±0.000	0.000±0.000	4.535±0.093	4.312±0.087	5.383±0.081	5.639±0.048
Combined	4.042±0.034	4.088±0.033	4.056±0.033	4.029±0.033	−2.821±0.009	−2.832±0.009	−2.824±0.009	−2.817±0.009
Wine	0.641±0.009	0.645±0.009	0.637±0.008	0.637±0.009	−0.984±0.016	−0.997±0.019	−0.980±0.016	−0.988±0.018
Protein	4.300±0.018	4.186±0.017	3.997±0.018	3.852±0.014	−2.874±0.004	−2.846±0.003	−2.796±0.004	−2.755±0.003
Year	8.630±0.007	8.686±0.010	8.637±0.005	8.594±0.009	−3.568±0.002	−3.577±0.002	−3.569±0.001	−3.561±0.002

4.4. Sentence Classification With Recurrent Neural Networks (RNN)

Settings

We consider the sentence classification task on four datasets: MR (Pang & Lee, 2005), CR (Hu & Liu, 2004), SUBJ (Pang & Lee, 2004), and MPQA (Wiebe et al., 2005). We use a recurrent neural network (RNN) based model, $p(y|\mathit{x})=\mathrm{softmax}\left({\mathit{w}}_y^{\top }{\mathit{h}}_{RNN}(\mathit{x},\mathit{v})\right)$, where x is the input sentence, y is a discrete-valued label of the sentence, and w_y is a weight coefficient related to label class y. And ${\mathit{h}}_{RNN}(\mathit{x},\mathit{v})$ is an RNN function with parameter v using a one-layer bidirectional GRU model (Cho et al., 2014) with 50 hidden units. We apply matrix SVGD to infer the posterior of $\mathit{w}=\{{\mathit{w}}_y:\forall y\}$, while updating the RNN weights v using typical stochastic gradient descent. In all experiments, we use the pre-processed text data provided in Gan et al. (2016). For all the datasets, we conduct 10-fold cross-validation for evaluation. We use n = 10 particles for all the methods. For training, we use a mini-batch size of 50 and run all the algorithms for 20 epochs. We use the Fisher information matrix for preconditioning.

Results

Table 2 shows the results of testing classification errors. We can see that Matrix-SVGD (mixture) generally performs the best among all algorithms.

Table 2:

Sentence classification errors measured with four benchmarks.

Method	MR	CR	SUBJ	MPQA
SGLD	20.52	18.65	7.66	11.24
pSGLD	19.75	17.50	6.99	10.80
Vanilla SVGD	19.73	18.07	6.67	10.58
Matrix-SVGD (average)	19.22	17.29	6.76	10.79
Matrix-SVGD (mixture)	19.09	17.13	6.59	10.71

5. Conclusion

We present a generalization of SVGD by leveraging general matrix-valued positive definite kernels, which allows us to flexibly incorporate various preconditioning matrices, including Hessian and Fisher information matrices, to improve exploration in the probability landscape. We test our practical algorithms on various practical tasks and demonstrate its efficiency compared to various existing methods.

A. Proof

Proof of Theorem 1

Let $e_{\mathcal{l}}$ be the column vector with 1 in ${\mathcal{l}}^{th}$ coordinate and 0 elsewhere. By the RKHS reproducing property (7) we have

$$\begin{gathered} {\mathbb{E}}_{\mathit{x}~q}\left[{\mathcal{P}}^{\top }\mathit{\varphi }(\mathit{x})\right]={\mathbb{E}}_{\mathit{x}~q}\left[{\nabla }_{\mathit{x}}\log p{(\mathit{x})}^{\top }\mathit{\varphi }(\mathit{x})+{\nabla }_{\mathit{x}}^{\top }\mathit{\varphi }(\mathit{x})\right] \\ ={\mathbb{E}}_{\mathit{x}~q}\left[\mathit{\varphi }{(\mathit{x})}^{\top }{\nabla }_{\mathit{x}}\log p(\mathit{x})+\sum _{\mathcal{l}=1}^d{\nabla }_{{\mathit{x}}^{\mathcal{l}}}\mathit{\varphi }{(\mathit{x})}^{\top }{\mathit{e}}_{\mathcal{l}}\right] \\ ={\mathbb{E}}_{\mathit{x}~q}\left[{\langle \mathit{\varphi }(\cdot ),\mathit{K}(\cdot ,\mathit{x}){\nabla }_{\mathit{x}}\log p(\mathit{x})\rangle }_{{\mathcal{H}}_{\mathit{K}}}+\sum _{\mathcal{l}=1}^d{\nabla }_{{\mathit{x}}^{\mathcal{l}}}{\langle \mathit{\varphi }(\cdot ),\mathit{K}(\cdot ,\mathit{x}){\mathit{e}}_{\mathcal{l}}\rangle }_{{\mathcal{H}}_{\mathit{K}}}\right] \\ ={\langle \mathit{\varphi }(\cdot ),{\mathbb{E}}_{\mathit{x}~q}\left[\mathit{K}(\cdot ,\mathit{x}){\nabla }_{\mathit{x}}\log p(\mathit{x})+\sum _{\mathcal{l}=1}^d{\nabla }_{{\mathit{x}}^{\prime }}\mathit{K}(\cdot ,\mathit{x}){\mathit{e}}_{\mathcal{l}}\right]\rangle }_{{\mathcal{H}}_{\mathit{K}}} \\ ={\langle \mathit{\varphi }(\cdot ),{\mathbb{E}}_{\mathit{x}~q}\left[\mathit{K}(\cdot ,\mathit{x}){\nabla }_{\mathit{x}}\log p(\mathit{x})+\mathit{K}(\cdot ,\mathit{x}){\nabla }_{\mathit{x}}\right]\rangle }_{{\mathcal{H}}_{\mathit{K}}} \\ ={\langle \mathit{\varphi }(\cdot ),{\mathbb{E}}_{\mathit{x}~q}[\mathit{K}(\cdot ,\mathit{x})\mathcal{P}]\rangle }_{{\mathcal{H}}_{\mathit{K}}}, \end{gathered}$$

The optimization in (8) is hence

$$\underset{\mathit{\varphi }\in {\mathcal{H}}_{\mathit{K}}}{\max }{\langle \mathit{\varphi }(\cdot ),{\mathbb{E}}_{\mathit{x}~q}[\mathit{K}(\cdot ,\mathit{x})\mathcal{P}]\rangle }_{{\mathcal{H}}_{\mathit{K}}},\mathrm{s.t.}\Vert \mathit{\varphi }{\Vert }_{{\mathcal{H}}_{\mathit{K}}}\le 1,$$

whose solution is ${\mathit{\varphi }}^{*}(\cdot )\propto {\mathbb{E}}_{\mathit{x}~q}[\mathit{K}(\cdot ,\mathit{x})\mathcal{P}]$

Proof of Lemma 2

This is a basic result of RKHS, which can be found in classical textbooks such as Paulsen & Raghupathi (2016). The key idea is to show that $\mathit{K}\left(\mathit{x},{\mathit{x}}^{\prime }\right)$ satisfies the reproducing property for $\mathcal{H}$. Recall the reproducing property of ${\mathcal{H}}_0$:

$${\mathit{\varphi }}_0{(\mathit{x})}^{\top }\mathit{c}={\langle {\mathit{\varphi }}_0,{\mathit{K}}_0(\cdot ,\mathit{x})\mathit{c}\rangle }_{{\mathcal{H}}_0},\forall \mathit{c}\in {ℝ}^d.$$

Taking $\mathit{\varphi }(\mathit{x})=\mathit{M}(\mathit{x}){\mathit{\varphi }}_0(\mathit{t}(\mathit{x}))$, we have

$$\begin{gathered} \mathit{\varphi }{(\mathit{x})}^{\top }\mathit{c}={\langle {\mathit{\varphi }}_0,{\mathit{K}}_0(\cdot ,\mathit{t}(\mathit{x}))\mathit{M}{(\mathit{x})}^{\top }\mathit{c}\rangle }_{{\mathcal{H}}_0} \\ ={\langle \mathit{\varphi },\mathit{M}(\cdot ){\mathit{K}}_0(\mathit{t}(\cdot ),\mathit{t}(\mathit{x}))\mathit{M}{(\mathit{x})}^{\top }\mathit{c}\rangle }_{\mathcal{H}} \\ ={\langle \mathit{\varphi },\mathit{K}(\cdot ,\mathit{x})\mathit{c}\rangle }_{\mathcal{H}}, \end{gathered}$$

where the second step follows ${\langle \mathit{\varphi },{\mathit{\varphi }}^{\prime }\rangle }_{\mathcal{H}}={\langle {\mathit{\varphi }}_0,{\mathit{\varphi }}_0^{\prime }\rangle }_{{\mathcal{H}}_0}\text{with}{\mathit{\varphi }}_0^{\prime }(\cdot )={\mathit{K}}_0(\cdot ,\mathit{t}(\mathit{x}))\mathit{M}{(\mathit{x})}^T\mathit{c}$.

Proof of Theorem 3

Proof.

Note that KL divergence is invariant under invertible variable transforms, that is,

$$\text{KL}(q_{[ϵ\mathit{\varphi }]}\Vert p)=\text{KL}(q_{[ϵ\mathit{\varphi }]0}\Vert p_0).$$ (20)

where p₀ denotes the distribution of x₀ = t(x) when x ~ p, and $q^{[ϵ\mathit{\varphi }]0}$ denotes the distribution of ${\mathit{x}}_0^{\prime }=\mathit{t}({\mathit{x}}^{\prime })\text{when}{\mathit{x}}^{\prime }~q_{[ϵ\mathit{\varphi }]}$. Recall that $q_{[ϵ\mathit{\varphi }]}$ is defined as the distribution of ${\mathit{x}}^{\prime }=\mathit{x}+ϵ\varphi (\mathit{x})$ when $\mathit{x}~q$.

Denote by t⁻¹ the inverse map of t, that is, ${\mathit{t}}^{-1}(\mathit{t}(\mathit{x}))=\mathit{x}$. We can see that ${\mathit{x}}_0^{\prime }~q_{[ϵ\mathit{\varphi }]0}$ can be obtained by

$$\begin{gathered} {\mathit{x}}_0^{\prime }=\mathit{t}\left({\mathit{x}}^{\prime }\right)//{\mathit{x}}^{\prime }~q_{[ϵ\mathit{\varphi }]} \\ =\mathit{t}(\mathit{x}+ϵ\mathit{\varphi }(\mathit{x}))//\mathit{x}~q \\ =\mathit{t}\left({\mathit{t}}^{-1}\left({\mathit{x}}_0\right)+ϵ\mathit{\varphi }\left({\mathit{t}}^{-1}\left({\mathit{x}}_0\right)\right)\right)//{\mathit{x}}_0~q_0 \\ ={\mathit{x}}_0+ϵ\nabla \mathit{t}\left({\mathit{t}}^{-1}\left({\mathit{x}}_0\right)\right)\mathit{\varphi }\left({\mathit{t}}^{-1}\left({\mathit{x}}_0\right)\right)+\mathcal{O}\left({ϵ}^2\right) \\ ={\mathit{x}}_0+ϵ{\mathit{\varphi }}_0\left({\mathit{x}}_0\right)+\mathcal{O}\left({ϵ}^2\right), \end{gathered}$$ (21)

where we used the definition that $\mathit{\varphi }(\mathit{x})=\nabla \mathit{t}{(\mathit{x})}^{-1}{\mathit{\varphi }}_0(\mathit{t}(\mathit{x}))$ in (11), and $\mathcal{O}(\cdot )$ is the big-O notation.

From Theorem 3.1 of Liu & Wang (2016), we have

$${\frac{d}{dϵ}\text{KL}(q_{[ϵ\mathit{\varphi }]}\Vert p)|}_{ϵ=0}=-{\mathbb{E}}_q[{\mathcal{P}}^{\top }\mathit{\varphi }].$$

Using Equation (21) and derivation similar to Theorem 3.1 of Liu & Wang (2016), we can show

$${\frac{d}{dϵ}\text{KL}(q_{[ϵ\mathit{\varphi }]0}\Vert p)|}_{ϵ=0}=-{\mathbb{E}}_{q0}[{\mathcal{P}}_0^{\top }{\mathit{\varphi }}_0].$$

Combining these with (20) proves (11).

Following Lemma 2, when ϕ₀ is in ${\mathcal{H}}_0$ with kernel ${\mathit{K}}_0\left(\mathit{x},{\mathit{x}}^{\prime }\right),\mathit{\varphi }$ is in $\mathcal{H}$ with kernel $\mathit{K}\left(\mathit{x},{\mathit{x}}^{\prime }\right)$. Therefore, maximizing ${\mathbb{E}}_q[{\mathcal{P}}^{\top }\mathit{\varphi }]$ in $\mathcal{H}$ is equivalent to ${\mathbb{E}}_{q_0}[{\mathcal{P}}_0^{\top }{\mathit{\varphi }}_0]$ in ${\mathcal{H}}_0$. This suggests the trajectory of SVGD on p₀ with K₀ and that on p with K are equivalent. □

B. Toy Examples

Figure 3 and Figure 4 show results of different algorithms on three 2D toy distributions: Star, Double banana and Sine. Detailed information of these distributions and more results are shown in Section B.1-B.3.

We can see from Figure 3–4 that both variants of matrix SVGD consistently outperform SVN and vanilla SVGD. We also find that Matrix SVGD(mixture) tends to outperform Matrix SVGD (average), which is expected since Matrix SVGD (average) uses a constant preconditioning matrix for all the particles, and can not capture different curvatures at different locations. Matrix SVGD (mixture) yields the best performance in general.

Figure 3:

The particles obtained by various methods at the 30/100/30-th iteration on three toy 2D distributions.

Figure 4:

The MMD vs. training iteration of different algorithms on the three toy distributions.

B.1. Sine

The density function of the “Sine” distribution is defined by

$$p\left(x_1,x_2\right)\propto \exp \left(\frac{-{(x_2+\sin (\alpha x_1))}^2}{2{\sigma }_1}-\frac{x_1^2+x_2^2}{2{\sigma }_2}\right),$$

where we choose σ = 1, σ₁ = 0:003, σ₂ = 1.

Figure 5:

The particles obtained by various methods on the toy Sine distribution.

B.2. Double Banana

We use the “double banana” distribution constructed in Detommaso et al. (2018), whose probability density function is

$$p(\mathit{x})\propto \exp \left(-\frac{\Vert \mathit{x}{\Vert }_2^2}{2{\sigma }_1}-\frac{{(y-F(\mathit{x}))}^2}{2{\sigma }_2}\right),$$

where $\mathit{x}=\left[x_1,x_2\right]\in {ℝ}^2$ and $F(\mathit{x})=\log ({(1-x_1)}^2+100{(x_2-x_1^2)}^2)$ and $y=\log (30)$, ${\sigma }_1=1.0$, ${\sigma }_2=0.09$.

Figure 6:

The particles obtained by various methods on the double banana distribution.

B.3. Star

We construct the “star” distribution with a Gaussian mixture model, whose density function is

$$p(\mathit{x})=\frac{1}{K}\sum _{i=1}^K\mathcal{N}\left(\mathit{x};{\mathit{\mu }}_i,{\mathbf{\Sigma }}_i\right),$$

with $\mathit{x}\in {ℝ}^2{\mu }_1=[0;1.5]$, ${\mathbf{\Sigma }}_1=\mathrm{diag}\left(\left[1;\frac{1}{100}\right]\right)$, and the other means and covariance matrices are defined by rotating their previous mean and covariance matrix. To be precise,

$${\mathit{\mu }}_{i+1}=\mathit{U}{\mathit{\mu }}_i,{\mathbf{\Sigma }}_{i+1}=\mathit{U}{\mathbf{\Sigma }}_i{\mathit{U}}^{\top },\mathit{U}=\left[\begin{array}{cc}\cos (\theta )& \sin (\theta )\\ -\sin (\theta )& \cos (\theta )\end{array}\right],$$

with angle $\theta =\frac{2\pi }{K}$. We set the number of component K to be 5.

Figure 7:

The particles obtained by various methods on the star-shaped distribution.