深入浅出GAMP算法（中）： GAMP

前言

在上一篇博客中我们简单介绍了消息传递算法的雏形。 然而， 当变量维度较大时，对每个标量进行逐个的消息计算需要不可能承担的复杂度。 因此， 通过一些近似来简化计算， 则是 AMP (approximate message passing) 类算法的核心思想。这篇博客， 我们将通过详细的数学推导， 来展示 GAMP 算法是如何对消息传递进行近似的。

模型背景

我们旨在解决上图这样的问题， 已知输入$\mathbf{q}$ （先验信息）， 已知输出 $\mathbf{y}$（后验信息)， 已知变换矩阵$\mathbf{A}$， 反推出变量$\mathbf{x}$。 以AWGN信道举例：
$$\mathbf{y}=\mathbf{z}+\mathbf{w}=\mathbf{A}\mathbf{x}+\mathbf{w}$$
$\mathbf{y}$已知而我们试图反推出$\mathbf{x}$。 此时， 如果$\mathbf{x}$有一些先验分布信息， 如稀疏分布， 即$\mathbf{x}$中只有少量元素非零， 那么， 便可以通过概率的方式， 以GAMP算法进行求解。

推导

我们先以 MAP 版本的 GAMP为例。 即最大化后验概率：
$${\hat{\mathbf{x}}}^{\text{map }}:=\underset{\mathbf{x}\in {\mathbb{R}}^n}{\mathrm{argmax}}F(\mathbf{x},\mathbf{z},\mathbf{q},\mathbf{y}),\hat{\mathbf{z}}=\mathbf{A}\hat{\mathbf{x}}$$
其中，
$$F(\mathbf{x},\mathbf{z},\mathbf{q},\mathbf{y}):=\sum _{j=1}^n{f}_{\mathrm{i}\mathrm{n}}\left(x_j,q_j\right)+\sum _{i=1}^m{f}_{\text{out }}\left(z_i,y_i\right)，$$
$f_{\mathrm{i}\mathrm{n}}$ 和 $f_{\mathrm{o}\mathrm{u}\mathrm{t}}$ 分别为：
$$\begin{array}{rl}{f}_{\text{out }}(z,y)& :=\log p_{Y\mid Z}(y\mid z)\\ f_{\text{in }}(x,q)& :=\log p_{X\mid Q}(x\mid q)\end{array}$$
因此 MAP 估计就是找出 可能性最大（对应后验概率最大）的 $x$ 作为估计值。

