AMP State Evolution的计算：以伯努利-高斯先验为例

AMP State Evolution (SE)的计算

$t = 1$ 时， $\mathcal E^{(t)} = \mathbb E [X^2]$ ，SE的迭代式为
$\begin{aligned} \tau^{(t)}_r &= \sigma^2 + \frac{1}{\delta} \mathcal E^{(t)} \\ \mathcal E^{(t+1)} &= \mathbb E \left | \eta^{(t)} \left ( X + Z \right) - X \right |^2, \ Z \sim \mathcal N(0, \tau^{(t)}_r) \end{aligned}$

撰写的时候存在一定的符号乱用，之后的 $\tau$ 即指 $\tau^{(t)}_r$ 。

注意到， $\mathcal E^{(t+1)} = \mathbb E \left | \eta^{(t)} \left ( X + Z \right) - X \right |^2$ 是关于随机变量 $X, Z$ 求期望，这里我们认为 $X, Z$ 之间相互独立，因此
$\mathcal E^{(t+1)} = \int \int p_{X, Z}(X, Z) \left | \eta^{(t)} \left ( X + Z \right) - X \right |^2 dX dZ$

令 $R = X + Z$ ，不难得到 $p_{X, R}(X, R)$ 等价于 $p_{X, Z}(X, Z)$ （因为 $p_{X, Z}(X=x_0, Z=z_0) = p_{X, R}(X=x_0, R=x_0 + z_0)$ ）,因此我们可以把 $\mathcal E^{(t+1)}$ 写为（这里我们考虑 $\eta^{(t)}$ 为MMSE函数）
$\begin{aligned} \mathcal E^{(t+1)} &= \int \int p_{X, R}(X, R) \left | \eta^{(t)} \left ( R \right) - X \right |^2 dX dR \\ &= \int \int p_R(R) p_{X|R}(X|R) \left | \mathbb E\left [ X|R \right] - X \right |^2 dX dR \\ &= \int p_R(R) \mathrm{var}[X|R] dR \end{aligned}$

因为 $p_{X, Z}(X, Z) = p_X(X) \cdot p_X(Z) = p_X(X) \mathcal N(z;0, \tau^{(t)}_r)$ ，因此可以得到
$\begin{aligned} p_{X, R}(X, R) &= p_X(X) \mathcal N(R-X;0, \tau^{(t)}_r) \\ &= p_X(X) \mathcal N(R;X, \tau^{(t)}_r) \end{aligned}$

X的先验为为伯努利高斯分布时

$p_X(X) \equiv (1-\rho) \delta(X) + \rho \mathcal N(\mu, \nu)$

那么， $X, R$ 的联合分布可写为
$\begin{aligned} p_{X, R}(X, R) &= \left ((1-\rho) \delta(X) + \rho \mathcal N(X; \mu, \nu) \right) \cdot \mathcal N(X;R, \tau^{(t)}_r) \\ &= (1-\rho) \delta(X) \mathcal N(X;R, \tau^{(t)}_r) + \rho \mathcal N(X; \mu, \nu) \mathcal N(X;R, \tau^{(t)}_r) \end{aligned}$

进一步，我们可以得到关于 $R$ 的边缘分布
$\begin{aligned} p_R(R) &= \int p_{X, R}(X, R) dX \\ &= (1-\rho) \mathcal N(0;R, \tau^{(t)}_r) + \rho \int \mathcal N(X; \mu, \nu) \mathcal N(X;R, \tau^{(t)}_r) dX \\ &= (1-\rho) \mathcal N(0;R, \tau^{(t)}_r) + \rho \frac{1}{\sqrt{2 \pi (\nu + \tau^{(t)}_r)}} \exp \left \{ - \frac{(R - \mu)^2}{2 (\nu + \tau^{(t)}_r)} \right \} \\ &= (1-\rho) \mathcal N(R; 0, \tau^{(t)}_r) + \rho \mathcal N(R; \mu, \nu + \tau^{(t)}_r) \end{aligned}$

我们计算后验均值 $\mathbb E[X|R]$
$\begin{aligned} \mathbb E[X|R] &= \int X p_{X|R} (X|R) dX \\ &= \frac{1}{p_R(R)} \int X p_{X,R} (X,R) dX \\ &= \frac{1}{p_R(R)} \cdot \rho \cdot \int X \mathcal N(X; \mu, \nu) \mathcal N(X;R, \tau^{(t)}_r) dX \\ &= \frac{1}{p_R(R)} \cdot \rho \cdot \mathcal N (R; \mu, \nu + \tau^{(t)}_r) \cdot \frac{\mu \tau^{(t)}_r + R \nu}{ \nu + \tau^{(t)}_r } \\ &= \frac{ \rho \cdot \mathcal N (R; \mu, \nu + \tau^{(t)}_r)}{(1-\rho) \mathcal N(R; 0, \tau^{(t)}_r) + \rho \mathcal N(R; \mu, \nu + \tau^{(t)}_r) } \cdot \frac{\mu \tau^{(t)}_r + R \nu}{ \nu + \tau^{(t)}_r } \\ &= \frac{1}{ 1+ \frac{1-\rho}{\rho} \cdot \frac{\mathcal N(R; 0, \tau^{(t)}_r) }{ \mathcal N(R; \mu, \nu + \tau^{(t)}_r)} } \cdot \frac{\mu \tau^{(t)}_r + R \nu}{ \nu + \tau^{(t)}_r } \\ \end{aligned}$

