A Simple Derivation of AMP and its State Evolution via First-Order Cancellation

TL;DR

This paper offers a heuristic, simplified derivation of the AMP algorithm and its state evolution for linear regression, providing new insights into why AMP works well with large i.i.d. matrices.

Contribution

It introduces a first-order cancellation approach to derive AMP and its state evolution, simplifying understanding beyond traditional belief propagation methods.

Findings

Heuristic derivation of AMP using first-order cancellation

Insights into the importance of large i.i.d. matrices for AMP

Simplified explanation of the state evolution formalism

Abstract

We consider the linear regression problem, where the goal is to recover the vector $\boldsymbol{x}\in\mathbb{R}^n$ from measurements $\boldsymbol{y}=\boldsymbol{A}\boldsymbol{x}+\boldsymbol{w}\in\mathbb{R}^m$ under known matrix $\boldsymbol{A}$ and unknown noise $\boldsymbol{w}$. For large i.i.d. sub-Gaussian $\boldsymbol{A}$, the approximate message passing (AMP) algorithm is precisely analyzable through a state-evolution (SE) formalism, which furthermore shows that AMP is Bayes optimal in certain regimes. The rigorous SE proof, however, is long and complicated. And, although the AMP algorithm can be derived as an approximation of loop belief propagation (LBP), this viewpoint provides little insight into why large i.i.d. $\boldsymbol{A}$ matrices are important for AMP, and why AMP has a state evolution. In this work, we provide a heuristic derivation of AMP and its state evolution, based on the idea of "first-order cancellation," that provides insights missing from the LBP derivation while being much shorter than the rigorous SE proof.

Tables

Table 1. Table 1: Numerical evidence that std ​ ( ℰ n ( t ) ) ​ n std superscript subscript ℰ 𝑛 𝑡 𝑛 \text{std}(\mathcal{E}_{n}^{\scriptscriptstyle(t)})\sqrt{n} and std ​ ( τ r , n ( t ) ) ​ n std superscript subscript 𝜏 𝑟 𝑛 𝑡 𝑛 \text{std}(\tau_{r,n}^{\scriptscriptstyle(t)})\sqrt{n} are approximately constant with n 𝑛 n , implying that std ​ ( ℰ n ( t ) ) std superscript subscript ℰ 𝑛 𝑡 \text{std}(\mathcal{E}_{n}^{\scriptscriptstyle(t)}) and std ​ ( τ r , n ( t ) ) std superscript subscript 𝜏 𝑟 𝑛 𝑡 \text{std}(\tau_{r,n}^{\scriptscriptstyle(t)}) scale as approximately 1 / n 1 𝑛 1/\sqrt{n} .

	$n$	1000	3000	10000	30000
MMSE denoiser	$\text{std}\left({ℰ}_n^{\left(29\right)}\right)\sqrt{n}$	0.0011	0.0008	0.0010	0.0010
	$\text{std}\left({\tau }_{r,n}^{\left(29\right)}\right)\sqrt{n}$	0.0026	0.0017	0.0022	0.0020
Soft Threshold	$\text{std}\left({ℰ}_n^{\left(29\right)}\right)\sqrt{n}$	0.0055	0.0044	0.0037	0.0040
	$\text{std}\left({\tau }_{r,n}^{\left(29\right)}\right)\sqrt{n}$	0.0118	0.0098	0.0106	0.0093

Equations

$$\mathbf{y}=\mathbf{A}\mathbf{x}+\mathbf{w}\in {\mathbb{R}}^m,$$

$${\mathbf{v}}^{(t)}$$

$${\mathbf{x}}^{(t+1)}$$

$${\mathbf{\mu }}^{(t)}=0,$$

$${\mathbf{\mu }}^{(t)}$$

$${\mathbf{e}}^{(t)}\triangleq {\mathbf{r}}^{(t)}-\mathbf{x}.$$

$${\mathbf{e}}^{(t)}$$

$$=(\mathbf{I}-{\mathbf{A}}^{\top }\mathbf{A}){\mathbf{x}}^{(t)}+{\mathbf{A}}^{\top }(\mathbf{A}\mathbf{x}+\mathbf{w}+{\mathbf{\mu }}^{(t)})-\mathbf{x}$$

$$=(\mathbf{I}-{\mathbf{A}}^{\top }\mathbf{A}){\mathbf{x}}^{(t)}-(\mathbf{I}-{\mathbf{A}}^{\top }\mathbf{A})\mathbf{x}+{\mathbf{A}}^{\top }(\mathbf{w}+{\mathbf{\mu }}^{(t)}).$$

$$[(\mathbf{I}-{\mathbf{A}}^{\top }\mathbf{A}){\mathbf{x}}^{(t)}{]}_j$$

$$\displaystyle=\Big{(}1-\sum_{i=1}^{m}a_{ij}^{2}\Big{)}x^{\scriptscriptstyle(t)}_{j}-\sum_{i}a_{ij}\sum_{l\neq j}a_{il}x_{l}^{\scriptscriptstyle(t)}$$

$$=-\sum _i{a}_{ij}\sum _{l\ne j}{a}_{il}{x}_l^{(t)}$$

$$[(\mathbf{I}-{\mathbf{A}}^{\top }\mathbf{A}){\mathbf{x}}^{(t)}{]}_j$$

$$\displaystyle=-\sum_{i}a_{ij}\sum_{l\neq j}a_{il}\eta^{\scriptscriptstyle(t-1)}\Big{(}\underbrace{x_{l}^{\scriptscriptstyle(t-1)}+\sum_{k\neq i}a_{kl}v_{k}^{\scriptscriptstyle(t-1)}}_{\displaystyle\triangleq r_{il}^{\scriptscriptstyle(t-1)}}+a_{il}v_{i}^{\scriptscriptstyle(t-1)}\Big{)},$$

$${\eta }^{(t-1)}(r_{il}^{(t-1)}+a_{il}{v}_i^{(t-1)})$$

$$={\eta }^{(t-1)}(r_{il}^{(t-1)})+a_{il}{v}_i^{(t-1)}{\eta }^{(t-1)\prime }(r_{il}^{(t-1)})+\underset{O(1/m)}{\underbrace{\frac{1}{2}{a}_{il}^2(v_i^{(t-1)}{)}^2{\eta }^{(t-1)\prime \prime }(r_{il}^{(t-1)})+\text{H.O.T.}}},$$