传统的消息传递算法， 在每次迭代中， 实际要计算两个消息算子：
$$\begin{array}{r}{\Delta }_{i\to j}\left(t,x_j\right)=\mathrm{c}\mathrm{o}\mathrm{n}\mathrm{s}\mathrm{t}+\underset{\mathbf{x}\backslash x_j}{\max }{f}_{\text{out }}\left(z_i,y_i\right)+\sum _{r\ne j}{\Delta }_{i\leftarrow r}\left(t,x_r\right)\end{array}\text{\hspace{1em}(1)}$$
和
$$\begin{array}{r}{\Delta }_{i\leftarrow j}\left(t+1,x_j\right)=\mathrm{c}\mathrm{o}\mathrm{n}\mathrm{s}\mathrm{t}+f_{\mathrm{i}\mathrm{n}}\left(x_j,q_j\right)+\sum _{\ell \ne i}{\Delta }_{\ell \to j}\left(t,x_j\right)\end{array}\text{\hspace{1em}(2)}$$
简单理解(1)和(2)的意义： $i\to j$， 代表的就是由于$z_i=a_i\mathbf{x}$这一限制的存在， 从而对$x_j$进行的修正。 $j\to i$ 则是$x_j$ 向 $i$ 节点传递当前自己 根据 除了 $z_i=a_i\mathbf{x}$ 这条限制以外的 其他限制 进行修正后的 结果。 这里是为了防止消息被重复累加。 在第一篇博客中对消息传递算法有更详细的介绍。
接下来进行近似简化， 我们先有如下一些将用到的定义：
$$\begin{array}{rl}{\hat{x}}_j(t)& :=\underset{x_j}{\mathrm{argmax}}{\Delta }_j\left(t,x_j\right)\\ {\hat{x}}_{i\leftarrow j}(t)& :=\underset{x_j}{\mathrm{argmax}}{\Delta }_{i\leftarrow j}\left(t,x_j\right)\\ \frac{1}{{\tau }_j^x(t)}& :=-{\frac{{\partial }^2}{\partial x_j^2}{\Delta }_j\left(t,x_j\right)|}_{x_j={\hat{x}}_j(t)}\\ \frac{1}{{\tau }_{i\leftarrow j}^x(t)}& :=-{\frac{{\partial }^2}{\partial x_j^2}{\Delta }_{i\leftarrow j}\left(t,x_j\right)|}_{x_j={\hat{x}}_{i\leftarrow j}(t)}.\end{array}$$
其中，
$${\Delta }_j\left(x_j\right):=\underset{\mathbf{x}\backslash x_j}{\max }F(\mathbf{x},\mathbf{z},\mathbf{q},\mathbf{y}),\hat{\mathbf{z}}=\mathbf{A}\hat{\mathbf{x}}$$
其中$\mathbf{x}\backslash x_j$表示$x_j$不变， 优化$\mathbf{x}$的其他元素们。因此${\hat{x}}_j$事实上就对应于我们在消息传递过程中不断迭代得到的最大后验估计。因此在迭代中， 我们有：
$${\Delta }_j\left(t+1,x_j\right)=f_{\mathrm{i}\mathrm{n}}\left(x_j,q_j\right)+\sum _i{\Delta }_{i\to j}\left(t,x_j\right)$$
和(2)相比， ${\Delta }_j$ 就是考虑了所有节点的限制及先验信息，再给出的整体的一个$x_j$的估计。
我们先对(2)进行近似：
将(2)在${\hat{x}}_{i\leftarrow r}(t)$点进行二次泰勒展开， 可以得到：
$${\Delta }_{i\leftarrow r}\left(t,x_r\right)\approx {\Delta }_{i\leftarrow r}\left(t,{\hat{x}}_{i\leftarrow r}(t)\right)-\frac{1}{2{\tau }_r^x(t)}{\left(x_r-{\hat{x}}_{i\leftarrow r}(t)\right)}^2\text{\hspace{1em}(3)}$$
注意，这里的展开里我们忽略了泰勒展开的一次项， 这是因为我们后续的推导中还会对$x_r$进行一次求导， 那么一次项便成了常数项， 不再影响后续的优化推导， 因此，在这里直接省略了。

将(3)代入(1)中，有：

$$\begin{array}{r}{\Delta }_{i\to j}\left(t,x_j\right)=\underset{z_i}{\max }\left[f_{\text{out }}\left(z_i,y_i\right)-\frac{1}{2{\tau }_{i\to j}^p(t)}{\left(z_i-{\hat{p}}_{i\to j}(t)-a_{ij}{x}_j\right)}^2,\right]+\mathrm{c}\mathrm{o}\mathrm{n}\mathrm{s}\mathrm{t}\end{array}\text{\hspace{1em}(4)}$$

其中
$$\begin{array}{rlr}{\hat{p}}_{i\to j}(t)& :=\sum _{r\ne j}{a}_{ir}{\hat{x}}_{i\leftarrow r}(t){\tau }_{i\to j}^p(t)& :=\sum _{r\ne j}{\left|a_{ir}\right|}^2{\tau }_r^x(t)\end{array}$$

推导(4)是一个以 $z_i=a_{ij}{x}_j+{\sum }_{r\ne j}{a}_{ir}{x}_r$ 为约束的优化问题，通过拉格朗日乘子法可以求得闭式解，推导出(4)。这里的详细步骤也不再展开。 注意到： 事实上通过观察 ${\hat{p}}_{i\to j}$ 和 ${\tau }_{i\to j}$的形式我们可以发现，前者对应于 $z_i$的均值， 而后者对应于$z_i$的方差 （各自少了一项）。

