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

news2024/9/29 15:34:59

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 \delta(X - \theta)$

那么， $X, R$ 的联合分布可写为
$\begin{aligned} p_{X, R}(X, R) &= \left ((1-\rho) \delta(X) + \rho \delta(X - \theta) \right) \cdot \mathcal N(X;R, \tau^{(t)}_r) \\ &= (1-\rho) \delta(X) \mathcal N(X;R, \tau^{(t)}_r) + \rho \delta(X - \theta) \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 \mathcal N(\theta;R, \tau^{(t)}_r) \\ &= (1-\rho) \mathcal N(R; 0, \tau^{(t)}_r) + \rho \mathcal N(R; \theta, \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 \delta(X - \theta) \mathcal N(X;R, \tau^{(t)}_r) dX \\ &= \frac{1}{p_R(R)} \cdot \rho \cdot \theta \cdot \mathcal N (R; \theta, \tau^{(t)}_r) \\ &= \theta \cdot \frac{ \rho \mathcal N (R; \theta, \tau^{(t)}_r) } {(1-\rho) \mathcal N(R; 0, \tau^{(t)}_r) + \rho \mathcal N(R; \theta, \tau^{(t)}_r) } \\ & = \theta \cdot \frac{ 1 } {1 + \frac{1-\rho}{\rho} \cdot \frac{ \mathcal N(R; 0, \tau^{(t)}_r) }{ \mathcal N(R; \theta, \tau^{(t)}_r) } } \\ &= \theta \cdot \frac{ 1 } {1 + \frac{1-\rho}{\rho} \cdot \exp \left \{\ - \frac{ 2 \theta R - \theta^2 }{2 \tau^{(t)}_r } \right \} } \end{aligned}$

同理可得，
$\mathbb E[X^2|R] = \theta \cdot \mathbb E[X|R]$

因此
$\mathrm{var}[X|R] = \mathbb E[X^2|R] - \mathbb E[X|R]^2$

令 $\psi(R) =\frac{1-\rho}{\rho} \cdot \exp \left \{\ - \frac{ 2 \theta R - \theta^2 }{2 \tau^{(t)}_r } \right \}$ ，则 $\mathrm{var}[X|R]$ 可写为
$\begin{aligned} \mathrm{var}[X|R] &= \theta^2 \frac{1}{ 2+ \psi(R) + \frac{1}{\psi(R) }} \\ &= \theta^2 \frac{1} { 2 + \frac{1- \rho}{\rho} \cdot \exp \left \{\ - \frac{ 2 \theta R - \theta^2 }{2 \tau^{(t)}_r } \right \} +\frac{\rho}{1-\rho} \exp \left \{\ \frac{ 2 \theta R - \theta^2 }{2 \tau^{(t)}_r } \right \}} \end{aligned}$

进一步， $\mathcal E^{(t+1)}$ 可以表征为
$\begin{aligned} \mathcal E^{(t+1)} &= \int p_R(R) \mathrm{var}[X|R] dR \\ &= \theta^2 \int \frac{1} { 2 + \frac{1- \rho}{\rho} \cdot \exp \left \{\ - \frac{ 2 \theta R - \theta^2 }{2 \tau^{(t)}_r } \right \} + \frac{\rho}{1-\rho} \cdot \exp \left \{\ \frac{ 2 \theta R - \theta^2 }{2 \tau^{(t)}_r } \right \}} \left ( (1-\rho) \mathcal N(R; 0, \tau^{(t)}_r) + \rho \mathcal N(R; \theta, \tau^{(t)}_r) \right ) dR \end{aligned}$

总结

当 $X$ 的先验是伯努利时： $\sim (1-\rho) \delta(X) + \rho \delta(X - \theta)$ ，可以把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)} &= \theta^2 \int \frac{1} { 2 + \frac{1- \rho}{\rho} \cdot \exp \left \{\ - \frac{ 2 \theta R - \theta^2 }{2 \tau^{(t)}_r } \right \} + \frac{\rho}{1-\rho} \cdot \exp \left \{\ \frac{ 2 \theta R - \theta^2 }{2 \tau^{(t)}_r } \right \}} \left ( (1-\rho) \mathcal N(R; 0, \tau^{(t)}_r) + \rho \mathcal N(R; \theta, \tau^{(t)}_r) \right ) dR \end{aligned}$

SE部分的MATLAB代码

Iteration = 40;
sigma2 = 0.2632;
rho = 0.1; % sparsity
v_g = 1 / rho; % variance/energy of the non-zero element (prior)
theta = sqrt(v_g);
delta = 0.6; % under-determined ratio
lim = inf;

SE_MSE = zeros(Iteration, 1);
SE_tau2 = zeros(Iteration, 1);

SE_MSE(1) = rho * v_g;
SE_tau2(1) = sigma2 + 1/delta * SE_MSE(1);

bound = 500;
fb = @(b) (b > bound) .* bound + (b < -bound) .* (-bound) + (abs(b) <= bound) .* b;  % bound < f(b) < bound
for it = 2: Iteration
    tau = SE_tau2( it - 1 );
    f = @(r) 1 ./ ...
        ( 2 + (1-rho)./rho .* exp( fb(-0.5 * ( 2 * theta .* r - theta .* theta) / tau ) )  + rho./(1-rho) .* exp( fb(0.5 * ( 2 * theta .* r - theta .* theta) / tau  )) ) ...
        .* ( (1-rho) .* normpdf(r, 0, sqrt(tau)) +  rho .* normpdf(r, theta, sqrt(tau)) );

    SE_MSE(it) = theta^2 * integral(f,-lim,lim);
    SE_tau2(it) = sigma2 + 1/delta * SE_MSE(it);
end