$$[(\mathbf{I}-{\mathbf{A}}^{\top }\mathbf{A}){\mathbf{x}}^{(t)}{]}_j$$

$$=-\sum _i{a}_{ij}\sum _{l\ne j}{a}_{il}{\eta }^{(t-1)}(r_{il}^{(t-1)})-\frac{1}{m}\sum _i{a}_{ij}{v}_i^{(t-1)}\sum _{l\ne j}{\eta }^{(t-1)\prime }(r_{il}^{(t-1)})$$

$$[(\mathbf{I}-{\mathbf{A}}^{\top }\mathbf{A})\mathbf{x}{]}_j$$

$$e_j^{(t)}$$

$$\displaystyle=\sum_{i}a_{ij}\sum_{l\neq j}a_{il}\big{[}x_{l}-\eta^{\scriptscriptstyle(t-1)}(r_{il}^{\scriptscriptstyle(t-1)})\big{]}+\sum_{i}a_{ij}w_{i}+\sum_{i}a_{ij}\Big{[}\mu_{i}^{\scriptscriptstyle(t)}-v_{i}^{\scriptscriptstyle(t-1)}\frac{1}{m}\sum_{l\neq j}\eta^{{\scriptscriptstyle(t-1)}\prime}(r_{il}^{\scriptscriptstyle(t-1)})\Big{]}.$$

$$\displaystyle\sum_{i}a_{ij}\Big{[}\frac{1}{m}v_{i}^{\scriptscriptstyle(t-1)}\sum_{l}\eta^{{\scriptscriptstyle(t-1)}\prime}(r_{l}^{\scriptscriptstyle(t-1)})-\frac{1}{m}v_{i}^{\scriptscriptstyle(t-1)}\sum_{l\neq j}\eta^{{\scriptscriptstyle(t-1)}\prime}(r_{il}^{\scriptscriptstyle(t-1)})\Big{]}$$

$$\displaystyle=\frac{1}{m}\sum_{i}a_{ij}v_{i}^{\scriptscriptstyle(t-1)}\Big{[}\eta^{{\scriptscriptstyle(t-1)}\prime}(r_{j}^{\scriptscriptstyle(t-1)})+\sum_{l\neq j}\Big{(}\eta^{{\scriptscriptstyle(t-1)}\prime}(r_{l}^{\scriptscriptstyle(t-1)})-\eta^{{\scriptscriptstyle(t-1)}\prime}(r_{il}^{\scriptscriptstyle(t-1)})\Big{)}\Big{]}$$

$$\displaystyle\approx\frac{1}{m}\sum_{i}a_{ij}v_{i}^{\scriptscriptstyle(t-1)}\Big{[}\eta^{{\scriptscriptstyle(t-1)}\prime}(r_{j}^{\scriptscriptstyle(t-1)})+\sum_{l\neq j}a_{il}v_{i}^{\scriptscriptstyle(t-1)}\eta^{{\scriptscriptstyle(t-1)}\prime\prime}(r_{il}^{\scriptscriptstyle(t-1)})\Big{]},$$

$${\eta }^{(t-1)\prime }(r_l^{(t-1)})={\eta }^{(t-1)\prime }(r_{il}^{(t-1)}+a_{il}{v}_i^{(t-1)})$$

$$\frac{1}{m}\sum _{i=1}^m\underset{O(1/\sqrt{m})}{\underbrace{a_{ij}{v}_i^{(t-1)}{\eta }^{(t-1)\prime }(r_j^{(t-1)})}}=O(1/\sqrt{m})$$

$$\frac{1}{m}\sum _{i=1}^m\underset{O(1/\sqrt{m})}{\underbrace{a_{ij}(v_i^{(t-1)}{)}^2\underset{O(1)}{\underbrace{\sum _{l\ne j}{a}_{il}{\eta }^{(t-1)\prime \prime }(r_{il}^{(t-1)})}}}}=O(1/\sqrt{m}),$$

$$e_j^{(t)}$$

$$\displaystyle\operatorname{\mathbb{E}}\bigg{[}\sum_{i}A_{ij}\sum_{l\neq j}A_{il}\epsilon_{il}^{\scriptscriptstyle(t)}\bigg{]}$$

$$\displaystyle\operatorname{\mathbb{E}}\Big{[}\Big{(}\sum_{i}A_{ij}\sum_{l\neq j}A_{il}\epsilon_{il}^{\scriptscriptstyle(t)}\Big{)}^{2}\Big{]}$$

$${ϵ}_{il}^{(t)}$$

$$=\underset{\triangleq {ϵ}_l^{(t)}}{\underbrace{x_l-{\eta }^{(t-1)}(r_l^{(t-1)})}}+\underset{O(1/\sqrt{m})}{\underbrace{a_{il}{v}_i^{(t-1)}{\eta }^{(t-1)\prime }(r_{il}^{(t-1)})+O(1/m)}},$$

$$\frac{1}{m^2}\sum _i\sum _{l\ne j}({ϵ}_{il}^{(t)}{)}^2$$

$$=\underset{O(1)}{\underbrace{\frac{n}{m}\frac{1}{n}\sum _{l=1}^n({ϵ}_l^{(t)}{)}^2}}-\underset{O(1/m)}{\underbrace{\frac{1}{m}({ϵ}_j^{(t)}{)}^2}}\approx {\delta }^{-1}{\mathcal{E}}^{(t)},$$

$${\mathcal{E}}^{(t)}\triangleq \underset{n\to \infty }{\lim }\frac{1}{n}\sum _{l=1}^n({ϵ}_l^{(t)}{)}^2$$

$$\displaystyle\operatorname{\mathbb{E}}\Big{[}\sum_{i}A_{ij}w_{i}\Big{]}$$

$$\displaystyle\operatorname{\mathbb{E}}\bigg{[}\Big{(}\sum_{i}A_{ij}w_{i}\Big{)}^{2}\bigg{]}$$

$${\tau }_w\triangleq \underset{m\to \infty }{\lim }\frac{1}{m}\sum _{i=1}^m{w}_i^2.$$

Full text

1 Introduction

We consider the standard linear regression problem, where the goal is to recover the vector $\boldsymbol{x}\in{\mathbb{R}}^{n}$ from measurements

