Source-linked AI summary
State Evolution for Approximate Message Passing with Non-Separable Functions
Raphael Berthier, Andrea Montanari, Phan-Minh Nguyen
TL;DR
Earlier state-evolution theory covered separable nonlinearities, while important AMP applications use non-separable functions. This paper proves state evolution for Lipschitz non-separable functions with Gaussian matrices using conditioning and approximation arguments, introducing LAMP as an intermediate algorithm. The theory is applied to compressed sensing settings, including matrix recovery and image denoising.
Problem
Earlier state-evolution results assumed separable nonlinearities, whereas important AMP applications require non-separable functions.
Method
The proof uses Bolthausen’s conditioning technique, random perturbations, continuity arguments, and the intermediate Long AMP algorithm.
Results
The paper establishes state evolution for Lipschitz continuous non-separable nonlinearities in the stated Gaussian AMP setting.
Takeaways & Limitations
State evolution can characterize AMP behavior for broader non-separable denoisers, including the matrix and image compressed-sensing applications discussed in the paper.
Takeaways & Limitations
The theory assumes Gaussian sensing matrices and the regularity and technical conditions specified in the paper; the image example explicitly uses an i.i.d. Gaussian model.
Abstract
from arXiv · showhide
Given a high-dimensional data matrix ${\boldsymbol A}\in{\mathbb R}^{m\times n}$, Approximate Message Passing (AMP) algorithms construct sequences of vectors ${\boldsymbol u}^t\in{\mathbb R}^n$, ${\boldsymbol v}^t\in{\mathbb R}^m$, indexed by $t\in\{0,1,2\dots\}$ by iteratively applying ${\boldsymbol A}$ or ${\boldsymbol A}^{\sf T}$, and suitable non-linear functions, which depend on the specific application. Special instances of this approach have been developed --among other applications-- for compressed sensing reconstruction, robust regression, Bayesian estimation, low-rank matrix recovery, phase retrieval, and community detection in graphs. For certain classes of random matrices ${\boldsymbol A}$, AMP admits an asymptotically exact description in the high-dimensional limit $m,n\to\infty$, which goes under the name of `state evolution.' Earlier work established state evolution for separable non-linearities (under certain regularity conditions). Nevertheless, empirical work demonstrated several important applications that require non-separable functions. In this paper we generalize state evolution to Lipschitz continuous non-separable nonlinearities, for Gaussian matrices ${\boldsymbol A}$. Our proof makes use of Bolthausen's conditioning technique along with several approximation arguments. In particular, we introduce a modified algorithm (called LAMP for Long AMP) which is of independent interest.
1 Introduction
AMP applies iterative linear operations and nonlinear functions across statistical estimation tasks, with state evolution providing an asymptotically exact high-dimensional characterization. This paper extends that theory to Lipschitz non-separable functions, using perturbation arguments and LAMP, and illustrates the resulting predictions in matrix and image compressed sensing.
- AMP and state evolution: AMP generates vector sequences by repeatedly applying a random data matrix or its transpose and application-specific nonlinear functions.Its applications include compressed sensing, robust regression, Bayesian estimation, low-rank recovery, phase retrieval, and graph community detection.
- AMP and state evolution: State evolution asymptotically characterizes AMP in high dimensions, with the relevant Gaussian covariance computed through a one-dimensional recursion.Earlier results assumed Gaussian matrices with i.i.d. entries and separable, Lipschitz nonlinearities.
- Generalization to non-separable functions: The paper establishes state evolution for Lipschitz continuous nonlinearities that are not necessarily separable.The proof uses Bolthausen’s conditioning technique and addresses possible collinearity among iterates through a random perturbation followed by a vanishing-perturbation argument.
- Long AMP: LAMP provides a streamlined proof route: state evolution is first proved for LAMP, then LAMP is shown to closely approximate the original AMP.The authors identify LAMP as potentially independently useful.
- Applications: In matrix compressed sensing, state evolution predicts lim_n→∞ NMSE(t; n), with the prediction already very accurate for n1 = n2 = 170.The matrix setting uses Gaussian linear measurements and singular value thresholding, a non-separable operator.
- Applications: For image compressed sensing with NLM-AMP, state evolution appears to track simulation results closely as normalized square error evolves across iterations.NLM averages similar patches and uses the residual norm as an effective noise-level measure.
2 Further related work
The paper situates its contribution among extensions of AMP state evolution to broader matrix classes, algorithmic settings, and asymptotic regimes. It distinguishes its general non-separable result from specialized non-asymptotic guarantees.
- AMP state evolution has been extended to subgaussian matrices with separable polynomial nonlinearities and to matrix-valued iterates with fixed width.
- Right-invariant and unitarily invariant matrix models broaden the analyzable matrix classes, although some related analyses use non-rigorous statistical-physics methods.
- The conditioning technique is asymptotic but can potentially be sharpened into non-asymptotic results through central-limit and concentration arguments.
- A separate result establishes non-asymptotic guarantees for compressed sensing with a specialized non-separable sliding-window denoiser, making it not directly comparable here.
3 Main results
The paper states state-evolution results for asymmetric AMP with uniformly Lipschitz, non-separable functions under Gaussian-matrix assumptions. The characterization extends to empirical estimates of algorithm coefficients when those estimates are consistent.
- State evolution defines covariance arrays recursively from the nonlinearities and initialization, with Gaussian variables and Onsager memory terms characterizing the AMP iteration.
- The asymmetric AMP theorem assumes iid Gaussian matrix entries, uniformly Lipschitz nonlinearities, bounded normalized initialization, and finite limiting moment conditions.
- For any fixed iteration and uniformly pseudo-Lipschitz test sequence, the AMP iterates are characterized by the corresponding state-evolution variables.
- The proof reduces the asymmetric theorem to the symmetric case and uses a perturbation argument whose effect vanishes by uniform continuity.
- Empirical estimators of the algorithm coefficients preserve state evolution when they are consistent; divergence-based estimates satisfy the required conditions under uniform pseudo-Lipschitz regularity.
4 Symmetric AMP
The paper develops an analogous state-evolution characterization for symmetric AMP with Gaussian orthogonal ensemble matrices and uniformly Lipschitz nonlinearities. Its recursion defines covariance iterates that describe the symmetric AMP trajectory.
- Symmetric AMP uses a sequence of functions and deterministic initialization, with the matrix sampled from the Gaussian orthogonal ensemble.
- The symmetric setting assumes uniformly Lipschitz nonlinearities, bounded normalized initialization, and finite limiting conditions for the state-evolution recursion.
- State evolution constructs a doubly infinite covariance array recursively from jointly Gaussian variables and the initialization, with Onsager terms defined from its diagonal entries.
- Theorem 3 provides the symmetric AMP state-evolution characterization under these assumptions and positive diagonal covariance entries.
- The symmetric result also permits consistent estimators of the Onsager coefficient in an analogue of the asymmetric corollary.
5 Proof of Theorem 3 (Symmetric AMP)
The proof replaces the original symmetric AMP recursion with LAMP, whose Gaussian-conditioning structure enables state-evolution analysis. A perturbation argument then extends the result from non-degenerate iterates to the general case.
- LAMP construction: LAMP is introduced as a different recursion that satisfies Theorem 3 and approximates the AMP recursion asymptotically.Its initialization is q0 = f0(x0), h1 = Aq0.
- Non-degeneracy: LAMP requires the iterate family q0, q1, . . . , qt−1 to be linearly independent so that QTt−1Qt−1 is invertible.Generic Lipschitz functions can map iterates into a lower-dimensional common subspace, creating degeneracy.
- Conditioning argument: Gaussian conditioning decomposes each new iterate into past iterates and a conditionally understood Gaussian component, while decoupling the matrix from the past.This preserves the asymptotic geometry of the LAMP iterates and supports state evolution.
- Non-degenerate case: Under non-degeneracy, Theorem 7 establishes state evolution for LAMP and transfers it to AMP through an asymptotic approximation.The proof compares AMP and LAMP iterates using uniformly pseudo-Lipschitz test functions.
- Degenerate case: For ill-conditioned iterates, the proof randomly perturbs the functions, proves the perturbed system is non-degenerate, and sends the perturbation to zero.The perturbed state evolution converges to the original one, and triangle-inequality bounds transfer the result back.
6 Proof of Theorem 1 and Corollary 2 (Asymmetric AMP)
The asymmetric AMP theorem is proved by embedding the recursion into the symmetric framework and applying the established state-evolution result. An induction argument then transfers the characterization to practical coefficient estimates.
- Reduction: The asymmetric AMP recursion is reduced to the symmetric setting by constructing a larger Gaussian system whose state evolution matches the original recursion.The proof uses an independent GOE representation and identifies the corresponding iterates.
- Induction: The induction proof controls the difference between the original and state-evolution recursions using Lipschitz continuity and high-probability norm bounds.The Bai–Yin law supplies one of the required operator-norm bounds.
7 Application to general compressed sensing
The general theory applies to compressed sensing with Gaussian sensing matrices and broad Lipschitz denoisers, yielding asymptotic predictions for reconstruction errors. For convex projection denoisers, the paper further derives exponential convergence in the high-dimensional limit.
- General compressed sensing: Compressed sensing reconstructs θ0 from noisy linear measurements using AMP with a sequence of denoisers ηt : Rn → Rn.The problem is underdetermined when m < n, so prior information is encoded through the denoisers.
- General theory: Theorem 14 gives state-evolution predictions for pseudo-Lipschitz observables of the AMP iterates under Gaussian sensing and regularity assumptions.The assumptions include uniformly Lipschitz denoisers and existence of the relevant limits.
- General denoisers: The theory permits denoisers that are fairly general and need not arise from an underlying optimization problem.A smoothed Onsager-coefficient choice can handle denoisers whose divergence is not uniformly Lipschitz.
- Convex projection: For convex projection denoisers, AMP fixed points are stationary points of the associated constrained least-squares problem.The projection onto a closed convex set is a 1-Lipschitz denoiser.
- Geometric interpretation: The convergence threshold is tied to the statistical dimension of the tangent cone, which quantifies the structure captured by the convex constraint.A smaller tangent cone corresponds to a more structured signal.
- Convergence: AMP converges exponentially fast in the high-dimensional limit when m ≥ (1 + η)∆n, with rate ∆n/m.In the noiseless case, ε accuracy requires approximately log(1/ε) / log(m/∆n) iterations.
A.1 Matrix compressed sensing
The matrix compressed-sensing application uses singular value soft thresholding and requires a weak interpretation of its divergence at matrices with repeated singular values.
- Singular value thresholding: The divergence of the singular value soft-thresholding operator can be computed from its singular-value decomposition.The formula is used in the matrix compressed-sensing state-evolution analysis.
- Regularity caveat: The divergence formula is interpreted weakly because it is undefined on the negligible set of matrices with repeated singular values.This is a scope condition on the displayed divergence expression, not on the thresholding operator itself.
A.2 Compressed sensing with images
The simulation approximates state-evolution iterates and expectations numerically for a 170 × 170 Lena image, using Monte Carlo samples with Gaussian noise and NLM denoising.
- The image size is n = 170 × 170.
- The expectation in equation (233) is approximated at each iteration by the mean over 10 Monte Carlo samples.
- Each sample adds Gaussian noise of variance ˆτ 2_t, applies NLM denoising to the Lena image, and computes the squared error.The resulting state evolution is shown in Figure 3.
B Some useful tools
This section collects Gaussian-matrix facts, pseudo-Lipschitz closure properties, and concentration tools used in the analysis.
- For A ∼ GOE(n), the operator norm converges almost surely to 2 as n →∞.
- Stein’s lemma relates Gaussian inner products involving ϕ(Z2) to the divergence of ϕ under the stated covariance and integrability conditions.
- The Gaussian Poincaré inequality provides a concentration bound for continuous, weakly differentiable functions of a standard Gaussian vector.
C Proof of Theorem 15
The proof proceeds by contradiction and subsequence refinement, establishing the required assumptions before applying the theorem to derive the contradiction.
- Assuming a limsup above B_t, the proof selects a subsequence converging to B_t + ε.
- A further subsequence is constructed on which Assumptions (C3), (C5), and (C6) hold with η_s(·) and η_t(·) equal to P_K(·).
- For Assumption (C6), the proof defines functions F_n on the cone of 2 × 2 positive semidefinite matrices using Gaussian expectations.
- Uniform convergence on compact sets and further subsequence refinements allow Theorem 14 to be applied, yielding the desired contradiction.