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}^{\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.