AMP软阈值函数进阶

问题背景

AMP(Approximate Message Passing)算法与BP(Belief Propagation)算法，或者松弛BP(relaxed-BP)算法有一个区别就是AMP可以不知道估计向量每个元素的先验分布。问题描述如下

$$y\text{y}y=A\text{A}Ax\text{x}x+n\text{n}n\text{\hspace{1em}(1)}$$

其中$y\text{y}y$为观测向量，$n\text{n}n$为噪声，且已知似然分布$p(y_m|(A\text{A}Ax\text{x}x{)}_m),\forall m$，不要求知道$x_n(\forall n)$的先验分布$p(x_n)$，AMP依然可以对$x\text{x}x$做回归估计。

本质上，AMP所要解决的问题是一个压缩感知问题，所以向量$x\text{x}x$是稀疏的，即$\exists K\in {\mathbb{N}}^{+},{\Vert x\text{x}x\Vert }_0\le K$。据此，我们可以引入一个拉普拉斯分布（AMP的推导和理解-Part-1中有介绍）来近似$x_n(\forall n)$的先验分布$p(x_n)$，记为

$$p(x_n)\cong \mathrm{e}\mathrm{x}\mathrm{p}(-\beta x_n)\text{\hspace{1em}(2)}$$

这样描述的意义何在，我们以LASSO(BPDN)问题下的联合概率密度函数举例，如式(3)所示，先不考虑其他参数，当$\beta \to \infty$时，可以发现，LASSO的解应与该分布的mode统一（实在不知道mode怎么翻译），所以考虑拉普拉斯分布是合理的。

$$\mu (\text{d}s\text{s}s)=\frac{1}{Z}\prod _{i=1}^{N}exp(-\beta \lambda |s_i|)\underset{\text{likelihood function}}{\underbrace{\prod _{a=1}^{n}exp\{-\frac{\beta }{2}(y_a-(As{)}_a{)}^2\}}}\text{\hspace{1em}(3)}$$

下面所有的叙述与AMP的推导和理解-Part-1中的“拉普拉斯分布参数的极限近似”一致，是对其做更具体地解释。（因为我写Part-1的时候这部分没有理清楚）

软阈值函数与其等价形式

