Intro

这一工作是清华大学liu qiang老师提出的，相关论文从2016年开始也一直在更新，分别发表在NIPS、ICLR等顶会上。
Stein Variational Gradient Descent as Gradient Flow
Stein Variational Gradient Descent as Moment Matching
Stein Variational Gradient Descent: A General Purpose Bayesian Inference Algorithm

从定义上来说，SVGD是一种确定性的采样算法，用一组粒子来近似给定的分布。基于这两点，它和MCMC以及VI都有共通之处。但SVGD即保证了在大量数据下的计算速度，也比变分推断具有更高的准确性。

从整体上来看，这一工作通过引入Stein discrepancy来度量两个分布之间的距离，再借助RKHS使其容易计算，最后借助gradient descent进行优化。因此下文也就从这三部分一一介绍。

Background

Stein's method

首先我们需要引入几个定义：

Stein score function \[ \boldsymbol{s}_{p}=\nabla_{x} \log p(x)=\frac{\nabla_{x} p(x)}{p(x)} \] 这一函数被称为$q(x)$的Stein score function
Stein class
当函数$f: \mathcal{X} \rightarrow \mathbb{R}$满足下式时则称其在stein class中 \[ \int_{x \in \mathcal{X}} \nabla_{x}(f(x) p(x)) d x=0 \] 其中$\mathcal{X}$ 是$\mathbb{R}^{d}$下的子集，而$p(x)$则是在$\mathcal{X}$ 下连续可微的分布。
Stein's operator：作用在$p$上的线性操作 \[ \mathcal{A}_{p} f(x)=\boldsymbol{s}_{p}(x) f(x)+\nabla_{x} f(x) \] 其中$s_{p}$和$\mathcal{A}_{p} f$ 都是$d \times 1$ 函数(mapping from $\mathcal{X}$ to $\mathbb{R}^{d}.$)

有了以上三个定义后，我们可以尝试得到stein discrepancy。首先作为一个度量手段，必然需要满足一些条件。
当且仅当 \[ \mathbb{E}_{p}\left[\boldsymbol{s}_{q}(x) f(x)+\nabla_{x} f(x)\right]=0 \qquad (1) \] $p(x)$ 和 $q(x)$是相等的。而当两个分布$p=q$时又被称为stein identity。
借助 (1) 式，我们可以定义Stein discrepancy来度量两个分布$p$和$q$之间的差异： \[ \mathbb{S}(p, q)=\max _{f \in \mathcal{F}}\left(\mathbb{E}_{p}\left[\boldsymbol{s}_{q}(x) f(x)+\nabla_{x} f(x)\right]\right)^{2} \] 借助之前定义的stein operator，也可以把上式写为 \[\mathbb{S}(p, q)=\max _{f \in \mathcal{F}}\left(\mathbb{E}_{p}\left[\mathcal{A}_{q} f(x)\right]\right)^{2}\]

$\mathcal{F}$是一系列连续可微的且满足$\mathbb{S}(p, q)$不为0($p \neq q$时)函数集合。当$p \neq q$时，$\mathbb{S}(p,q)>0$，而$max$则是因为我们希望距离尽可能明显。

$\mathbb{S}(p, q)$并没有被广泛应用在机器学习中，因为其计算和优化的复杂性: $q(x)=f(x) / Z$ 而$Z=\int f(x) d x$的计算往往设计高维积分。
但是论文提出了将函数$\mathcal{F}$用核函数代替时，会得到易于计算的Stein discrepancy $\mathbb{S}(p, q)$。具体而言，我们令$\mathcal{F}$来源于希尔伯特再生核空间的一个球 (reproducing kernel Hilbert space (RKHS))。

Kernelized Stein Discrepancy