[TABLE]

where $\boldsymbol{\boldsymbol{A}}$ is a known matrix and $\boldsymbol{w}$ is an unknown disturbance. With high-dimensional random $\boldsymbol{\boldsymbol{A}}$, the approximate message passing (AMP) algorithm [1] remains one of the most celebrated and best understood iterative algorithms. In particular, when the entries of $\boldsymbol{\boldsymbol{A}}$ are drawn i.i.d. from a sub-Gaussian distribution and $m,n\rightarrow\infty$ with $m/n\rightarrow\delta\in(0,\infty)$, ensemble behaviors of AMP, such as the per-iteration mean-squared error (MSE), can be perfectly predicted using a state evolution (SE) formalism [2].111See also [3] for an earlier proof of AMP’s state evolution under i.i.d. Gaussian entries. Furthermore, the SE formalism shows that, in certain regimes, AMP’s MSE converges to the minimum MSE as predicted by the replica method [3, 2], which has been shown to coincide with the minimum MSE for linear regression under i.i.d. Gaussian $\boldsymbol{\boldsymbol{A}}$ [4, 5] as $m,n\rightarrow\infty$ with $m/n\rightarrow\delta\in(0,\infty)$. More recently, it has been proven that the state-evolution accurately characterizes AMP’s behavior for large but finite $m,n$ [6].

The rigorous SE proofs in [2, 3, 6], however, are long and complicated, and thus remain out of reach for many readers. And, although the AMP algorithm can be heuristically derived from an approximation of loop belief propagation (LBP) [7, 8], as outlined in [9], [3, App.A], and [10], and expectation propagation (EP) [11, 12], as outlined in [13], the LBP/EP perspective is lacking in several respects. First, LBP and EP are heuristic, making it surprising that further approximations of these approaches can be optimal. Second, the LBP and EP derivations provide little insight into why large i.i.d. $\boldsymbol{\boldsymbol{A}}$ matrices are important for AMP. For the LBP/EP derivations, it suffices that the entries of $\boldsymbol{A}$ have “roughly the same magnitude $O(1/\sqrt{m})$” [10, 13], suggesting that structured matrices (e.g., DCT, Hadamard, Fourier) should work as well as i.i.d. random ones. But, in practice, AMP often diverges with such structured matrices, and the reasons why are not evident from the LBP/EP viewpoint. Third, the LBP and EP derivations do not explain why AMP obeys a scalar state evolution when $\boldsymbol{A}$ is large and i.i.d. sub-Gaussian, nor do they describe how the estimation error variance can be predicted at each iteration.

In this work, we propose a heuristic derivation of AMP and its MSE state evolution that uses the simple idea of “first-order cancellation.” This derivation provides insights missing from the LBP and EP derivations, while being much more accessible than the rigorous SE proofs.

2 Problem Setup

In our treatment of the linear regression problem (1), $\boldsymbol{y}=[y_{1},\dots,y_{m}]^{\top}$, $\boldsymbol{x}=[x_{1},\dots,x_{n}]^{\top}$, and $\boldsymbol{w}=[w_{1},\dots,w_{m}]^{\top}$ are deterministic vectors and $\boldsymbol{\boldsymbol{A}}\in{\mathbb{R}}^{m\times n}$ is a deterministic matrix. Importantly, however, we assume that the components $\{a_{ij}\}$ of $\boldsymbol{\boldsymbol{A}}$ are realizations of i.i.d. Bernoulli222With additional work, our derivation can be extended to i.i.d. Gaussian $A_{ij}$, but doing so lengthens the derivation and provides little additional insight. random variables $A_{ij}\in\pm\frac{1}{\sqrt{m}}$ that are drawn independently of $\boldsymbol{x}$ and $\boldsymbol{w}$. Our model for $\boldsymbol{\boldsymbol{A}}$ is a special case of that considered in [2].

Throughout, we will focus on the following large-system limit.

Definition 1**.**

The “large system limit” is defined as $m,n\rightarrow\infty$ with $m/n\rightarrow\delta$ for some fixed sampling ratio $\delta\in(0,\infty)$.

We will assume that the components of $\boldsymbol{x}$, $\boldsymbol{w}$, and $\boldsymbol{y}$ scale as $O(1)$ in the large-system limit.

We consider a family of algorithms that, starting with $\boldsymbol{x}^{\scriptscriptstyle(0)}=\boldsymbol{0}$, iterates the following over iteration index $t=0,1,2,\dots$:

[TABLE]

where $\boldsymbol{\mu}^{\scriptscriptstyle(t)}$ is a correction term and $\eta^{\scriptscriptstyle(t)}(\cdot)$ is a component-wise separable function. That is, $[\eta^{\scriptscriptstyle(t)}(\boldsymbol{r})]_{j}=\eta^{\scriptscriptstyle(t)}(r_{j})~{}\forall j$, where, with some abuse of notation, we use the same notation for the function $\eta^{\scriptscriptstyle(t)}:{\mathbb{R}}^{n}\rightarrow{\mathbb{R}}^{n}$ and its component maps $\eta^{\scriptscriptstyle(t)}:{\mathbb{R}}\rightarrow{\mathbb{R}}$. The quantity $\boldsymbol{x}^{\scriptscriptstyle(t)}$ is iteration-$t$ estimate of the unknown vector $\boldsymbol{x}$. We refer to $\eta^{\scriptscriptstyle(t)}(\cdot)$ as a “denoiser” for reasons that will become clear in the sequel. For technical reasons, we will assume that $\eta^{\scriptscriptstyle(t)}(\cdot)$ is a polynomial function of bounded degree, similar to the assumption in [2].

The classical iterative shrinkage/thresholding (IST) algorithm [14] uses no correction, i.e.,

[TABLE]

for all iterations $t$, whereas the AMP algorithm [1] uses the “Onsager” correction

[TABLE]

initialized with $\boldsymbol{\mu}^{\scriptscriptstyle(0)}=\boldsymbol{0}$. In (4), $\eta^{{\scriptscriptstyle(t)}\prime}$ refers to the derivative of $\eta^{\scriptscriptstyle(t)}$. Our goal is to analyze the effect of $\boldsymbol{\mu}^{\scriptscriptstyle(t)}$ on the behavior of algorithm (2) in the large-system limit, and in particular to understand how and why the Onsager correction (4) is a good choice. To do this, we will analyze the errors on $\boldsymbol{r}^{\scriptscriptstyle(t)}$ and $\boldsymbol{x}^{\scriptscriptstyle(t)}$ in (2) and drop terms that vanish in the large-system limit.