我们继续向下推导。定义：
$$H\left(\hat{p},y,{\tau }^p\right):=\underset{z}{\max }\left[f_{\text{out }}(z,y)-\frac{1}{2{\tau }^p}(z-\hat{p}{)}^2\right]$$
那么(4)可以进一步写为：
$${\Delta }_{i\to j}\left(t,x_j\right)=H\left({\hat{p}}_{i\to j}(t)+a_{ij}{x}_j,y_i,{\tau }_{i\to j}^p(t)\right)+\text{ const }\text{\hspace{1em}(5)}$$
这里必须提一嘴的是： $H$的这个形式，你想起了什么吗？没错， 这就是高斯分布随机变量的概率密度函数的指数项——均值为$\hat{p}$， 方差为${\tau }_p$。与之对应， 也可以发现事实上消息 ${\Delta }_{i\to j}\left(t,x_j\right)$ 就是找一个最大后验概率的$x_j$。

然而我们的初衷是为了简化计算， 而（5）中所需要用到的 $\hat{p}$， 依然需要对多个标量进行求和。 但是我们注意到，
如果继续定义：${\hat{p}}_i(t):={\sum }_j{a}_{ij}{\hat{x}}_{i\leftarrow j}(t)$ 和 ${\tau }_i^p(t)={\sum }_j{\left|a_{ij}\right|}^2{\tau }_j^x(t)$， 有：
$$\begin{array}{rl}{\hat{p}}_{i\to j}(t)& ={\hat{p}}_i(t)-a_{ij}{\hat{x}}_{i\leftarrow j}(t)\\ {\tau }_{i\to j}^p(t)& ={\tau }_i^p(t)-a_{ij}^2{\tau }_j^x(t)\end{array}$$
如果矩阵$A$中的所有元素$a_{ij}$都服从同样的分布且已被归一化，我们接下来可以做以下近似： 忽略所有$O(a_{ij}^2)$项， 那么：
(5)可以进一步写为：
$${\Delta }_{i\to j}\left(t,x_j\right)\approx H\left({\hat{p}}_i(t)+a_{ij}\left(x_j-{\hat{x}}_j\right),y_i,{\tau }_i^p(t)\right)+\text{ const }$$
这里是因为我们忽略了 $O\left(a_{ij}^2\right)$级的项， 因此有了${\tau }_{i\to j}^p(t)={\tau }_i^p(t)$的效果。也可以注意到， (5)中的${\hat{x}}_{i\leftarrow j}$也变为了${\hat{x}}_j$, 同样也是用了这个近似。 这一近似在后续也将被持续用到。

许多人看到这里可能一头雾水了， GAMP在干嘛？ 一顿操作， 到底想干什么？ 是的， 其实他的核心就在于， 把涉及到 ${\hat{p}}_{i\to j}(t)$ 和 ${\hat{\tau }}_{i\to j}(t)$的计算， 通通替换为用 ${\hat{p}}_i^p(t)$ 和 ${\tau }_i^p(t)$替代。 因为后者的计算， 在后续可以看到， 能被高度简化。 如果看GAMP的算法框图的也可以发现， 尽管在推导的过程中， 由于是从原始的消息传递算法而来， 有许多$i\to j$这样的项， 然而在最后的算法中， 是没有的。

继续下面的推导。 我们希望进一步的近似也是利用了泰勒展开， 因此我们需要定义一阶导和二阶导：
$$\begin{array}{rl}{\hat{s}}_i(t)& =g_{\text{out }}\left(t,{\hat{p}}_i(t),y_i,{\tau }_i^p(t)\right)\\ {\tau }_i^s(t)& =-\frac{\partial }{\partial \hat{p}}{g}_{\text{out }}\left(t,{\hat{p}}_i(t),y_i,{\tau }_i^p(t)\right)\end{array}\text{\hspace{1em}(6)}$$
其中， $$g_{\text{out }}\left(\hat{p},y,{\tau }^p\right):=\frac{\partial }{\partial \hat{p}}H\left(\hat{p},y,{\tau }^p\right)$$
那么根据$H$函数的定义， 这里$g_{\mathrm{o}\mathrm{u}\mathrm{t}}$可化简为下式：
$$g_{\text{out }}\left(\hat{p},y,{\tau }^p\right):=\frac{1}{{\tau }^p}\left({\hat{z}}^0-\hat{p}\right)$$
其中 ${\hat{z}}^0:=\underset{z}{\mathrm{argmax}}{F}_{\text{out }}\left(z,\hat{p},y,{\tau }^p\right)$ 和 $F_{\text{out }}\left(z,\hat{p},y,{\tau }^p\right):=f_{\text{out }}(z,y)-\frac{1}{2{\tau }^p}(z-\hat{p}{)}^2$。 进一步的，二阶导可以求得：
$$-\frac{\partial }{\partial \hat{p}}{g}_{\text{out }}\left(\hat{p},y,{\tau }^p\right)=\frac{-f_{\text{out }}^{\prime \prime }\left({\hat{z}}^0,y\right)}{1-{\tau }^p{f}_{\text{out }}^{\prime \prime }\left({\hat{z}}^0,y\right)}$$
这里很明确了， (6)中的 ${\hat{s}}_i(t)$ 和 ${\tau }_i^s(t)$ 分别代表了 $H$ 函数 在 ${\hat{p}}_i(t)$的一阶导和二阶导， 那么继续泰勒展开：