对于映射后的函数，对应的正定核$k\left(x, x^{\prime}\right)$，我们有 \[ \mathbb{S}(p, q)=\mathbb{E}_{x, x^{\prime} \sim p}\left[u_{q}\left(x, x^{\prime}\right)\right] \] 其中$x, x^{\prime}$是$p$中独立同分布的两个变量，函数$u_{q}\left(x, x^{\prime}\right)$由$q$确定，如果展开的话实际上是： \[u_{q}\left(x, x^{\prime}\right)= \boldsymbol{s}_{q}(x)^{\top} k\left(x, x^{\prime}\right) \boldsymbol{s}_{q}\left(x^{\prime}\right)+\boldsymbol{s}_{q}(x)^{\top} \nabla_{x^{\prime}} k\left(x, x^{\prime}\right)+\nabla_{x} k\left(x, x^{\prime}\right)^{\top} \boldsymbol{s}_{q}\left(x^{\prime}\right)+\operatorname{trace}\left(\nabla_{x, x^{\prime}} k\left(x, x^{\prime}\right)\right)\]

当我们从未知分布$p(x)$采样出一个样本${x_i}$时，我们可以进行近似计算 \[\hat{\mathbb{S}}(p, q)=\frac{1}{n(n-1)} \sum_{i \neq j} u_{q}\left(x_{i}, x_{j}\right)\]

接下来我们详细介绍上述的过程。

Kernels and Reproducing Kernel Hilbert Spaces

Kernel and Hilbert Spaces介绍 (未完待续)

令$k\left(x, x^{\prime}\right)$为一个正定核，根据Mercer’s theorem我们对其进行谱分解： \[ k\left(x, x^{\prime}\right)=\sum_{j} \lambda_{j} e_{j}(x) e_{j}\left(x^{\prime}\right) \] 其中$\left\{e_{j}\right\},\left\{\lambda_{j}\right\}$分别是正交特征函数和正特征值，满足$\int e_{i}(x) e_{j}(x) d x=\mathbb{I}[i=j]$, for $\forall i, j$

对于一个正定核，它可以分解为RKHS中特征函数的线性组合（空间中任何一个函数可以用这组基的线性组合来表示）。由一个特定的核函数能产生一个唯一的Hilbert空间，有性质 \[ f(x)=\langle f, k(\cdot, x)\rangle_{\mathcal{H}}, \quad k\left(x, x^{\prime}\right)=\left\langle k(\cdot, x), k\left(\cdot, x^{\prime}\right)\right\rangle_{\mathcal{H}} \] 当我们定义 $\mathcal{H}^{d}=\mathcal{H} \times \mathcal{H} \times \cdots \mathcal{H}$ 为 $d$ 维向量函数 $\mathbf{f}=\left\{f_{i}: f_{i} \in \mathcal{H} \quad i=1, \cdots, d\right\}$ 组成的 Hilbert空间, $\mathcal{H}^{d}$ 上的内积定义为 $<\mathbf{f}, \mathbf{g}>_{\mathcal{H}^{d}}=\sum_{i=1}^{d}<f_{i}, g_{i}>_{\mathcal{H}}$ 。如果觉得上述的介绍太过抽象，可以看附录部分关于RKHS的一个toy example

lemmas

Stein's Identity ： \[ \mathbb{E}_{p}\left[\mathcal{A}_{p} \boldsymbol{f}(x)\right]=\mathbb{E}_{p}\left[\boldsymbol{s}_{p}(x) \boldsymbol{f}(x)^{\top}+\nabla \boldsymbol{f}(x)\right]=0 \] 证明的话根据$\boldsymbol{s}_{p}(x) \boldsymbol{f}(x)^{\top}+\nabla \boldsymbol{f}(x)=\nabla_{x}(\boldsymbol{f}(x) p(x)) / p(x)$和分布积分法则就可以推导出来。

有了上面的引理，我们可以得到 \[\mathbb{E}_{p}\left[\mathcal{A}_{q} \boldsymbol{f}(x)\right]=\mathbb{E}_{p}\left[\mathcal{A}_{q} \boldsymbol{f}(x)-\mathcal{A}_{p} \boldsymbol{f}(x)\right]=\mathbb{E}_{p}\left[\left(\boldsymbol{s}_{q}(x)-\boldsymbol{s}_{p}(x)\right) \boldsymbol{f}(x)^{\top}\right]\]

