Stein variational gradient descent（SVGD）

最近读了Liu Qiang组关于SVGD的几篇文章，和老师讨论了很久、查阅了很多文献后大概了解了它的思路。这篇文章整合一下我所理解的SVGD。由于太笨，不知道Katex怎么像Latex一样写公式标号····所以本文所有公式都没有标号Orz

近似推断被广泛用于概率机器学习与统计中，Stein variational gradient descent （SVGD）是由Liu Qiang等提出的一种近似推断算法。不同于MCMC，它是一种确定性的算法。不同于变分推断（VI），它采用粒子方法直接对目标概率分布进行逼近。

1.Stein’s Identity

首先Liu Qiang组关于SVGD的idea出发点都是Stein’s Identity。Stein’s Identity讲了这么一件事情：

$${\mathbb{E}}_p\left[{\nabla }_x\log q(x)f(x)+{\nabla }_{x}f(x)\right]=0ifp(x)=q(x)$$

我们考虑 ${\mathbb{R}}^d$上两个光滑分布 $p$和 $q$，上式成立的充分条件是 $p(x)=q(x)$. 这一等式的成立非常好证，只需要将积分展开用分部积分法就可以证明($\int u(x)v^{\prime }(x)dx=u(x)v(x)-\int u^{\prime }(x)v(x)dx$).

那么，当 $p(x)\ne q(x)$时，上式是不等于0的，这就有些像一个距离度量，我们进一步想到上式是不是也可以像KL divergence一样衡量两个分布之间的距离呢？并且通过优化这个距离做一些事情呢？这就是Liu Qiang组的研究工作。

2.向量情况下的距离度量

根据Stein’s identity我们可以定义一种度量来衡量分布p和q之间的距离，首先考虑向量情况下的推导，即$f(x)$属于${\mathbb{R}}^{1\times 1}$

我们定义算子 ${\mathcal{A}}_{p}f(x)=s_p(x)f(x)+{\nabla }_{x}f(x)$，以及 ${\mathbf{s}}_q={\nabla }_x\log q(x)$，之前的Stein’s identity就可以写成：
$${\mathbb{E}}_{x\sim p}\left[{\mathcal{A}}_{q}f(x)\right]=0where(q=p)$$

我们定义距离 $S(p,q)={\max }_{f\in \mathcal{F}}{\left({\mathbb{E}}_p\left[{\mathcal{A}}_{q}f(x)\right]\right)}^2$

这里$\mathcal{F}$是光滑函数的集合，但是从$\mathcal{F}$中找到$f$不是一件容易的事情，所以我们还需要别的trick（也就是kernel）。关于这里为什么是max，因为我们希望找到的度量，一定是在$p(x)\ne q(x)$时让这个距离值尽可能的大，这样才能在同一种度量下分的更开。

3.矩阵情况下的距离度量

下面我们尝试将${\mathcal{A}}_{p}f(x)$推广到矩阵的形式，也就是$f(x)$推广到${\mathbb{R}}^{d\times 1}$的形式。

这里为了区分我们把$f(x)$属于${\mathbb{R}}^{d\times 1}$记作$\varphi (x)$，令$\varphi (x)={\left[{\varphi }_1(x),\cdots ,{\varphi }_d(x)\right]}^{\top }$. 也就是说$\varphi (x)$代表着一组光滑函数，是一个向量值函数。

此时， ${\mathcal{A}}_p\varphi (x)=s_p(x){\varphi }^{\top }(x)+{\nabla }_x\varphi (x)$

Stein’s identity可以写成：
$${\mathbb{E}}_{x\sim p}\left[{\mathcal{A}}_q\varphi (x)\right]={\mathbb{E}}_{x\sim p}\left[{\mathcal{A}}_q\varphi (x)-{\mathcal{A}}_p\varphi (x)\right]={\mathbb{E}}_{x\sim p}\left[(s_q(x)-s_p(x)){\varphi }^{\top }(x)\right]$$

由于$s_q(x)$属于${\mathbb{R}}^{d\times 1}$，${\varphi }^{\top }(x)$属于${\mathbb{R}}^{1\times d}$，上式应该是一个$d\times d$的矩阵。但是我们想要的是距离度量，矩阵的期望是没有任何意义的，所以我们可以取标量为：

$${\mathbb{E}}_p\left[\mathrm{trace}\left({\mathcal{A}}_q\varphi (x)\right)\right]={\mathbb{E}}_p\left[{\left({\mathbf{s}}_q(x)-{\mathbf{s}}_p(x)\right)}^{\top }\varphi (x)\right]$$

