Approximate Message Passing algorithms for rotationally invariant matrices

Approximate Message Passing (AMP) algorithms have seen widespread use across a variety of applications. However, the precise forms for their Onsager corrections and state evolutions depend on properties of the underlying random matrix ensemble, limiting the extent to which AMP algorithms derived for white noise may be applicable to data matrices that arise in practice. In this work, we study a more general AMP algorithm for random matrices $W$ that satisfy orthogonal rotational invariance in law, where $W$ may have a spectral distribution that is different from the semicircle and Marcenko-Pastur laws characteristic of white noise. For a symmetric square matrix $W$, the forms of the Onsager correction and state evolution in this algorithm are defined by the free cumulants of its eigenvalue distribution. These forms were derived previously by Opper, Cakmak, and Winther using non-rigorous dynamic functional theory techniques, and we provide a rigorous proof. For rectangular matrices $W$ that are bi-rotationally invariant in law, we derive a similar AMP algorithm that is defined by the rectangular free cumulants of the singular value distribution of $W$, as introduced by Benaych-Georges. To illustrate the general algorithms, we discuss an application to Principal Components Analysis with prior information for the principal components (PCs). For sufficiently large signal strength and any prior distributions of the PCs that are not mean-zero Gaussian laws, we develop a Bayes-AMP algorithm that provably achieves better estimation accuracy than the sample PCs.