It has been shown [3, 10] that IST has an predictable and desirable behavior in the case that $\boldsymbol{A}$ is a large i.i.d. Gaussian matrix that is re-drawn at each iteration $t$ (with a corresponding update of $\boldsymbol{y}$). But this desirable behavior vanishes in the practical case that $\boldsymbol{A}$ is fixed over the iterations. In some sense, the goal of the Onsager correction (4) is to restore this desirable behavior when $\boldsymbol{A}$ is fixed over the iterations.

3 AMP Derivation

We will now analyze the error $\boldsymbol{e}^{\scriptscriptstyle(t)}$ on the input to the denoiser $\boldsymbol{r}^{\scriptscriptstyle(t)}$, i.e.,

[TABLE]

From (2) and (5) we have that

[TABLE]

Let us examine the $j$th component of $\boldsymbol{e}^{\scriptscriptstyle(t)}$ when $t\geq 1$. We have that

[TABLE]

since $a_{ij}^{2}=1/m~{}\forall ij$. Continuing,

[TABLE]

where $r_{il}^{\scriptscriptstyle(t-1)}$ omits the direct contribution of $a_{il}$ from $r_{l}^{\scriptscriptstyle(t-1)}$ and thus is only weakly dependent on $\{a_{ij}\}_{j=1}^{n}$. We formalize this weak dependence through Assumption 1, which is admittedly an approximation. In fact, the approximate nature of Assumption 1 is the main reason that our derivation is heuristic.

Assumption 1**.**

The matrix entry $a_{ij}$ is a realization of an equiprobable Bernoulli random variable $A_{ij}\in\pm\frac{1}{\sqrt{m}}$, where $\{A_{ij}\}$ are mutually independent and, $A_{ij}$ is independent of $\{r_{il}^{\scriptscriptstyle(t-1)}\}_{l=1}^{n}$, $\{x_{l}\}_{l=1}^{n}$, and $\{w_{k}\}_{k=1}^{m}$.

We now say a few words about Assumption 1. The assumption that $\{A_{ij}\}$ are i.i.d. and independent of $\boldsymbol{x}$ and $\boldsymbol{w}$ is rather common in the compressive sensing literature. For example, these assumptions are used in the rigorous AMP analyses [2, 3, 6]. The assumption of Bernoulli $A_{ij}$ is a bit stronger, but it is not critical, in that our analysis could be extended to handle other sub-Gaussian distributions on $A_{ij}$ with additional steps. Doing so, however, would complicate the derivation without providing much additional insight, and so we have elected not to take this path. The assumption that $A_{ij}$ is independent of $\{r_{il}^{\scriptscriptstyle(t-1)}\}_{l=1}^{n}$ is far stronger. In reality, there is a weak dependence between these quantities, but properly accounting for it seems to require completely different analysis methods, such as those in [2, 3, 6].

Assumption 1 will be used often when analyzing summations, as in the following lemma.

Lemma 1**.**

Consider the quantity $z_{i}=\sum_{j=1}^{n}a_{ij}u_{j}$, where $a_{ij}$ are realizations of i.i.d. random variables $A_{ij}$ with zero mean and $\operatorname{\mathbb{E}}[A_{ij}^{2}]=1/m$. If $\{A_{ij}\}$ are drawn independently of $\{u_{j}\}$, and $\{u_{j}\}$ scale as $O(1)$ in the large-system limit, then $z_{i}$ also scales as $O(1)$.

Proof.

First, note that $z_{i}$ is a realization of the random variable $Z_{i}\triangleq\sum_{j=1}^{n}A_{ij}u_{j}$. Furthermore, \operatorname{\mathbb{E}}[Z_{i}^{2}]=\operatorname{\mathbb{E}}[(\sum_{j=1}^{n}A_{ij}u_{j})^{2}]=\sum_{j=1}^{n}\sum_{l=1}^{n}\operatorname{\mathbb{E}}[A_{ij}A_{il}]u_{j}u_{l}=\frac{1}{m}\sum_{j=1}^{n}u_{j}^{2}={\color[rgb]{0,0,0}\frac{n}{m}}\frac{1}{n}\sum_{j=1}^{n}u_{j}^{2}, since $\operatorname{\mathbb{E}}[A_{ij}A_{il}]=1/m$ if $j=l$ and $\operatorname{\mathbb{E}}[A_{ij}A_{il}]=\operatorname{\mathbb{E}}[A_{ij}]\operatorname{\mathbb{E}}[A_{il}]=0$ if $j\neq l$. Clearly $m/n$ and $\frac{1}{n}\sum_{j=1}^{n}u_{j}^{2}$ are both $O(1)$ in the large-system limit. Thus we conclude that $\operatorname{\mathbb{E}}[Z_{i}^{2}]$ is $O(1)$. Finally, since $z_{i}$ is a realization of a random variable $Z_{i}$ whose second moment is $O(1)$, we conclude that $z_{i}$ scales as $O(1)$ in the large-system limit. ∎

We now emphasize the importance of the independence properties on $\{A_{ij}\}$ used in Assumption 1 and Lemma 1. For example, consider the case where $A_{ij}=\operatorname{sgn}(u_{j})/\sqrt{m}$, so that $\operatorname{\mathbb{E}}[A_{ij}^{2}]=1/m$ (as in Lemma 1) but where $A_{ij}$ depends on $\{u_{j}\}$. In this case, $z_{i}=\sum_{j=1}^{n}a_{ij}u_{j}=\frac{1}{\sqrt{m}}\sum_{j=1}^{n}|u_{j}|$, which scales as $O(\sqrt{m})$ in the large-system limit, and thus is fundamentally different from the $O(1)$ scaling observed in Lemma 1.