$$\begin{array}{rl}{\Delta }_{i\to j}\left(t,x_j\right)& \approx \mathrm{c}\mathrm{o}\mathrm{n}\mathrm{s}\mathrm{t}+s_i(t)a_{ij}\left(x_j-{\hat{x}}_j(t)\right)-\frac{{\tau }_i^s(t)}{2}{a}_{ij}^2{\left(x_j-{\hat{x}}_j(t)\right)}^2\\ & =\mathrm{const}+\left[s_i(t)a_{ij}+a_{ij}^2{\tau }_i^s(t){\hat{x}}_j(t)\right]x_j-\frac{{\tau }_i^s(t)}{2}{a}_{ij}^2{x}_j^2\end{array}\text{\hspace{1em}(7)}$$
可以看到， 此时${\Delta }_{i\to j}\left(t,x_j\right)$中已经没有任何$i\to j$这样的项了。 但是$\hat{p}$里其实还有， 不过这个我们后续再说。 接下来我们对另一个方向的消息传递进行分析：
此时， 把(7)代入到(2)中， (2)也可以被近似为：
$$\begin{array}{r}{\Delta }_{i\leftarrow j}\left(t+1,x_j\right)\approx \mathrm{c}\mathrm{o}\mathrm{n}\mathrm{s}\mathrm{t}+f_{\mathrm{i}\mathrm{n}}\left(x_j,q_j\right)-\frac{1}{2{\tau }_{i\leftarrow j}^r(t)}{\left({\hat{r}}_{i\leftarrow j}(t)-x_j\right)}^2\end{array}$$
其中，
$$\frac{1}{{\tau }_{i\leftarrow j}^r(t)}=\sum _{\ell \ne i}{a}_{\ell j}^2{\tau }_{\ell }^s(t)$$
和
$$\begin{array}{rl}{\hat{r}}_{i\leftarrow j}(t)& ={\tau }_{i\leftarrow j}^r(t)\sum _{\ell \ne i}\left[s_{\ell }(t)a_{\ell j}+a_{\ell j}^2{\tau }_{\ell }^s(t){\hat{x}}_j(t)\right]\\ & ={\hat{x}}_j(t)+{\tau }_{i\leftarrow j}^r(t)\sum s_{\ell }(t)a_{\ell j}\end{array}$$
注意， 这里算法又颇具技巧性地，加入了常数项， 从而把右边的形式再一次写成了高斯分布的概率密度函数的次方项！这无疑是故意而为之的。