对于大多数问题，我们有 $\mu=0$ ，令 $\phi (R) = \frac{1-\rho}{\rho} \cdot \frac{\mathcal N(R; 0, \tau^{(t)}_r) }{ \mathcal N(R; \mu, \nu + \tau^{(t)}_r)}$ , 当 $\mu=0$ 时，我们有
$\begin{aligned} \mathbb E[X|R] = \frac{1}{1 + \phi(R)} \cdot \frac{ R \nu}{ \nu + \tau^{(t)}_r } \end{aligned}$

进一步，可以得到
$\frac{\partial}{\partial R} \mathbb E[X|R] =\frac{\nu}{\nu + \tau^{(t)}_r } \cdot \frac{1+\phi(R) + R^2 ( \frac{1}{\tau } - \frac{1}{\tau + \nu} ) \phi(R) }{ (1+\phi(R))^2 }$

根据AWGN后验均值与方差的关系，
$\mathrm{var}[X|R] = \tau^{(t)}_r \cdot \frac{\partial}{\partial R} \mathbb E[X|R] =\frac{\nu \tau^{(t)}_r }{\nu + \tau^{(t)}_r } \cdot \frac{1+\phi(R) + R^2 ( \frac{1}{\tau } - \frac{1}{\tau + \nu} ) \phi(R) }{ (1+\phi(R))^2 }$

不难看出， $p_R(R)$ 可以表示为 $p_R(R) = (1 + \phi(R)) \cdot \rho \mathcal N(R; 0, \nu + \tau^{(t)}_r)$ ，因此
$\begin{aligned} \mathcal E^{(t+1)} &= \int p_R(R) \mathrm{var}[X|R] dR \\ &= \int (1 + \phi(R)) \cdot \rho \mathcal N(R; 0, \nu + \tau^{(t)}_r) \cdot \frac{\nu \tau^{(t)}_r }{\nu + \tau^{(t)}_r } \cdot \frac{1+\phi(R) + R^2 ( \frac{1}{\tau } - \frac{1}{\tau + \nu} ) \phi(R) }{ (1+\phi(R))^2 } \\ &= \rho \frac{\nu \tau^{(t)}_r }{\nu + \tau^{(t)}_r } \cdot \int \mathcal N(R; 0, \nu + \tau^{(t)}_r) \left \{ 1 + R^2 \frac{\nu}{\tau(\tau + \nu) } \left( 1- \frac{1}{1 + \phi (R)} \right) \right \} dR \\ &= \rho \frac{\nu \tau^{(t)}_r }{\nu + \tau^{(t)}_r } \cdot \left \{ 1 + \frac{\nu}{\tau} - \frac{\nu}{\tau(\tau + \nu) } \int \frac{R^2}{1 + \phi(R)} \mathcal N(R; 0, \nu + \tau^{(t)}_r) dR \right \} \\ & = \rho \cdot \nu - \rho \cdot \frac{\nu^2}{(\nu+ \tau)^2} \int \frac{R^2}{1 + \phi(R)} \mathcal N(R; 0, \nu + \tau^{(t)}_r) dR \end{aligned}$

其中
$\frac{R^2}{1 + \phi(R)} = \frac{R^2}{ 1 + \frac{1-\rho}{\rho} \cdot \sqrt{\frac{\nu + \tau}{\tau}} \exp \left ( - \frac{\nu}{2\tau(\tau + \nu)} R^2 \right) }$

总结

当 $X$ 的先验是伯努利高斯时： $\sim (1- \rho) \delta(X)+ \rho \mathcal N (0, \nu)$ ，可以把SE表征为
$t = 1$ 时， $\mathcal E^{(t)} = \mathbb E [X^2]= \rho \nu$ ，SE的迭代式为
$\begin{aligned} \tau^{(t)}_r &= \sigma^2 + \frac{1}{\delta} \mathcal E^{(t)} \\ \mathcal E^{(t+1)} &= \rho \cdot \nu - \rho \cdot \frac{\nu^2}{(\nu+ \tau^{(t)}_r)^2} \int \frac{R^2}{1 + \phi(R)} \mathcal N(R; 0, \nu + \tau^{(t)}_r) dR \end{aligned}$