此时，$s_q^{\top }(x)$属于${\mathbb{R}}^{1\times d}$，$\varphi (x)$属于${\mathbb{R}}^{d\times 1}$，我们得到了一个标量，这是我们所希望的距离值。同时，也可以很直接地看到我们取标量后的值其实就是原矩阵形式的迹。

总结一下，到现在为止，我们将Stein’s identity从向量形式拓展到了矩阵形式，在矩阵形式下，我们依旧可以将$S(p,q)$写成一个标量，为之后kernel的引入做好了铺垫。

4.Kernel的引入

4.1 RKHS

在引入kernel到Stein’s identity之前，我们先回顾一下RKHS(reproducing kernel hilbert space)。

每一个函数都可以看做一个无限维的向量，那么二元函数$K(x,y)$就可以看做是一个无限维的矩阵。当$K(x,y)$满足正定性和对称性的时候我们把它叫做核函数。与矩阵特征值和特征向量类似，核函数存在特征值和特征函数 (将函数看做无限维向量)。也就是

$$\int K(\mathrm{x},\mathrm{y})\psi (\mathrm{x})d\mathrm{x}=\lambda \psi (\mathrm{y})$$

那么根据Mercer定理，我们有
$$K(\mathrm{x},\mathrm{y})=\sum _{i=0}^{\infty }{\lambda }_i{\psi }_i(\mathrm{x}){\psi }_i(\mathrm{y})$$
这里一个核函数对应无穷个特征值${\left\{{\lambda }_i\right\}}_{i=1}^{\infty }$和无穷个特征方程${\left\{{\psi }_i\right\}}_{i=1}^{\infty }$
将${\left\{\sqrt{{\lambda }_i}{\psi }_i\right\}}_{i=1}^{\infty }$作为一组正交基构建一个希尔伯特空间 $\mathcal{H}$ 。这个空间中的任何一个函数（向量）都可以表示为这组基的线性组合。如$f={\left(f_1,f_2,\dots \right)}_{\mathcal{H}}^T$

用$K(\mathrm{x},\cdot )$表示固定核函数的一个参数为$x$，即矩阵第$x$行的一元函数或无限维向量。那么，我们有
$$\begin{array}{c}K(\mathrm{x},\cdot )={\left(\sqrt{{\lambda }_1}{\psi }_1(\mathrm{x}),\sqrt{{\lambda }_2}{\psi }_2(\mathrm{x}),\dots \right)}_{\mathcal{H}}^T\\ K(\mathrm{y},\cdot )={\left(\sqrt{{\lambda }_1}{\psi }_1(\mathrm{y}),\sqrt{{\lambda }_2}{\psi }_2(\mathrm{y}),\dots \right)}_{\mathcal{H}}^T\\ <K(\mathbf{x},\cdot ),K(\mathbf{y},\cdot ){>}_{\mathcal{H}}={\sum }_{i=0}^{\infty }{\lambda }_i{\psi }_i(\mathbf{x}){\psi }_i(\mathbf{y})=K(\mathbf{x},\mathbf{y})\end{array}$$

这就是RKHS，核心思想就是核函数的分解与重构。

4.2 Kernelized Stein Discrepancy（KSD）

现在，在RKHS的基础上，我们将kernel引入Stein’s identity

定义距离kernelized Stein discrepancy（KSD）:
$$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]$$

此时的$S(p,q)$本质上就是上一节矩阵情况中最后式子的变换，$k(\mathbf{x},\mathbf{y})\left(s_q(\mathbf{x})-s_p(\mathbf{x})\right)$可以看成$\varphi (x)$，这一点从维度分析上易知。同时，我们还可以换一种理解方式，$\varphi (x)$原来是一组函数，我们要找这组函数，那么现在变成了$k(\mathbf{x},\mathbf{y})\left(s_q(\mathbf{x})-s_p(\mathbf{x})\right)$，我们同样要找的是$x$固定时的一组列向量函数。

可行性证明

但是，这个时候出现了一个问题，我们原来的$S(p,q)={\max }_{f\in \mathcal{F}}{\left({\mathbb{E}}_p\left[{\mathcal{A}}_{q}f(x)\right]\right)}^2$中${\mathbb{E}}_p$内的式子只与$q$有关，而引入了kernel后，$S(p,q)$中${\mathbb{E}}_p$内的式子与$q,p$都有关了，$p$是我们不希望出现的，这会导致不可解。根据变换证明，这个问题是不存在的，引入kernel后的${\mathbb{E}}_p$中同样可以只与$q$有关，证明如下：