这时候，如果我们定义：
$$g_{\text{in }}\left(\hat{r},q,{\tau }^r\right):=\underset{x}{\mathrm{argmax}}{f}_{\text{in }}(x,q)-\frac{1}{2{\tau }^r}(\hat{r}-x{)}^2$$
那么有：
$${\hat{x}}_{i\leftarrow j}(t+1)\approx g_{\text{in }}\left({\hat{r}}_{i\leftarrow j}(t),q_j,{\tau }_{i\leftarrow j}^r(t)\right)\text{\hspace{1em}(8)}$$
**注意！这里我们第一次出现了$t+1$次迭代中的项！我们已经成功地完成了一次消息传递但， 还需要具体的简化， 即去除所有$i\to j$相关项。
再定义：
$$\begin{array}{rl}& {\tau }_j^r(t)={\left[\sum _i{\left|a_{ij}\right|}^2{\tau }_i^s(t)\right]}^{-1}\\ & {\hat{r}}_j(t)={\hat{x}}_j(t)+{\tau }_j^r(t)\sum _i{a}_{ij}{\hat{s}}_i(t)\end{array}$$
又可以有下列近似关系：
$$\begin{array}{rl}{\tau }_{i\leftarrow j}^r(t)& \approx {\tau }_j^r(t)\\ {\hat{r}}_{i\leftarrow j}(t)& \approx {\hat{x}}_j(t)+{\tau }_j^r(t)\sum _{\ell \ne i}{s}_{\ell }(t)a_{\ell j}\\ & ={\hat{r}}_j(t)-{\tau }_j^r(t)a_{ij}{s}_i(t)\end{array}$$
这里同样是忽略了 $O\left(a_{ij}^2\right)$ 级的项。 利用这一近似， （8）可以近似为：
$$\begin{array}{rl}& {\hat{x}}_{i\leftarrow j}(t+1)\\ & \stackrel{(a)}{\approx }{g}_{in}\left({\hat{r}}_j(t)-a_{ij}{s}_i(t){\tau }_j^r(t),q_j,{\tau }_j^T(t)\right)\\ & \stackrel{(b)}{\approx }{\hat{x}}_j(t+1)-a_{ij}{s}_j(t)D_j(t+1)\end{array}\text{\hspace{1em}(9)}$$
这里(a)就是直接代入了近似关系， 而(b)则是对进行了泰勒一次展开：
$$\begin{array}{rl}{\hat{x}}_j(t+1)& :=g_{\text{in }}\left({\hat{r}}_j(t),q_j,{\tau }_j^r(t)\right)\\ D_j(t+1)& :={\tau }_j^r(t)\frac{\partial }{\partial \hat{r}}{g}_{\text{in }}\left({\hat{r}}_j(t),q_j,{\tau }_j^r(t)\right)\end{array}$$
注意到又有：
$$\begin{array}{rl}& D_j(t+1)\stackrel{(a)}{\approx }{\tau }_j^r(t)\frac{\partial }{\partial \hat{r}}\left(\Gamma f_{\mathrm{i}\mathrm{n}}\left(\cdot ,q_j\right)\right)\left({\hat{r}}_j(t),{\tau }_j^r(t)\right)\\ & \stackrel{(b)}{=}\frac{{\tau }_j^T(t)}{1-{\tau }_j^r(t)f_{in}^{\prime \prime }\left({\hat{x}}_j(t+1),q_j\right)}\\ & \stackrel{(c)}{\approx }{\left[-\frac{{\partial }^2}{\partial x_j^2}{\Delta }_{i\leftarrow j}\left(t+1,{\hat{x}}_j(t+1)\right)\right]}^{-1}\\ & \stackrel{(d)}{\approx }{\tau }_j^x(t+1)\end{array}$$
这里具体的推导就不展开了。 注意到， 根据(9)我们其实已经得到了这一次迭代的结果 ${\hat{x}}_j(t+1)$。并且， 将上式代入(9)和${\hat{p}}_i(t)$的定义中，我们有：
$${\hat{p}}_i(t)=\sum _j{a}_{ij}{\hat{x}}_j(t)-{\tau }_i^p(t){\hat{s}}_i(t-1)$$

也就是说， 整个算法运行中， 不再涉及任何$i\to j$项的操作了！

至此，推导结束。
总结以下几个点：

• 整个GAMP推导的核心任务就是尽可能地把所有操作变为 关于 $x_j$, $p_i$这样的计算， 而不涉及 $x_{i\leftarrow j}$， $p_{i\leftarrow j}$， 从而大幅地降低计算复杂度。
• 使用的方法就是不断假设整个矩阵$\mathbf{A}$的元素服从同分布的高斯，从而可以进行不断的泰勒展开近似。

最终的算法流程可以被简化为：

我们可以将算法流程整理一下：

• 初始化
• 计算 ${\tau }_i^p(t)$ 和 ${\hat{p}}_i(t)$, ${\hat{s}}_i$ 和 ${\tau }_i^s(t)$, 从而计算 ${\Delta }_{i\to j}\left(t,x_j\right)$， 在GAMP算法里最后体现在对 ${\hat{r}}_j$的计算中。
• 计算 ${\tau }_j^r$ 和 ${\hat{r}}_j$， 用于最后对 ${\hat{x}}_j$的计算。