The distinction between $\{a_{ij}\}$ being independent of other quantities, versus $\{a_{ij}\}$ simply being the right size, is one of the major differences between our derivation and other simple derivations based on LBP and EP. For example, both [10] and [13] require only that “$\{a_{ij}\}$ are roughly the same magnitude $O(1/\sqrt{m})$,” which suggests that properly normalized structured matrices (e.g, Hadamard, Fourier, Discrete Cosine Transform, Radon Transform) should work well with AMP. But when used to recover natural images/signals with these matrices, AMP often diverges. And the LBP/EP viewpoint does not explain why.

In the sequel, we will make use of the following lemma, whose proof is postponed because it is a bit long and does not provide much insight.

Lemma 2**.**

Under Assumption 1 and the Onsager choice of $\boldsymbol{\mu}^{\scriptscriptstyle(t)}$ from (4), the elements of $\boldsymbol{v}^{\scriptscriptstyle(t)}$, $\boldsymbol{r}^{\scriptscriptstyle(t)}$, $\boldsymbol{x}^{\scriptscriptstyle(t)}$, and $\boldsymbol{\mu}^{\scriptscriptstyle(t)}$ scale as $O(1)$ in the large-system limit for all iterations $t$.

Proof.

See the appendix. ∎

We now perform a Taylor series expansion of the $\eta^{\scriptscriptstyle(t-1)}$ term in (13) about $r_{il}^{\scriptscriptstyle(t-1)}$:

[TABLE]

where the $O(1/m)$ scaling follows from the fact that $a_{il}^{2}=1/m~{}\forall il$, that both $v_{i}^{\scriptscriptstyle(t-1)}$ and $r_{il}^{\scriptscriptstyle(t-1)}$ scale as $O(1)$ via Lemma 1, and $\eta^{\scriptscriptstyle(t-1)}(\cdot)$ is polynomial of bounded degree, which implies that $\eta^{{\scriptscriptstyle(t-1)}\prime\prime}(r_{il}^{\scriptscriptstyle(t-1)})$ also scales as $O(1)$. For similar reasons, the second term in (14) scales as $O(1/\sqrt{m})$. We will ignore the $O(1/m)$ term in (14) since it vanishes relative to the $O(1/\sqrt{m})$ component in the large-system limit. Thus we have

[TABLE]

using $a_{il}^{2}=1/m~{}\forall il$.

Similar to (11), we have

[TABLE]

which, combined with (8) and (16), yields

[TABLE]

We are now in a position to observe the principal mechanism of AMP. As we argue below (using the central limit theorem), the first and second terms in (19) behave like realizations of zero-mean Gaussians in the large-system limit, because $\{a_{il}\}$ are realizations of i.i.d. zero-mean random variables $\{A_{il}\}$ that are independent of $x_{l}$, $w_{i}$, and $\{r_{il}^{\scriptscriptstyle(t-1)}\}$ under Assumption 1. But the same cannot be said in general about the third term in (19), because $v_{i}^{\scriptscriptstyle(t-1)}$ is strongly coupled to $a_{ij}$. Consequently, the denoiser input-error $e_{j}^{\scriptscriptstyle(t)}$ is difficult to characterize for general $\mu_{i}^{\scriptscriptstyle(t)}$.

With AMP’s choice of $\mu_{i}^{\scriptscriptstyle(t)}$, however, the 3rd term in (19) vanishes in the large-system limit. In particular, with the Onsager choice (4), the 3rd term in (19) takes the form

[TABLE]

where for the last step we used the Taylor-series expansion

[TABLE]

and dropped the $O(1/m)$ term, since it will vanish relative to the $a_{il}v_{i}^{\scriptscriptstyle(t-1)}\eta^{{\scriptscriptstyle(t-1)}\prime\prime}(r_{il}^{\scriptscriptstyle(t-1)})$ term in the large-system limit. Looking at (21), the first term is

[TABLE]

since $a_{ij}\in\pm 1/\sqrt{m}$ and $v_{i}^{\scriptscriptstyle(t-1)}\eta^{{\scriptscriptstyle(t-1)}\prime}(r_{j}^{\scriptscriptstyle(t-1)})$ is $O(1)$ due to Lemma 1. Thus the first term in (21) will vanish in the large-system limit. The second term in (21) is

[TABLE]

which will also vanish in the large-system limit. The $O(1)$ scaling in (24) follows from Lemma 1 under Assumption 1, and the $O(1/\sqrt{m})$ scaling follows from the fact that $a_{il}\in\pm 1/\sqrt{m}$ and $(v_{i}^{\scriptscriptstyle(t-1)})^{2}=O(1)$.

Thus, for large $m$ and the AMP choice of $\mu_{i}^{\scriptscriptstyle(t)}$, equation (19) becomes

[TABLE]

Recall that, under Assumption 1, $a_{ij}$ is a realization of equiprobable $A_{ij}\in\pm\frac{1}{\sqrt{m}}$ that is independent of $\{A_{il}\}_{l\neq j}$, $\{x_{l}\}_{l=1}^{n}$, and $\{r_{il}^{\scriptscriptstyle(t-1)}\}_{l=1}^{n}$. This implies that $A_{ij}$ is also independent of $\epsilon_{il}^{\scriptscriptstyle(t)}=x_{l}-\eta^{\scriptscriptstyle(t-1)}(r_{il}^{\scriptscriptstyle(t-1)})$ for any $l$. Thus we can apply the central limit theorem to say that, for any fixed $\{\epsilon_{il}^{\scriptscriptstyle(t)}\}$, the first term converges to a Gaussian with mean and variance

[TABLE]

From the Taylor expansion (14), we have

[TABLE]

where the $O(1/\sqrt{m})$ scaling follows from the facts that $a_{il}\in\pm 1/\sqrt{m}$ and $v_{i}^{\scriptscriptstyle(t-1)}\eta^{{\scriptscriptstyle(t-1)}\prime}(r_{il}^{\scriptscriptstyle(t-1)})$ is $O(1)$. Notice that $\epsilon_{l}^{\scriptscriptstyle(t)}$ is the denoiser output error. Because both $x_{l}$ and $r_{l}^{\scriptscriptstyle(t-1)}$ are $O(1)$, it follows that $\epsilon_{l}^{\scriptscriptstyle(t)}$ is also $O(1)$. Because the $O(1/\sqrt{m})$ term in (29) vanishes in the large-system limit, we see that (27) becomes

[TABLE]

where

[TABLE]

