Exercises

$\hat{\mathbf{x}}^{t+1} = \eta_{\mathrm{st}}(\mathbf{A}^{H}\mathbf{r}^t+\hat{\mathbf{x}}^t;\lambda_t)KATEXPLACEHOLDER0END\mathbf{r}^{t+1} = \mathbf{y} - \mathbf{A}\hat{\mathbf{x}}^{t+1} + \frac{1}{\delta}\cdot\frac{\|\hat{\mathbf{x}}^{t+1}\|_0}{N}\cdot\mathbf{r}^t.$$The term$\delta^{-1}|\hat{\mathbf{x}}^{t+1}|_0/N\cdot\mathbf{r}^t$ is the Onsager correction.

$\eta_{\mathrm{st}}(u;\lambda) = u-\lambda$ for $u>\lambda$ (derivative $1$); $=0$ for $|u|\le\lambda$ (derivative $0$); $=u+\lambda$ for $u<-\lambda$ (derivative $1$). Thus $\eta'_{\mathrm{st}}(u;\lambda)=\mathbf{1}\{|u|>\lambda\}$ for $u \ne \pm\lambda$.

SE requires $\eta$ Lipschitz, not differentiable everywhere. Soft- thresholding is 1-Lipschitz, so SE applies. The measure-zero set $\{u = \pm\lambda\}$ does not affect the expectation $\mathbb{E}[\eta'(X+\tau Z;\lambda)]$ because $Z$ has a continuous density.

$\tau_{t+1}^2 = \sigma^2 + \frac{1}{\delta}\!\left[(1-\rho)\,\mathbb{E}[\eta_{\mathrm{st}}(\tau_t Z;\lambda_t)^2] + \rho\,\mathbb{E}[(\eta_{\mathrm{st}}(X+\tau_t Z;\lambda_t)-X)^2]\right]$$with$X \sim \mathcal{N}(0,1)$,$Z \sim \mathcal{N}(0,1)$ independent.

Each expectation reduces to Gaussian integrals involving the standard normal CDF $\Phi$ and PDF $\varphi$. The non-zero term gives, with $s^2=1+\tau_t^2$ and $\lambda=\lambda_t$, $\mathbb{E}[(\eta_{\mathrm{st}}(X+\tau_t Z;\lambda)-X)^2] = (s^2)(1-2\Phi(\lambda/s)+...) + \tau_t^2\cdot(\text{soft-threshold noise term})$, which can be fully expanded using standard Gaussian identities (see Donoho--Maleki--Montanari 2011).

With $\lambda=\alpha\tau$ the MSE of soft-thresholding scales as $\tau^2$ because both the input noise and the threshold scale identically. Precisely, $\mathrm{MSE}(\tau^2;\alpha\tau) = \tau^2 M(\rho,\alpha)$.

The recursion $\tau_{t+1}^2 = \delta^{-1}M(\rho,\alpha)\tau_t^2$ is geometric with ratio $q = M(\rho,\alpha)/\delta$. It converges to zero iff $q<1$, i.e., $M(\rho,\alpha)<\delta$. Optimising $\alpha$ gives $\rho$ below the Donoho--Tanner curve as the recovery region.

$\pi(u) = \frac{\rho\,\mathcal{N}(u;0,\sigma_x^2+\tau^2)}{(1-\rho)\mathcal{N}(u;0,\tau^2)+\rho\,\mathcal{N}(u;0,\sigma_x^2+\tau^2)}.$$

$\eta_{\mathrm{mmse}}(u;\tau^2) = \pi(u)\cdot\frac{\sigma_x^2}{\sigma_x^2+\tau^2}\,u.$$

By the Stein identity, $\eta'_{\mathrm{mmse}}(u;\tau^2) = \mathrm{Var}(X|U=u)/\tau^2$. The posterior variance is $\pi(u)\cdot\frac{\sigma_x^2\tau^2}{\sigma_x^2+\tau^2} + \pi(u)(1-\pi(u))\!\left(\frac{\sigma_x^2 u}{\sigma_x^2+\tau^2}\right)^2$, giving a smooth, bounded $\eta'_{\mathrm{mmse}}$.

Write $U = X + \tau Z$, so conditional on $X$, $U|X \sim \mathcal{N}(X,\tau^2)$. $\mathbb{E}[\eta(U)(U-X)|X] = \tau^2 \mathbb{E}[\eta'(U)|X]$ by Stein's lemma applied to the conditional Gaussian.

Take the expectation over $X$: $\mathbb{E}[\eta(U)(U-X)] = \tau^2 \mathbb{E}[\eta'(U)]$. Since $U - \mathbb{E}[U] = (U-X) + (X - \mathbb{E}[X])$ and a direct computation shows the cross terms equal $\mathrm{Cov}(\eta(U),X)$, rearrangement yields the claimed identity.

import numpy as np
rng = np.random.default_rng(0)
N, M, rho, sigma = 500, 250, 0.1, 1e-2
A = rng.normal(0, 1/np.sqrt(M), size=(M, N))
x = rng.normal(size=N) * (rng.uniform(size=N) < rho)
y = A @ x + sigma * rng.normal(size=M)
xh = np.zeros(N); r = y.copy(); rprev = np.zeros(M)
for t in range(20):
    tau = np.linalg.norm(r)/np.sqrt(M)
    lam = 1.5 * tau
    u = A.T @ r + xh
    xh = np.sign(u) * np.maximum(np.abs(u) - lam, 0)
    b = np.mean(np.abs(u) > lam) * (N/M)
    r = y - A @ xh + b * r
mse = np.linalg.norm(xh - x)**2 / np.linalg.norm(x)**2

Typical run produces NMSE around $10^{-3}$ in 10--15 iterations. If the Onsager term is removed (set $b=0$), MSE plateaus at a much larger value — the ISTA residual.

At an AMP fixed point: $\mathbf{r}^\infty(1-\delta^{-1}\langle\eta'\rangle_\star) = \mathbf{y}-\mathbf{A}\hat{\mathbf{x}}^\infty$. Denote the scaling factor by $c = 1-\delta^{-1}\langle\eta'\rangle_\star$.

LASSO KKT: $\mathbf{A}^{H}(\mathbf{y}-\mathbf{A}\hat{\mathbf{x}}) \in \lambda_{\mathrm{eff}}\partial\|\hat{\mathbf{x}}\|_1$. Substituting AMP's fixed-point gives $c\cdot\mathbf{A}^{H}\mathbf{r}^\infty \in \lambda_{\mathrm{eff}}\partial\|\hat{\mathbf{x}}^\infty\|_1$, while AMP says $\lambda_\star \partial\|\hat{\mathbf{x}}^\infty\|_1 \ni \mathbf{A}^{H}\mathbf{r}^\infty$. Comparing: $\lambda_{\mathrm{eff}} = c\lambda_\star$.

The effective LASSO penalty is a fraction of the AMP threshold because the Onsager term absorbs part of the residual. Practically: if you want LASSO at penalty $\lambda_{\mathrm{eff}}$, set AMP's threshold to $\lambda_\star = \lambda_{\mathrm{eff}}/c$, where $c$ is recomputed from the current Onsager coefficient.

For a Lipschitz denoiser, the input distribution $X+\tau Z$ becomes more spread-out as $\tau$ increases, so the best achievable MSE cannot decrease. A formal proof couples two copies $\tau_1 < \tau_2$ and uses the data-processing inequality.

Since $\Psi$ is monotone non-decreasing, $\tau_{t+1}^2 = \Psi(\tau_t^2)$ is a monotone dynamical system: either $\tau_1^2 \ge \tau_0^2$ (and the sequence is non-decreasing) or $\tau_1^2 \le \tau_0^2$ (and the sequence is non-increasing). In either case, by monotone-convergence, $\tau_t^2$ converges to a fixed point.

F = np.fft.fft(np.eye(N)) / np.sqrt(N)
idx = rng.choice(N, M, replace=False)
A = np.sqrt(N/M) * F[idx]  # random rows of DFT

Apply the same AMP loop as in Exercise 7.

The residual norm typically grows super-linearly after a few iterations and overflows. Damping with $\theta \approx 0.3$ stabilises the iteration but the MSE trajectory no longer matches the SE prediction.

For each $(\rho,\alpha)$, compute $M(\rho,\alpha)$ numerically by Monte Carlo or Gaussian quadrature. Minimise over $\alpha$. The curve $\rho^\star(\delta)$ is the largest $\rho$ such that the minimum equals $\delta$.

The resulting curve is monotone, concave, starts at $\rho^\star(0)=0$ and approaches $\rho^\star(1) = 1$. For $\delta=0.5$ one finds $\rho^\star \approx 0.19$, in agreement with Donoho--Tanner combinatorial-geometry predictions.

Isolate column $i$: $\mathbf{A} = \mathbf{A}_{-i} + \mathbf{a}_i \mathbf{e}_i^T$. The pseudo-data at index $i$ is $u_i = (\mathbf{A}^{H}\mathbf{r})_i + \hat{x}_i = \mathbf{a}_i^H \mathbf{r} + \hat{x}_i$. Linearising $\mathbf{r} = \mathbf{r}_{-i} + \mathbf{a}_i(\hat{x}_i - \hat{x}_i^{\mathrm{cavity}})$ introduces a self-feedback term of size $\mathbf{a}_i^H \mathbf{a}_i \approx 1/\delta$.

Tracking the Taylor expansion to first order, one finds that the bias in the denoiser input due to self-feedback is $\delta^{-1}\eta'(u_i^{t-1})\,r_i^{t-1}$ at each coordinate. Aggregating: the residual at step $t$ should have added $\delta^{-1}\langle\eta'\rangle\,\mathbf{r}^{t-1}$ to cancel this bias — precisely the Onsager term.

The damped update $\hat{x}^{t+1} = (1-\theta)\hat{x}^t + \theta\eta(u^t)$ is an affine combination of the old iterate and the raw denoiser output. If we include $\hat{x}^t$ as a side input to the denoiser, the iteration looks like a single AMP step with a history-dependent $\tilde\eta$.

SE demands that the denoiser output at iteration $t$ depend only on $u^t$ (a Gaussian conditionally on history). Damped AMP's effective denoiser explicitly depends on $\hat{x}^t$, introducing correlation across iterations that changes the covariance structure of the pseudo-data and breaks the scalar recursion.

In the continuous-time limit (many damped steps per ``true'' step) the iterate traces a gradient flow on the AMP fixed-point equation. SE re-emerges in this limit but with a different time-scaling.

Compute $p(x|u) \propto p_X(x)\mathcal{N}(u;x,\tau^2)$ for $x \ge 0$: the posterior is a mixture of a point mass at $0$ and a truncated Gaussian. The posterior mean is available in closed form via the error function.

Evaluate the posterior mean in closed form; its derivative is (posterior variance)/$\tau^2$ by the Stein identity. Replace the soft-threshold in the AMP loop with this denoiser and compute the Onsager coefficient as the empirical mean of the derivative.

The MSE trajectory should closely match the scalar $\tau_{t+1}^2 = \sigma^2 + \delta^{-1}\mathbb{E}[(\eta_{\mathrm{mmse}}(X+\tau_t Z)-X)^2]$. For $\rho=0.1$, $\delta=0.5$, this MSE is typically 2--3 dB below the L1-AMP MSE.

With $\eta(u)=u$: $\mathbb{E}[(\eta(X+\tau Z)-X)^2] = \tau^2\mathbb{E}[Z^2] = \tau^2$. So $\tau_{t+1}^2 = \sigma^2 + \delta^{-1}\tau_t^2$.

The fixed point is $\tau_\star^2 = \sigma^2/(1-1/\delta) = \sigma^2\delta/(\delta-1)$ when $\delta > 1$. For $\delta \le 1$ the recursion diverges — the identity denoiser does not exploit any prior information, so AMP reduces to an underdetermined pseudo-inverse which cannot average down the noise.

For any $\tau$, $\mathbb{E}[(\eta(X+\tau Z)-X)^2]$ is minimised over all measurable $\eta$ by the posterior mean $\eta_{\mathrm{mmse}}(u;\tau^2)$.

Denote by $\tau_{\mathrm{B}}^2$ the fixed point achieved with $\eta_{\mathrm{mmse}}$ and by $\tau_{\mathrm{L}}^2$ the fixed point with optimal soft-thresholding. Since at every $\tau$ the MMSE curve is below the L1 curve, the fixed points satisfy $\tau_{\mathrm{B}}^2 \le \tau_{\mathrm{L}}^2$, with strict inequality whenever the prior is non-Gaussian.

The replica-symmetric cavity calculation yields an equation identical to the Bayes-AMP state-evolution fixed point (Reeves-- Pfister 2019). Thus the AMP MSE equals the conjectured (and rigorously proved in the i.i.d. Gaussian case) information- theoretic MMSE limit.

MLP with input $(u,\tau)$, two hidden layers of 64 units each, ReLU activations, scalar output. Train with MSE loss on synthetic pairs $(X+\tau Z, X)$.

Use the D-AMP pseudo-code: apply the network to the pseudo-data vector, estimate divergence by the Monte-Carlo rule, update the residual with the estimated divergence.

A universal approximator of the posterior mean would make D-AMP equivalent to Bayes-AMP. In practice, (i) the network is trained over a finite set of noise levels, introducing bias for intermediate $\tau$; (ii) the MC divergence estimate has variance; (iii) training imbalance favours large signal components. Each introduces a fixed- point gap.