$$\begin{array}{rl}S(p,q)& ={\mathbb{E}}_{\mathbf{x},\mathbf{y}\sim p}\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 p}\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 p}\left[{\left(s_q-s_p\right)}^{T}v(\mathbf{x},\mathbf{y})\right]\end{array}$$

这里$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 )$

为了让$p$消失，将$s_p$带入
$$\begin{array}{rl}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[\mathrm{tr}{\nabla }_{\mathbf{x}}v(\mathbf{x},\mathbf{y})\right]\\ & ={\mathbb{E}}_{\mathbf{x},\mathbf{y}\sim p}\left[u_q(\mathbf{x},\mathbf{y})\right]\end{array}$$

这里$u_q(\mathbf{x},\mathbf{y})=s_q(\mathbf{x}{)}^{T}k(\mathbf{x},\mathbf{y})s_q(\mathbf{y})+s_q(\mathbf{x}{)}^T{\nabla }_{\mathbf{y}}k(\mathbf{x},\mathbf{y})+{\left({\nabla }_{\mathbf{x}}k(\mathbf{x},\mathbf{y})\right)}^T{s}_q(\mathbf{y})+\mathrm{tr}\left({\nabla }_{\mathbf{x}}{\nabla }_{\mathbf{y}}k(\mathbf{x},\mathbf{y})\right)$

虽然很长Orz但是只和q有关了ao。到现在为止，我们证明了引入kernel不影响求解，引入kernel后与引入前的原式形式等价。

易于求解证明

引入kernel自然是有引入的好处，好处在刚开始就已经说过了因为易与求解。光滑函数函数集$\mathcal{F}$中直接找使距离最大的$f$是很难做到的，但是引入kernel后我们就很容易找到了。我们下面来看为什么容易找到。

令$\mathbf{\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]$

根据RKHS的分析，我们可以将KSD写成
$$\begin{array}{rl}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}}\\ & =\Vert \mathbf{\beta }{\Vert }_{{\mathcal{H}}^d}^2\end{array}$$

也就是说，我们可以将$S(p,q)$看成希尔伯特空间中$\beta$函数的内积了。此时易知，任意的希尔伯特空间中的归一化（归一化目的：防止max导致无穷大）函数$\varphi$与$\beta$的内积都小于$\beta$与自己的内积。那么，我们想找的max，就可以写成

$$\Vert \mathbf{\beta }{\Vert }_{{\mathcal{H}}^d}=S(p,q)=\underset{\varphi \in {\mathcal{H}}^d}{\max }\left\{{\mathbb{E}}_{\mathbf{x}\sim p}\left[\mathrm{tr}\left({\mathcal{A}}_q\varphi (\mathbf{x})\right)\right],\text{ s.t. }\Vert \varphi {\Vert }_{{\mathcal{H}}^d}\le 1\right\}$$

此时最大值${\varphi }^{\ast }=\mathbf{\beta }/\Vert \mathbf{\beta }{\Vert }_{{\mathcal{H}}^d}$

这也就是说我们用kernel可以直接导出想找的$\varphi$

5.SVGD算法

那么有了KSD之后，我们能做什么呢？一个简单的想法是用它来优化KL散度，导师说这个式子${{\nabla }_{ϵ}\mathrm{K}\mathrm{L}\left(q_{[T]}\Vert p\right)|}_{ϵ=0}=-{\mathbb{E}}_{x\sim q}\left[\mathrm{trace}\left({\mathcal{A}}_p\varphi (x)\right)\right]$在learning领域是很常见的，也就是说KL散度变分求导是等于KSD的。所以我其实在想是不是这些论文都是根据KL散度反推回去的…其实它们都是拿结果往回推起源…然后再填充证明…

5.1 KL divergence的联系

下面我们来证明${{\nabla }_{ϵ}\mathrm{K}\mathrm{L}\left(q_{[T]}\Vert p\right)|}_{ϵ=0}=-{\mathbb{E}}_{x\sim q}\left[\mathrm{trace}\left({\mathcal{A}}_p\varphi (x)\right)\right]$