is the average squared error on the denoiser output $\boldsymbol{x}^{\scriptscriptstyle(t)}$. We have thus deduced that, in the large-system limit, the first term in (25) behaves like a zero-mean Gaussian with variance $\delta^{-1}\mathcal{E}^{\scriptscriptstyle(t)}$. For the second term in (25), we can again use the central limit theorem to say that, for any fixed $\{w_{i}\}$, the second term converges to a Gaussian with mean and variance

[TABLE]

where $\tau_{w}$ denotes the empirical second moment of the noise:

[TABLE]

To summarize, with AMP’s choice of $\boldsymbol{\mu}^{\scriptscriptstyle(t)}$ from (4), the $j$th component of the denoiser input-error behaves like

[TABLE]

in the large-system limit, where $\mathcal{N}(\mu,\sigma^{2})$ denotes a Gaussian random variable with mean $\mu$ and variance $\sigma^{2}$. With other choices of $\boldsymbol{\mu}^{\scriptscriptstyle(t)}$ (e.g., IST’s choice of $\boldsymbol{\mu}^{\scriptscriptstyle(t)}=\boldsymbol{0}~{}\forall t$), it is difficult to characterize the denoiser input-error $\boldsymbol{e}^{\scriptscriptstyle(t)}$ and in general it will not be Gaussian.

4 AMP State Evolution

In Section 3, we used Assumption 1 to argue that the AMP algorithm yields a denoiser input-error $\boldsymbol{e}^{\scriptscriptstyle(t)}$ whose components are $\mathcal{N}(0,\tau_{r}^{\scriptscriptstyle(t)})$ in the large system limit. Here, $\tau_{r}^{\scriptscriptstyle(t)}=\delta^{-1}\mathcal{E}^{\scriptscriptstyle(t)}+\tau_{w}$ where $\mathcal{E}^{\scriptscriptstyle(t)}$ is the average squared-error at the denoiser output in the large-system limit.

Recalling the definition of $\mathcal{E}^{\scriptscriptstyle(t)}$ from (32), we can write

[TABLE]

where $X$ is a scalar random variable defined from the empirical distribution

[TABLE]

with $\delta(\cdot)$ denoting the Dirac delta function. Thus we can argue that, in the large-system limit,

[TABLE]

where $X$ now is distributed according to the $n\rightarrow\infty$ limit of the empirical distribution. Combining (40) with the update equation for $\tau_{r}^{\scriptscriptstyle(t)}$ gives the following recursion for $t=0,1,2,\dots$:

[TABLE]

initialized with $\mathcal{E}^{\scriptscriptstyle(0)}=\operatorname{\mathbb{E}}[X^{2}]$. The recursion (41) is known as AMP’s “state evolution” for the mean-squared error [1, 3, 2].

The reason that we call $\eta^{\scriptscriptstyle(t)}(\cdot)$ a “denoiser” should now be clear. To minimize the mean-squared error $\mathcal{E}^{\scriptscriptstyle(t+1)}$, the function $\eta^{\scriptscriptstyle(t)}(\cdot)$ should remove as much of the noise from its input as possible. The smaller that $\mathcal{E}^{\scriptscriptstyle(t+1)}$ is, the smaller the input-noise variance $\tau_{r}^{\scriptscriptstyle(t+1)}$ will be during the next iteration.

5 AMP Variance Estimation

For best performance, the iteration-$t$ denoiser $\eta^{\scriptscriptstyle(t)}(\cdot)$ should be designed in accordance with the iteration-$t$ input noise variance $\tau_{r}^{\scriptscriptstyle(t)}$. With the AMP algorithm, there is an easy way to estimate the value of $\tau_{r}^{\scriptscriptstyle(t)}$ at each iteration $t$ from the $\boldsymbol{v}^{\scriptscriptstyle(t)}$ vector, i.e., $\tau_{r}^{\scriptscriptstyle(t)}\approx\|\boldsymbol{v}^{\scriptscriptstyle(t)}\|^{2}/m$ [10]. Below, we explain this approach using arguments similar to those used above.

Equation (58) shows that

[TABLE]

Ignoring the $O(1/m)$ term and plugging in AMP’s choice of $\mu_{i}^{\scriptscriptstyle(t)}$ from (4) yields

[TABLE]

where we used the Taylor series (22) in the second step and $a_{il}\in\pm 1/\sqrt{m}$ to justify the $O(1/\sqrt{m})$ scaling. Since the last term in (44) is the scaled average of $O(1/\sqrt{m})$ terms, with $O(1)$ scaling, the entire term is $O(1/\sqrt{m})$. We can thus drop it since it will vanish relative to the others in the large-system limit. Doing this and plugging in $\boldsymbol{y}=\boldsymbol{Ax}+\boldsymbol{w}$ yields

[TABLE]

recalling the definition of $\epsilon_{il}^{\scriptscriptstyle(t)}$ from (25). Squaring the result and averaging over $i$ yields

[TABLE]

We now examine the components of (46) in the large-system limit. By definition, the first term in (46) converges to $\tau_{w}$. By the law of large numbers, the second term converges to

[TABLE]

since $\operatorname{\mathbb{E}}[A_{il}A_{ij}]=1/m$ when $l=j$ and $\operatorname{\mathbb{E}}[A_{il}A_{ij}]=0$ when $l\neq j$. Using the relationship between $\epsilon_{il}^{\scriptscriptstyle(t)}$ and $\epsilon_{l}^{\scriptscriptstyle(t)}$ from (29), it can be seen that

[TABLE]

where $m$ is implicitly a function of $n$ because $m/n=O(1)$. In summary,

[TABLE]

which shows that $\tau_{r}^{\scriptscriptstyle(t)}$ is well estimated by $\|\boldsymbol{v}^{\scriptscriptstyle(t)}\|^{2}/m$ in the large-system limit.

6 Numerical Experiments

