MCMC算法大统一: Involutive MCMC

标准采样方法

先从最基本的方法说起，可能很多人都知道只要可以对分布函数$F(x)$求逆，并从均匀分布中采样u,并将u代进逆函数中就能得到x的样本，即$x=F^{-1}(u),u\sim U(0,1)$。他的原理是什么？其实他的出发点是找到一个从均匀分布到目标分布的可逆变换$g$：

$$\begin{gathered} x=g(u) \\ p(x)=p(u)\left|\frac{du}{dx}\right|=p_u\left(g^{-1}(x)\right)\left|\frac{d{g}^{-1}(x)}{dx}\right| \end{gathered}$$

因为我们要保证x是我们的目标分布，那么他一定要满足${\int }_{-\infty }^{x}p(x)dx=F(x)$，因为$p_u\left(f^{-1}(x)\right)$是0到1的均匀分布所以他等于1，于是必须有${\int }_{-\infty }^x\left|\frac{d{g}^{-1}(x)}{dx}\right|dx=F(x)⟹g^{-1}(x)=F(x)⟹g(x)=F^{-1}(x)$。
然而该方法的问题在于，并不是每个概率分布的都可以写出来并且可以求逆的。对于那些特别复杂的分布这样方法就无能为力了。

拒绝采样

现在我们的目标是要采样$p(x)$的分布。为了采样这个分布，实际上，我们只需要知道$\tilde{p}(x)$就足够了，$\tilde{p}(x)$是没有标准化的$p(x)=\tilde{p}(x)/Z$, 其中$Z=\int \tilde{p}(x)dx$ 是不需要知道的。那么拒绝采样的流程是这样的，首先随便找一个proposal distribution$q(x)$ 使得$Mq(x)⩾\tilde{p}(x)$，然后从q中随机一个样本$x\sim q(x)$，再从均匀分布中随机一个样本$u\sim U(0,1)$，如果$u⩾\frac{\tilde{p}(x)}{Mq(x)}$则拒绝，否则接受。为什么这样就能从p中采样？

订正：图中的$p$应该是$\tilde{p}$

首先，因为u是0,1区间，所以一定有$0⩽uMq(x)⩽Mq(x)$，那么u的作用其实就是选择0到$Mq(x)$之间的值，如果$\tilde{p}(x)⩽uMq(x)$，那么就拒绝掉样本，也就是上图白色部分，否则$\tilde{p}(x)⩾uMq(x)$就是落在黑色区域，则接受。显然，对于x轴上的任意一个点$x^{\prime }$，其对应接受的概率就是$\frac{\tilde{p}(x^{\prime })}{Mq(x^{\prime })}$。不过每个点，因为q，抽到的概率都不一样，所以我们要把q抽中的概率也算上去。

那么用这个方法我们可以得到一个序列的样本$(x_1,accep{t}_1,...,x_i,accep{t}_i)$，其中$x_i\sim q(x)$是由$q$产生的，而且每个pair$(x_i,accep{t}_i)$都是相互独立的，那么我们可以证明，$p(x|accept)$就等于我们想要的分布$p(x)$:

$$\begin{array}{rl}p(x|accept)& =\frac{p(accept|x)q(x)}{p(accpet)}\\ & =\frac{\frac{\tilde{p}(x)}{Mq(x)}q(x)}{\int \frac{\tilde{p}(x)}{Mq(x)}q(x)dx}\\ & =\frac{\tilde{p}(x)}{\int \tilde{p}(x)dx}\\ & =\frac{p(x)/Z}{\int p(x)/Zdx}\\ & =p(x)\end{array}$$

通过拒绝采样，我们可以知道，对于任意的分布，我们可以设计一种接受-拒绝的概率，从而使得所有接受的点都是该分布的样本。然而他的问题在于，他对propose distribution q的要求很高，试想下，如果q与真实p的差距过大，那么几乎每个样本都会被拒绝掉，效率很低。那如何设计更好的$q$呢？一个办法是设计一个$t(x^{\prime }|x)$，他从上一次接受的样本作为条件，然后经过转换propose一个新的样本，如果我们能够约束这个转换不会改变目标分布的，那么我们就能快速的找到$p(x)$的样本。这也正是MCMC的思想。

IMCMC

然而MCMC的方法少说也有成百上千种，如何统一起来呢？我们从另一个角度来考虑MCMC。首先MCMC最基本可以从离散马尔科夫链说起，A是一个转移矩阵，$\pi$是一个向量，如果满足下述公式

$$\pi A=\pi$$

则称$A$的平稳分布是$\pi$。类似的在连续的情况下：

$$\int t(x^{\prime }|x)p(x)dx=p(x^{\prime })$$

我们希望找到一个转换概率分布$t(x^{\prime }|x)$，且该转换不会改变来自$p(x)$的样本的分布。事实上这个t不一定是随机的，当他是确定的映射时，我们有$t(x^{\prime }|x)=\delta (x^{\prime }-f(x))$,相当于$x^{\prime }=f(x)$于是
$$\int \delta (x^{\prime }-f(x))p(x)dx=p(x^{\prime })$$

这意味着我们要找到一个映射函数，使得$x^{\prime }=f(x)$并且他们的分布一样，又因为根据分布变换公式