也就是说$\mathbb{E}_{p}\left[\mathcal{A}_{q} \boldsymbol{f}(x)\right]$是一个由$f(x)$加权的期望，对于$\left(s_{q}(x)-s_{p}(x)\right)$的期望。

KSD

借助上面的推导以及RKHS的性质，我们可以定义出核空间下的stein discrepancy： \[ S(p, q)=\mathbb{E}_{\mathbf{x}, \mathbf{y} \sim p}\left[\left(s_{q}(\mathbf{x})-s_{p}(\mathbf{x})\right)^{T} k(\mathbf{x}, \mathbf{y})\left(s_{q}(\mathbf{x})-s_{p}(\mathbf{x})\right)\right] \] （这里我们将Kernelized前后的stein discrepancy分别写为$\mathbb{S}(p,q)$和$S(p, q)$。另外需要注意后续部分推导是针对stein discrepancy中的期望项）
我们要用一个可采样的分布$q$拟合分布$p$，因此我们希望$S(p, q)$[]中的式子是与$p$无关的。把式子展开后可以发现 \[ \begin{aligned} S(p, q) &=\mathbb{E}_{\mathbf{x}, \mathbf{y} \sim}\left[\left(s_{q}-s_{p}\right)^{T} k(\mathbf{x}, \mathbf{y})\left(s_{q}-s_{p}\right)\right] \\ &=\mathbb{E}_{\mathbf{x}, \mathbf{y} \sim}\left[\left(s_{q}-s_{p}\right)^{T}\left(k(\mathbf{x}, \mathbf{y}) s_{q}+\nabla_{y} k(\mathbf{x}, \mathbf{y})-k(\mathbf{x}, \mathbf{y}) s_{p}-\nabla_{\mathbf{y}} k(\mathbf{x}, \mathbf{y})\right)\right] \\ &=\mathbb{E}_{\mathbf{x}, \mathbf{y} \sim}\left[\left(s_{q}-s_{p}\right)^{T} v(\mathbf{x}, \mathbf{y})\right] \end{aligned} \] 其中$v(\mathbf{x}, \mathbf{y})=k(\mathbf{x}, \mathbf{y}) s_{q}(\mathbf{y})+\nabla_{\mathbf{y}} k(\mathbf{x}, \mathbf{y})=\mathcal{A}_{q} k_{\mathbf{x}}(\mathbf{y}), k_{\mathbf{x}}(\cdot)=k(\mathbf{x}, \cdot)$

而对于固定的 $\mathbf{y}$, 容易证明 $v(\cdot, \mathbf{y})$ 是Stein class of $p(\mathbf{x})$, 即满足 $\int_{\mathbf{x} \in \mathcal{X}} \nabla_{\mathbf{x}}(v(\mathbf{x}, \mathbf{y}) p(\mathbf{x})) d \mathbf{x}=0$ 。因此就可以进一步将上式写开为 \[ \begin{aligned} S(p, q) &=\mathbb{E}_{\mathbf{x}, \mathbf{y} \sim p}\left[s_{q}^{T} v(\mathbf{x}, \mathbf{y})-\left(\nabla_{\mathbf{x}} \ln p(\mathbf{x})\right)^{T} v(\mathbf{x}, \mathbf{y})\right] \\ &=\mathbb{E}_{\mathbf{x}, \mathbf{y} \sim p}\left[s_{q}^{T} v(\mathbf{x}, \mathbf{y})\right]-\int d \mathbf{x} d \mathbf{y} p(\mathbf{x}) p(\mathbf{y})\left(\nabla_{\mathbf{x}} \ln p(\mathbf{x})\right)^{T} v(\mathbf{x}, \mathbf{y}) \\ &=\mathbb{E}_{\mathbf{x}, \mathbf{y} \sim p}\left[s_{q}^{T} v(\mathbf{x}, \mathbf{y})\right]+\mathbb{E}_{\mathbf{x}, \mathbf{y} \sim p}\left[\operatorname{tr} \nabla_{\mathbf{x}} v(\mathbf{x}, \mathbf{y})\right] \end{aligned} \]