We now present numerical experiments that demonstrate the AMP behaviors discussed above. In all experiments, we used a sampling ratio of $\delta=0.5$, $\{A_{ij}\}$ drawn i.i.d. zero-mean Gaussian with variance $1/m$, $\{x_{j}\}$ drawn i.i.d. from the Bernoulli-Gaussian distribution with sparsity rate $\beta=0.1$ (i.e., $p_{X}(x_{j})=(1-\beta)\delta(x_{j})+\beta\mathcal{N}(x_{j};0,1)~{}\forall j$, where $\delta(\cdot)$ denotes the Dirac delta distribution), and $\{w_{i}\}$ drawn i.i.d. zero-mean Gaussian with variance $\beta 10^{-\text{SNRdB}/10}$ and $\text{SNRdB}=20$, so that $\operatorname{\mathbb{E}}[\|\boldsymbol{Ax}\|^{2}]/\operatorname{\mathbb{E}}[\|\boldsymbol{w}\|^{2}]\approx 20$ dB. We experimented with two denoisers: the MMSE denoiser $\eta^{\scriptscriptstyle(t)}(r_{j})=\operatorname{\mathbb{E}}[X\,|\,r_{j}=X+\mathcal{N}(0,\tau_{r}^{\scriptscriptstyle(t)})]$ and the soft-thresholding denoiser $\eta^{\scriptscriptstyle(t)}(r_{j})=\operatorname{sgn}(r_{j})\max\{0,|r_{j}|-\alpha\sqrt{\tau_{r}^{\scriptscriptstyle(t)}}\}$ with $\alpha=1.14$, which is the minimax choice, i.e., the value of $\alpha$ that minimizes the maximum MSE over all $0.1$-sparse signals (see [10] for more details). With the soft-thresholding denoiser, AMP solves the LASSO problem “$\operatorname*{arg\,min}_{\boldsymbol{x}}\{\frac{1}{2}\|\boldsymbol{y}-\boldsymbol{Ax}\|^{2}+\lambda\|\boldsymbol{x}\|_{1}\}$” for some value of $\lambda$ [1, 10].

Figures 1-6 plot finite-dimensional versions of the denoiser output MSE $\mathcal{E}^{\scriptscriptstyle(t)}$ and the denoiser input-error variance $\tau_{r}^{\scriptscriptstyle(t)}$ versus iteration $t$ for both the AMP algorithm (2) and the AMP state evolution (41). For the AMP algorithm, the iteration-$t$ denoiser output MSE was computed as $\mathcal{E}_{n}^{\scriptscriptstyle(t)}=\frac{1}{n}\sum_{j=1}^{n}(x_{j}-x_{j}^{\scriptscriptstyle(t)})^{2}$ and the denoiser input-error variance was computed as $\tau_{r,n}^{\scriptscriptstyle(t)}=\|\boldsymbol{v}^{\scriptscriptstyle(t)}\|^{2}/m$, where the subscript $n$ indicates the dimensional dependence of these quantities. For the AMP state evolution, the denoiser output MSE was computed as

[TABLE]

with the expectation evaluated using the $n$-term empirical distribution for $X$, and the iteration-$t$ denoiser input-error variance was computed as $\tau_{r,n}^{\scriptscriptstyle(t)}=\delta^{-1}\mathcal{E}_{n}^{\scriptscriptstyle(t)}+\tau_{w,n}$ using the empirical noise variance $\tau_{w,n}=\frac{1}{m}\sum_{i=1}^{m}w_{i}^{2}$. Each figure plots the empirical mean and standard deviation over $T\in\{100,1000,10000\}$ random draws of $\boldsymbol{\boldsymbol{A}}$ for a single fixed draw of $\boldsymbol{x}$ and $\boldsymbol{w}$.

Figure 1 shows the results for the MMSE denoiser at dimension $n=300$. The figure shows a good, but not great, agreement between the state evolution and average AMP quantities, where the average was computed over $T=10000$ realizations of $\boldsymbol{\boldsymbol{A}}$. Furthermore, the error bars in Fig. 1, which show the empirical standard deviation over the $T$ realizations, indicate that there was considerable dependence of the trajectories $\{\mathcal{E}_{n}^{\scriptscriptstyle(t)}\}_{t=1}^{30}$ and $\{\tau_{r,n}^{\scriptscriptstyle(t)}\}_{t=1}^{30}$ on the realization of $\boldsymbol{\boldsymbol{A}}$ when $n=300$. Figure 2 and 3 plot the same quantities for dimensions $n=3000$ and $n=30000$, respectively. These figures show that, as the dimension $n$ increases, the agreement between the state-evolution and average AMP trajectories improves and the standard deviation of the AMP trajectories decreases. Table 1 suggests that the standard deviation decreases proportional to $1/\sqrt{n}$. Figure 3 shows that, when $n=30000$, the trajectories $\{\mathcal{E}_{n}^{\scriptscriptstyle(t)}\}_{t=1}^{30}$ and $\{\tau_{r,n}^{\scriptscriptstyle(t)}\}_{t=1}^{30}$ are nearly invariant to changes in $\boldsymbol{\boldsymbol{A}}$.

Figures 4-6 are similar to Figures 1-3, except that they show the results for the soft-thresholding denoiser. As expected, the use of the soft-thresholding denoiser results in larger MSEs than the MMSE denoiser. But, otherwise, the trends are the same: the agreement between the state-evolution and average AMP trajectories increases with the dimension $n$, and the standard deviation of the AMP trajectories decreases proportional to $1/\sqrt{n}$.

Note that, because the state evolution was computed using the empirical distributions of $\{x_{j}\}_{j=1}^{n}$ and $\{w_{i}\}_{i=1}^{m}$, which change from one figure to the next (e.g., as $n$ and $m$ change), the state evolution trajectories vary across Figures 1-6.

Finally, to give evidence that the denoiser input error $\{e_{j}^{\scriptscriptstyle(t)}\}$ is approximately Gaussian, we show quantile-quantile (QQ) plots in Figures 7 and 8 at iteration $t=5$ and dimension $n=3000$ for the MMSE and soft-thresholding denoisers, respectively. The figures show that the quantiles of $\{e_{j}^{\scriptscriptstyle(t)}\}$ are very close to those of a zero-mean Gaussian random variable. Although not shown here, QQ plots at other iterations $t$ look similar, and the QQ plots become more linear (i.e., $\{e_{j}^{\scriptscriptstyle(t)}\}$ looks more Gaussian) as $n$ grows larger.

7 Conclusion