我们定义$z=T_{ϵ}(x)$,$T_{ϵ}(x)=T(x)$,$T(x)=x+ϵ\varphi (x)$，当$x\sim q$时，$q_{[T]}$表示z的概率密度。

同样地，取反函数，当$x\sim p$时，$p_{[T^{-1}]}$表示$z=T^{-1}(x)$的概率密度。

也可以理解成用微小扰动$ϵ\varphi (x)$来让$q$逼近$p$

我们有
$$\begin{array}{c}\mathrm{K}\mathrm{L}\left(q_{[T]}\Vert p\right)=\mathrm{K}\mathrm{L}\left(q\Vert p_{\left[T^{-1}\right]}\right)\\ {\nabla }_{ϵ}\mathrm{K}\mathrm{L}\left(q_{[T]}\Vert p\right)=-{\mathbb{E}}_{x\sim q}\left[{\nabla }_{ϵ}\log p_{\left[T^{-1}\right]}(x)\right]\end{array}$$

将${\nabla }_{ϵ}\log p_{\left[T^{-1}\right]}(x)$展开得到
$${\nabla }_{ϵ}\log p_{\left[T^{-1}\right]}(x)=s_p(\mathbf{T}(x){)}^{\top }{\nabla }_{ϵ}\mathbf{T}(x)+\mathrm{trace}\left({\left({\nabla }_x\mathbf{T}(x)\right)}^{-1}\cdot {\nabla }_{ϵ}{\nabla }_x\mathbf{T}(x)\right)$$

当扰动$ϵ$很小，$ϵ=0$的时候
$\mathbf{T}(x)=x,{\nabla }_{ϵ}\mathbf{T}(x)=\mathbf{\varphi }(x),{\nabla }_x\mathbf{T}(x)=I,{\nabla }_{ϵ}{\nabla }_x\mathbf{T}(x)={\nabla }_x\mathbf{\varphi }(x)$.

代入上式可得${\nabla }_{ϵ}\log p_{\left[T^{-1}\right]}(x)=\mathrm{trace}\left({\mathcal{A}}_p\varphi (x)\right)$

所以有${{\nabla }_{ϵ}\mathrm{K}\mathrm{L}\left(q_{[T]}\Vert p\right)|}_{ϵ=0}=-{\mathbb{E}}_{x\sim q}\left[\mathrm{trace}\left({\mathcal{A}}_p\varphi (x)\right)\right]$

即证明了KL散度变化最快的方向就是KSD所对应的向量函数${\varphi }^{\ast }=\mathbf{\beta }/\Vert \mathbf{\beta }{\Vert }_{{\mathcal{H}}^d}$

5.2 Algorithm

有了梯度方向，有了kernel，我们就可以设计算法了，就是SVGD的实现。
算法流程如下所示：

算法的实现比较清晰明了，相比之前的kernel引入推导更容易快速理解一些。 $p(\mathbf{x})$ 是我们想要逼近的 ${\mathbb{R}}^d$ 分布, 我们想要用若干粒子来做 $p(\mathbf{x})$上的采样。选取粒子群 ${\left\{{\mathbf{x}}_i\right\}}_{i=1}^n\subset {\mathbb{R}}^d$ . 使用梯度下降法，设置学习率为$ϵ$, 每一步的梯度为 $\varphi \left({\mathbf{x}}_i\right)$.不断迭代直至$l$达到指定次数或梯度变化达到阈值。

总的来说，这个idea感觉很赞，我觉得从${\mathbb{E}}_{x\sim p}\left[(s_q(x)-s_p(x)){\varphi }^{\top }(x)\right]$看到${\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]$这一步真的很强，毕竟是ICML、NeurIPS连着中的组（

5.实验

关于实验部分，A General Purpose Bayesian Inference Algorithm这篇中做了三个SVGD应用的实验，高斯分布逼近，逻辑回归以及贝叶斯神经网络的训练。效果好像都不错。

Github实现代码在这里：https://github.com/DartML/Stein-Variational-Gradient-Descent

关于一维分布拟合的迭代过程如下：

References
[1]Liu Q , Wang D . Stein Variational Gradient Descent: A General Purpose Bayesian Inference Algorithm[C]// 2016.
[2]Liu Q , Jason D.Lee . A Kernelized Stein Discrepancy for Goodness-of-fit Tests[C]// 2016.
[3]Liu Q , Stein Variational Gradient Descent as Gradient Flow[C]// 2016.
[4]Stein variational gradient descent https://zhuanlan.zhihu.com/p/114489837