其中$\nabla_{\mathbf{x}} v(\mathbf{x}, \mathbf{y})=\nabla_{\mathbf{x}} k(\mathbf{x}, \mathbf{y}) s_{q}(\mathbf{y})^{T}+\nabla_{\mathbf{x}} \nabla_{\mathbf{y}} k(\mathbf{x}, \mathbf{y})$为一个矩阵。

这样就得到了前文提到了 \[ S(p, q)=\mathbb{E}_{\mathbf{x}, \mathbf{y} \sim p}\left[u_{q}(\mathbf{x}, \mathbf{y})\right] \] 其中$u_{q}(\mathbf{x}, \mathbf{y})$仅与分布$q$有关。

上述的推导说明了什么？说明引入核方法的可行性。而另一方面，我们要用核方法来加速内积的计算，如何体现？
借助RKHS的对称性和再生性，我们可以将$S(p,q)$进行变化： \[ \begin{aligned} S(p, q) &=\mathbb{E}_{\mathbf{x}, \mathbf{y} \sim p}\left[\left(s_{q}(\mathbf{x})-s_{p}(\mathbf{y})\right)^{T} k(\mathbf{x}, \mathbf{y})\left(s_{q}(\mathbf{x})-s_{p}(\mathbf{y})\right)\right] \\ &=\mathbb{E}_{\mathbf{x}, \mathbf{y} \sim p}\left[\left(s_{q}(\mathbf{x})-s_{p}(\mathbf{y})\right)^{T}<k(\mathbf{x}, \cdot), k(\cdot, \mathbf{y})>_{\mathcal{H}}\left(s_{q}(\mathbf{x})-s_{p}(\mathbf{y})\right)\right] \\ &=\sum_{i=1}^{d}<\mathbb{E}_{\mathbf{x} \sim p}\left[\left(s_{q}^{i}(\mathbf{x})-s_{p}^{i}(\mathbf{x})\right) k(\mathbf{x}, \cdot)\right], \mathbb{E}_{\mathbf{y} \sim p}\left[\left(s_{q}^{i}(\mathbf{y})-s_{p}^{i}(\mathbf{y})\right) k(\cdot \mathbf{y})\right]>_{\mathcal{H}} \\ &=\sum_{i=1}^{d}<\beta_{i}, \beta_{i}>_{\mathcal{H}} \\ &=\|\boldsymbol{\beta}\|_{\mathcal{H}^{d}}^{2} \end{aligned} \] 其中$\boldsymbol{\beta}(\mathbf{y})$是一个向量函数 \[ \boldsymbol{\beta}(\mathbf{y})=\mathbb{E}_{\mathbf{x} \sim p}\left[\mathcal{A}_{q} k_{\mathbf{y}}(\mathbf{x})\right]=\mathbb{E}_{\mathbf{x} \sim p}\left[\left(s_{q}(\mathbf{x})-s_{p}(\mathbf{x})\right) k_{\mathbf{y}}(\mathbf{x})\right] \]