For the linear regression problem, we presented a simple derivation of AMP and its MSE state-evolution based on the idea of “first-order cancellation.” In particular, our derivation writes the linear transform of the denoiser output, $\boldsymbol{\boldsymbol{A}}\boldsymbol{x}^{\scriptscriptstyle(t)}$, as the sum of a term that is weakly dependent on the previous iteration and another term that is strongly dependent but canceled by AMP’s Onsager correction term in the large-system limit. Our derivation provides insights that are missing from the usual loopy belief-propagation derivation of AMP, while being much more accessible than Bayati et al.’s rigorous analysis of AMP.

Here we establish Lemma 2, which says that the elements of $\boldsymbol{v}^{\scriptscriptstyle(t)},\boldsymbol{r}^{\scriptscriptstyle(t)},\boldsymbol{x}^{\scriptscriptstyle(t)}$, and $\boldsymbol{\mu}^{\scriptscriptstyle(t)}$ scale as $O(1)$ in the large-system limit under Assumption 1. We do this by induction.

From the initialization $\boldsymbol{x}^{\scriptscriptstyle(0)}=\boldsymbol{0}=\boldsymbol{\mu}^{\scriptscriptstyle(0)}$, we have that $\boldsymbol{v}^{\scriptscriptstyle(0)}=\boldsymbol{y}$, whose elements are $O(1)$, and we have $\boldsymbol{r}^{\scriptscriptstyle(0)}=\boldsymbol{A}^{\top}\boldsymbol{y}=\boldsymbol{A}^{\top}\boldsymbol{Ax}+\boldsymbol{A}^{\top}\boldsymbol{w}$. Examining the $j$th entry, we see that

[TABLE]

since $a_{ij}^{2}=1/m~{}\forall ij$ and where the $O(1)$ scalings follow from Lemma 1 under Assumption 1. Thus the elements of $\boldsymbol{r}^{\scriptscriptstyle(0)}$ are $O(1)$. Because $x_{j}^{\scriptscriptstyle(1)}=\eta^{\scriptscriptstyle(0)}(r_{j}^{\scriptscriptstyle(0)})$, the elements of $\boldsymbol{x}^{\scriptscriptstyle(1)}$ are also $O(1)$. And from (4),

[TABLE]

where $\frac{n}{m}$, $v_{i}^{\scriptscriptstyle(0)}$, and $\frac{1}{n}\sum_{j=1}^{n}\eta^{{\scriptscriptstyle(0)}\prime}(r_{j}^{\scriptscriptstyle(0)})$ are all $O(1)$, implying that the elements of $\boldsymbol{\mu}^{\scriptscriptstyle(1)}$ are $O(1)$.

Now, suppose that the elements of $\boldsymbol{r}^{\scriptscriptstyle(t-1)},\boldsymbol{v}^{\scriptscriptstyle(t-1)},\boldsymbol{x}^{\scriptscriptstyle(t)},\boldsymbol{\mu}^{\scriptscriptstyle(t)}$ are all $O(1)$, which we know occurs when $t=1$. Then from (2) we have that

[TABLE]

where $r_{il}^{\scriptscriptstyle(t-1)}$ is only weakly dependent on $\{a_{ij}\}_{j=1}^{n}$, leading us to invoke Assumption 1. The Taylor expansion (14) (which is admissible since, under the induction hypothesis, the elements of $\boldsymbol{v}^{\scriptscriptstyle(t-1)}$ and $\boldsymbol{r}^{\scriptscriptstyle(t-1)}$ scale as $O(1)$) then yields

[TABLE]

where we used the fact that $a_{il}^{2}=1/m~{}\forall il$. In (58), the $O(1)$ scaling of the first and last sums follows from Lemma 1 under Assumption 1. Thus the elements of $\boldsymbol{v}^{\scriptscriptstyle(t)}$ are $O(1)$.

Next, we establish that the elements of $\boldsymbol{r}^{\scriptscriptstyle(t)}$ are $O(1)$ when the elements of $\boldsymbol{v}^{\scriptscriptstyle(t-1)},\boldsymbol{x}^{\scriptscriptstyle(t)}$ and $\boldsymbol{\mu}^{\scriptscriptstyle(t)}$ are. From (5)-(8) we have

[TABLE]

where we previously established that the elements in $\boldsymbol{r}^{\scriptscriptstyle(0)}$ are $O(1)$. As for the remaining term in (60), we have from (13) and (4) that

[TABLE]

where, for the last step, we applied the Taylor expansion (14). Applying $a_{il}^{2}=1/m~{}\forall il$ and rearranging, we get

[TABLE]

where the $O(1)$ term follows from the fact that $a_{ij}\in\pm 1/\sqrt{m}$. For the 2nd-to-last term in (66), we have

[TABLE]

which follows from Lemma 1 under Assumption 1. For the remaining term in (66), we apply the Taylor expansion (22) and rearrange terms as follows:

[TABLE]

where the $O(1)$ scaling of the sum follows from Lemma 1 under Assumption 1. The final $O(1/\sqrt{m})$ scaling in (70) follows since each of the three terms in (69) is the average of $O(1/\sqrt{m})$ terms, due to $a_{ij}\in\pm 1/\sqrt{m}$. In conclusion, we have established that the elements of $\boldsymbol{r}^{\scriptscriptstyle(t)}$ are $O(1)$ when the elements of $\boldsymbol{v}^{\scriptscriptstyle(t-1)},\boldsymbol{x}^{\scriptscriptstyle(t)}$ and $\boldsymbol{\mu}^{\scriptscriptstyle(t)}$ are.

To complete the induction proof, we need to establish that the elements of $\boldsymbol{x}^{\scriptscriptstyle(t+1)}$ and $\boldsymbol{\mu}^{\scriptscriptstyle(t+1)}$ are $O(1)$ when those of $\boldsymbol{v}^{\scriptscriptstyle(t)}$ and $\boldsymbol{r}^{\scriptscriptstyle(t)}$ are. But this follows straightforwardly from (2) and (4), i.e.,

[TABLE]

since $n/m$ is $O(1)$ in the large-system limit.

In summary, we have established Lemma 2, which says the elements of $\boldsymbol{v}^{\scriptscriptstyle(t)},\boldsymbol{r}^{\scriptscriptstyle(t)},\boldsymbol{x}^{\scriptscriptstyle(t)}$, and $\boldsymbol{\mu}^{\scriptscriptstyle(t)}$ scale as $O(1)$ in the large-system limit under Assumption 1.