$$\begin{gathered} p_x(x)=p_{x^{\prime }}(x^{\prime }) \\ {p}_x(x)=p_{x^{\prime }}(f(x))\left|\frac{df(x)}{dx}\right|=p_x(f(x))\left|\frac{df(x)}{dx}\right| \\ =p_{x^{\prime }}(x^{\prime })=p_x\left(f^{-1}(x)\right)\left|\frac{d{f}^{-1}(x)}{dx}\right| \end{gathered}$$
（其实从上述公式中我们就能大致的能看出，$f(x)=f^{-1}(x)$这个条件的成立，这就是所谓的Involutive function，IMCMC就是围绕如何设计这个f函数来做的）
不过这样的f不好找，所以一般我们可以用以下这个：

$$\begin{array}{r}t(x^{\prime }|x)=\delta (x^{\prime }-f(x))\underset{P_{\text{accept}}}{\min \{1,\frac{p(f(x))}{p(x)}|\frac{\partial f}{\partial x}\left|\right\}}+\\ +\delta (x^{\prime }-x)\underset{P_{\text{reject}}}{(1-\min \{1,\frac{p(f(x))}{p(x)}\left|\frac{\partial f}{\partial x}\right|\left\}\right)},\end{array}$$

我们一般可以通过拒绝采样的方式来近似这个转换概率(这部分我不是很理解他跟拒绝采样的关系，没看懂)。这意味着，当$x^{\prime }=f(x)$时，$t(x^{\prime }|x)=\min \{1,\frac{p(f(x))}{p(x)}\left|\frac{\partial f}{\partial x}\right|\}$，而当$x^{\prime }=x$时，$t(x^{\prime }|x)=(1-\min \{1,\frac{p(f(x))}{p(x)}|\frac{\partial f}{\partial x}\left|\right\})$，显然，如果$x\sim p(x)$，且$x^{\prime }=f(x)$，我们写的稍微简单一点，那么$p_{x^{\prime }}(x^{\prime })=t(x^{\prime }|x)p(x)=p(f(x))|\frac{\partial f}{\partial x}|=p_x\left(f^{-1}(x)\right)\left|\frac{d{f}^{-1}(x)}{dx}\right|$。事实上，这等价于认为，$f$满足$f(x)=f^{-1}(x)$，这样的函数称为involutive function。然而，这个分布只是在两个点之间反复横跳，我们可以引入一个辅助变量$v$，我们随便抽取$v$的样本，从而得到不同的$x^{\prime }$，这样就不会反复横跳了。可以证明，引入一个辅助变量后，只要满足$f(x,v)=f^{-1}(x,u)$，那么x的平稳分布仍然是目标分布。

$$t(x^{\prime },v^{\prime }|x,v)p(x,v)=t(x,v|x^{\prime },v^{\prime })p(x^{\prime },v^{\prime }).$$

可以证明，t的边缘分布仍然服从detailed balance
$$\begin{array}{r}\hat{t}(x^{\prime }|x)=\int dvd{v}^{\prime }t(x^{\prime },v^{\prime }|x,v)p(v|x)\end{array}\begin{array}{r}\hat{t}(x^{\prime }|x)p(x)=\hat{t}(x|x^{\prime })p(x^{\prime }).\end{array}$$

通过引入辅助变量，我们可以设计出iMCMC的算法，只需保证f是involutive function。

MH

一个特例是$f(x,v)=v,x$，他将$x,v$的位置交换了一下，这时候iMCMC等价于MH算法。显然他的接收概率为

$$P=\min \{1,\frac{p(f(x,v))}{p(x)q(v|x)}\}=\min \{1,\frac{p(v)q(x|v)}{p(x)q(v|x)}\}$$

HMC

另外一个经典的例子是HMC算法，他通过构造一个复杂的函数$f$，从而使得接受率非常高，几乎不会拒绝（实际做的时候就直接接受，不判断拒不拒绝），但是$q$又很简单，实际上启发了很多基于神经网络的MCMC就是在搞这个$f$。

HMC依赖于用leap-frog算子$L$来求解Hamiltonian dynamics的数值积分。对于$L:[x(t),v(t)]\to [x(t+\epsilon ),v(t+\epsilon )]$，

$$\begin{array}{l}v(t+\epsilon /2)=v(t)-\frac{\epsilon }{2}{\nabla }_x(-\log p(x(t)))\\ x(t+\epsilon )=x(t)+\epsilon {\nabla }_v(-\log p(v(t+\epsilon /2)))\\ v(t+\epsilon )=v(t+\epsilon /2)-\frac{\epsilon }{2}{\nabla }_x(-\log p(x(t+\epsilon )))\end{array}$$

后$F$则是一个简单的将辅助变量取负号的函数$F:[x,v]=[x,-v]$，可以论文中给出了证明$F{L}^k={\left(F{L}^k\right)}^{-1}$，而且他的jacobian矩阵的行列式也是为1。里面还有很多算法，这里就不写了，有兴趣自己去看。

参考资料

Neklyudov, Kirill, et al. “Involutive MCMC: a Unifying Framework.” arXiv preprint arXiv:2006.16653 (2020).

How does the proof of Rejection Sampling make sense?

(ML 17.13) Proof of rejection sampling (part 1)

Murphy, Kevin P. Machine learning: a probabilistic perspective. MIT press, 2012.

Bishop, Christopher M. Pattern recognition and machine learning. springer, 2006.