在这里再展示一下最开始的stein discrepancy \[ \mathbb{S}(p, q)=\max _{f \in \mathcal{F}}\left(\mathbb{E}_{p}\left[\mathcal{A}_{q} f(x)\right]\right)^{2} \] 对于任意向量 $\mathbf{f} \in \mathcal{H}^{d}$ 与 $\boldsymbol{\beta}$ 的内积为 \[ \begin{aligned} <\mathbf{f}, \boldsymbol{\beta}>_{\mathcal{H}^{d}} &=\sum_{i=1}^{d}<f_{i}, \mathbb{E}_{\mathbf{x} \sim p}\left[s_{q}^{i}(\mathbf{x}) k(\mathbf{x}, \cdot)+\nabla_{x_{i}} k(\mathbf{x}, \cdot)\right]>_{\mathcal{H}} \\ &=\sum_{i=1}^{d} \mathbb{E}_{\mathbf{x} \sim p}\left[s_{q}^{i}(\mathbf{x})<f_{i}, k(\mathbf{x}, \cdot)>_{\mathcal{H}}+<f_{i}, \nabla_{x_{i}} k(\mathbf{x}, \cdot)>_{\mathcal{H}}\right] \\ &=\sum_{i=1}^{d} \mathbb{E}_{\mathbf{x} \sim p}\left[s_{q}^{i}(\mathbf{x}) f_{i}(\mathbf{x})+\nabla_{x_{i}} f_{i}(\mathbf{x})\right] \\ &=\mathbb{E}_{\mathbf{x} \sim p}\left[\operatorname{tr}\left(\mathcal{A}_{q} \mathbf{f}(\mathbf{x})\right)\right] \leq\|\boldsymbol{\beta}\|_{\mathcal{H}^{d}} (因为任意两个向量内积小于它与自身的内积) \end{aligned} \] 因此我们有 \[ \|\boldsymbol{\beta}\|_{\mathcal{H}^{d}}=S(p, q)=\max _{f \in \mathcal{H}^{d}}\left\{\mathbb{E}_{\mathbf{x} \sim p}\left[\operatorname{tr}\left(\mathcal{A}_{q} f(\mathbf{x})\right)\right]\right \}. \] 上式中$\|f\|_{\mathcal{H}^{d}} \leq 1$，对应于最初提到的映射到希尔伯特空间中的一个球中的最优向量。这个$S(p,q)$最大值所对应的向量为 ${f}^{*}=\boldsymbol{\beta} /\|\boldsymbol{\beta}\|_{\mathcal{H}^{d}}$

简单来说，引入kernel之后，我们可以直接得到stein discrepancy定义式中的函数。也就是说，KSD非常容易就能求得。

Stein Variational Gradient Descent (SVGD)

SVGD的另一个核心公式在于 \[ \left.\nabla_{\epsilon} \mathrm{KL}\left(q_{[T]} \| p\right)\right|_{\epsilon=0}=-\mathbb{E}_{x \sim q}\left[\operatorname{tr}\left(\mathcal{A}_{p} \phi(x)\right)\right] \] 也就是说KL散度变分求导等于KSD（具体推导过程见附录），这意味着KL散度变化最快的方向就是KSD所对应的向量函数 $\phi^{*}=\boldsymbol{\beta} /\|\boldsymbol{\beta}\|_{\mathcal{H}^{d}}$ 算法伪代码具体而言粒子$\left\{x_{i}^{l}\right\}_{i=1}^{n}$ 表示第$l$次达代的第$i$个粒子, 一共$n$ 个。粒子最开始是从分布$q_{0}$中采样的, 最初的分布$q$可以是任意的。也就是说，该算法不依赖于初始的分布。

算法中的更新项包含了两个部分 \[ k\left(x_{j}^{\ell}, x\right) \nabla_{x_{j}^{\ell}} \log p\left(x_{j}^{\ell}\right)+\nabla_{x_{j}^{\ell}} k\left(x_{j}^{\ell}, x\right) \] 其中第一项意味着粒子会朝$p$分布概率高的地方移动，而第二项代表着粒子将会朝着远离当前迭代轮数$l ll的粒子，从而减轻局部最优的风险。

SVGD拟合一维分布

回顾与总结

从上文繁杂的推导中，SVGD算法确保粒子的移动是朝着KL散度的减小最快方向，而这个方向可以有核化的stein discrepancy导出我们回看一下KL divergence的定义式： \[\mathrm{KL}(P \| Q)=\int P(x) \log \frac{P(x)}{Q(x)} d x\]

