State evolution for approximate message passing with non-separable functions

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}^{{\textsf T}}$ and suitable nonlinear 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 nonlinearities (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 LoAMP for Long AMP), which is of independent interest.