软阈值函数(soft-threshold function)形式如下
$$\eta (x;b)=\{\begin{array}{rl}& x-bifx>b\\ & 0if-b\le x\le b\\ & x+bifx<-b\end{array}\text{\hspace{1em}(4)}$$

软阈值函数与下式，即式(5)是等价的

$$\eta (x;b)=\underset{s\in \mathbb{R}}{\mathrm{argmin}}\{|s|+\frac{1}{2b}(s-x{)}^2\}\text{\hspace{1em}(5)}$$

对于两者等价的解释：式(5)想要做的就是

$$\mathrm{f}\mathrm{i}\mathrm{n}\mathrm{d}s,\mathrm{s}\mathrm{u}\mathrm{c}\mathrm{h}\mathrm{t}\mathrm{h}\mathrm{a}\mathrm{t},\min |s|+\frac{1}{2b}(s-x{)}^2$$

下面就是中学数学了，我们把式(5)括号中的内容展开得到

$$\{\begin{array}{rl}& \frac{1}{2b}(s-(x-b){)}^2+const,s>0\\ & \frac{1}{2b}(s-(x+b){)}^2+const,s<0\\ & \text{corresponding to }s=0\end{array}\text{\hspace{1em}(6)}$$

联系式(5)和式(6)可以得到式(4)，式(4)也可以得到式(6)，也就是说式(4)与式(5)等价。

拉普拉斯分布参数下的极限近似均值($\beta \to \infty$)

在Part-1中我们知道，从因子节点到变量节点的消息通过Berry-Esseen中心极限定理被近似为高斯分布，进而从变量节点到因子节点的消息可以被表示为一组高斯分布乘积与拉普拉斯先验分布（假设的先验）的乘积，该函数形式可以被表示为如下形式（Part-1中也有叙述）：

f β ( s ; x , b ) ≡ 1 z β ( x , b ) e x p { − β ∣ s ∣ − β 2 b ( s − x ) 2 } (7) f_{\beta}(s;x,b) \equiv \frac{1}{z_{\beta}(x,b)}exp\{ -\beta |s|-\frac{\beta}{2b}(s-x)^2 \} \tag{7} fβ​(s;x,b)≡zβ​(x,b)1​exp{−β∣s∣−2bβ​(s−x)2}(7)

若随机变量$Z$的概率密度函数为$f_{\beta }(\cdot ;x,b)$，那么定义其均值和方差为

$$F_{\beta }(x;b)\equiv {\mathbb{E}}_{f_{\beta }(\cdot ;x,b)}(Z),G_{\beta }(x;b)\equiv {\text{Var}}_{f_{\beta }(\cdot ;x,b)}(Z)\text{\hspace{1em}(8)}$$

首先，对概率密度函数为$f_{\beta }(\cdot ;x,b)$，一定满足

∫ 1 z β ( x , b ) e x p { − β ( ∣ s ∣ + 1 2 b ( s − x ) 2 ) } d s = 1 (9) \int \frac{1}{z_{\beta}(x,b)}exp\{ -\beta(|s|+\frac{1}{2b}(s-x)^2 )\} \mathrm{d} s = 1 \tag{9} ∫zβ​(x,b)1​exp{−β(∣s∣+2b1​(s−x)2)}ds=1(9)

接下来考虑当$\beta \to \infty$时，$f_{\beta }(s;x,b)$被$|s|+\frac{1}{2b}(s-x{)}^2$的最小值主导，假设$\log f_{\beta }(s_{\ast };x,b)\cong \min |s|+\frac{1}{2b}(s-x{)}^2$，即式(5)，那么式(9)可以被近似为

$$\begin{gathered} \int \delta (s-s_{\ast })\mathrm{d}s=1 \\ or{f}_{\beta }(s;x,b)\to \delta (s-s_{\ast })(\beta \to \infty )\text{\hspace{1em}(10)} \end{gathered}$$

所以期望$F_{\beta }(x;b)$可以表示为

$$\underset{\beta \to \infty }{\lim }{F}_{\beta }(x;b)=\underset{\beta \to \infty }{\lim }\int s{f}_{\beta }(s;x,b)\mathrm{d}s=\int s\delta (s-s_{\ast })\mathrm{d}s=s_{\ast }=\eta (x;b)\text{\hspace{1em}(11)}$$

即均值为

$$\underset{\beta \to \infty }{\lim }{F}_{\beta }(x;b)=\eta (x;b)\text{\hspace{1em}(12)}$$

从均值到方差

式(8)中还有一项没有考虑， 方差项$G_{\beta }(x;b)\equiv {\text{Var}}_{f_{\beta }(\cdot ;x,b)}(Z)$。再回过头来看式(7)，我们可以把式(7)改写为

f β ( s ; x , b ) ≡ 1 z β ( x , b ) e x p ( − β ∣ s ∣ ) ⏟ assume prior  p ( s ) ⋅ e x p { − β 2 b ( s − x ) 2 } ⏟ assume likelihood p(x|s) (13) f_{\beta}(s;x,b) \equiv \frac{1}{z_{\beta}(x,b)} \underbrace{exp(-\beta |s|) }_{\text{assume prior } p(s)}\cdot \underbrace{exp\{-\frac{\beta}{2b}(s-x)^2 \}}_{\text{assume likelihood p(x|s)}} \tag{13} fβ​(s;x,b)≡zβ​(x,b)1​assume prior p(s) exp(−β∣s∣)​​⋅assume likelihood p(x|s) exp{−2bβ​(s−x)2}​​(13)

由式(13)可知，将$f_{\beta }(s;x,b)$拆为先验分布和似然分布（观测量为$x$）的乘积，其中似然分布为高斯分布。形式上可以将$f_{\beta }(s;x,b)$理解为是模型$x=s+n,s\sim p(s),n\sim \mathcal{N}(0,\frac{b}{\beta })$的后验概率密度函数，根据AWGN后验估计下的均值与方差关系，即

$$\frac{\partial }{\partial x}{f}_{\beta }(s;x,b)=\frac{\beta }{b}{G}_{\beta }(x;b)$$

即

$$\underset{\beta \to \infty }{\lim }\beta G_{\beta }(x;b)=b{\eta }^{\prime }(x;b)$$