对于目标分布$p(x)$，变分推断 (VI)目标是从一类分布族$\mathcal{Q}$中找到最优的$q(x)$ \[ q^{*}=\underset{q \in \mathcal{Q}}{\arg \min }\left\{K L(q \| p)=\mathbb{E}_{q}[\log q(x)]-\mathbb{E}_{q}[\log \bar{p}(x)]+\log Z\right\} \] 而SVGD在再生核希尔伯特空间下给出了使得KL散度下降最快的确定性方向，类似经典的梯度下降算法，可以理解为迭代构建增量变化的方法。

Appendix

The reproducing kernel Hilbert space

先回顾一下kernel的定义： Let $\mathcal{X}$ be a non-empty set. A function $k: \mathcal{X} \times \mathcal{X} \rightarrow \mathbb{R}$ is called a kernel if there exists an $\mathbb{R}$-Hilbert space and a map $\phi: \mathcal{X} \rightarrow \mathcal{H}$ such that $\forall x, x^{\prime} \in \mathcal{X}$ \[ k\left(x, x^{\prime}\right):=\left\langle\phi(x), \phi\left(x^{\prime}\right)\right\rangle_{\mathcal{H}} \]

在此基础上，我们用一个异或问题的例子来介绍RKHS。考虑特征映射

\[\phi: \mathbb{R}^{2} \rightarrow \mathbb{R}^{3}\] \[x=\left[\begin{array}{l} x_{1} \\ x_{2} \end{array}\right] \quad \mapsto \quad \phi(x)=\left[\begin{array}{c} x_{1} \\ x_{2} \\ x_{1} x_{2} \end{array}\right]\] 特征空间和特征映射。希尔伯特空间的元素一般是函数，而函数可以被视为无穷维的向量。因此事实上希尔伯特空间的基底是一组无限维的函数，可以参考傅立叶变化或泰勒展开

kernel \[ k(x, y)=\left[\begin{array}{c} x_{1} \\ x_{2} \\ x_{1} x_{2} \end{array}\right]^{\top}\left[\begin{array}{c} y_{1} \\ y_{2} \\ y_{1} y_{2} \end{array}\right] \]

接下来我们可以定义一个特征函数： \[ f(x)=a x_{1}+b x_{2}+c x_{1} x_{2} \] 这个函数属于从$\mathcal{X}=\mathbb{R}^{2}$映射到$\mathbb{R}$的函数空间。此时，我们也可以把函数$f$等价表示为： \[ f(\cdot)=\left[\begin{array}{l} a \\ b \\ c \end{array}\right] \] 至此，我们可以把$f(x)$写为： \[ \begin{aligned} f(x) &=f(\cdot)^{\top} \phi(x) \\ &:=\langle f(\cdot), \phi(x)\rangle_{\mathcal{H}} \end{aligned} \] 也就是说，特征函数$f$在$x$的值可以被写为特征空间中的内积。$\mathcal{H}$是一个将$\mathbb{R}^{2}$映射到$\mathbb{R}$的函数空间。上面这些乱七八糟的怎么体现再生性呢？我们仔细看下面的等式 \[ k(\cdot, y)=\left[\begin{array}{c} y_{1} \\ y_{2} \\ y_{1} y_{2} \end{array}\right]=\phi(y) \] 上式我们参考$f(\cdot)$类似的定义。具体来说，如果我们令$a=y_{1}, b=y_{2}$, and $c=y_{1} y_{2}$，就有 \[ \langle k(\cdot, y), \phi(x)\rangle_{\mathcal{H}}=a x_{1}+b x_{2}+c x_{1} x_{2} \]

总的来说RKHS两个特性：
每个点的特征映射在特征空间中 \[ \forall x \in \mathcal{X}, \quad k(\cdot, x) \in \mathcal{H} \] 再生性：
\[ \forall x \in \mathcal{X}, \forall f \in \mathcal{H},\langle f, k(\cdot, x)\rangle_{\mathcal{H}}=f(x) \] \[ k(x, y)=\langle k(\cdot, x), k(\cdot, y)\rangle_{\mathcal{H}} \]

Stein variational gradient descent (NIPS2018)