其中
$\frac{R^2}{1 + \phi(R)} = \frac{R^2}{ 1 + \frac{1-\rho}{\rho} \cdot \sqrt{\frac{\nu + \tau^{(t)}_r}{ \tau^{(t)}_r}} \exp \left ( - \frac{\nu}{ 2\tau^{(t)}_r( \tau^{(t)}_r + \nu)} R^2 \right) }$

SE部分的MATLAB代码

sigma2 = 0.0165;
rho = 0.1;
nu = 1 / rho;
lim = inf;
Iteration = 50;
delta = 0.6;   % underdetermined ratio
SE_MSE = zeros(Iteration, 1);
SE_tau2 = zeros(Iteration, 1);SE_MSE(1) = rho * nu;
SE_tau2(1) = sigma2 + 1/delta * SE_MSE(1);
for it = 2: Iterationtau = SE_tau2( it - 1 );f = @(r) r.*r ./ ...(    1   +   (1-rho)./rho .* sqrt( 1 + nu./tau )   .*   exp( max(-60, -0.5 * r .* r .* nu./ tau./(nu + tau) ) )     ) ....* normpdf(r, 0, sqrt(nu + tau))  ;  % -60 is a bound for numerical robustnessSE_MSE(it) = rho * nu - rho * nu^2 / (nu + tau)^2 * integral(f,-lim,lim);SE_tau2(it) = sigma2 + 1/delta * SE_MSE(it);
end

当N=20000时，AMP的MSE性能与SE一致

附录：两个高斯随机变量相乘

$\begin{aligned} & \ \ \ \ \mathcal N(x; \mu_1, \nu_1) \mathcal N(x; \mu_2, \nu_2) \\ &= \frac{1}{\sqrt{2 \pi \nu_1}} \exp \left ( - \frac{ (x - \mu_1)^2 }{2 \nu_1} \right ) \cdot \frac{1}{\sqrt{2 \pi \nu_2}} \exp \left ( - \frac{ (x - \mu_2)^2 }{2 \nu_2} \right ) \\ & = \frac{1}{2 \pi \sqrt{ \nu_1 \nu_2}} \exp \left \{ - \frac{ \nu_2 ( x^2 - 2 \mu_1 x + \mu_1^2 ) }{2 \nu_1 \nu_2} - \frac { \nu_1 ( x^2 - 2 \mu_2 x + \mu_2^2 ) }{2 \nu_1 \nu_2} \right \} \\ & = \frac{1}{2 \pi \sqrt{ \nu_1 \nu_2}} \exp \left \{ - \frac{ (\nu_1 + \nu_2) x^2 - 2 (\mu_1 \nu_2 + \mu_2 \nu_1) x + \mu^2_1 \nu_2 + \mu^2_2 \nu_1 }{2 \nu_1 \nu_2} \right \} \\ & = \frac{1}{2 \pi \sqrt{ \nu_1 \nu_2}} \exp \left \{ - \frac{ x^2 - 2 \frac{(\mu_1 \nu_2 + \mu_2 \nu_1)}{\nu_1 + \nu_2} x + \left ( \frac{(\mu_1 \nu_2 + \mu_2 \nu_1)}{\nu_1 + \nu_2} \right)^2 + \frac{\mu^2_1 \nu_2 + \mu^2_2 \nu_1}{\nu_1 + \nu_2} - \left ( \frac{(\mu_1 \nu_2 + \mu_2 \nu_1)}{\nu_1 + \nu_2} \right)^2 }{2 \frac{\nu_1 \nu_2} { \nu_1 + \nu_2 }} \right \} \\ &= \frac{1}{ \sqrt{2 \pi (\nu_1 + \nu_2)} } \exp \left \{ - \frac{\mu^2_1 \nu_2 + \mu^2_2 \nu_1 - \frac {(\mu_1 \nu_2 + \mu_2 \nu_1)^2}{ \nu_1 + \nu_2} } { 2 \nu_1 \nu_2 } \right \} \cdot \frac{1}{\sqrt{2 \pi \frac{\nu_1 \nu_2} { \nu_1 + \nu_2 } }} \exp \left \{ - \frac{ \left( x- \frac{\mu_1 \nu_2 + \mu_2 \nu_1}{\nu_1 + \nu_2} \right)^2 }{ 2 \frac{\nu_1 \nu_2} { \nu_1 + \nu_2 } } \right \} \\ &= \frac{1}{ \sqrt{2 \pi (\nu_1 + \nu_2)} } \exp \left \{ - \frac{ (\mu_1 - \mu_2)^2 }{2 (\nu_1 + \nu_2)} \right \} \cdot \frac{1}{\sqrt{2 \pi \frac{\nu_1 \nu_2} { \nu_1 + \nu_2 } }} \exp \left \{ - \frac{ \left( x- \frac{\mu_1 \nu_2 + \mu_2 \nu_1}{\nu_1 + \nu_2} \right)^2 }{ 2 \frac{\nu_1 \nu_2} { \nu_1 + \nu_2 } } \right \} \\ \end{aligned}$