[Paper Review] Approximate Message Passing algorithms for rotationally invariant matrices

Fan Zhou|arXiv (Cornell University)|Aug 27, 2020

Blind Source Separation Techniques9 citations

TL;DR

This paper rigorously derives Onsager corrections and state evolution for Approximate Message Passing (AMP) algorithms applied to rotationally invariant random matrices, using free cumulants of the matrix's spectral distribution. It establishes that for large signal strengths and non-Gaussian priors, the proposed Bayes-AMP algorithm achieves higher estimation accuracy than standard sample principal components in structured PCA.

ABSTRACT

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 more general AMP algorithms 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. The Onsager corrections and state evolutions in these algorithms are defined by the free cumulants or rectangular free cumulants of the spectral distribution of $W$. Their forms were derived previously by Opper, Çakmak, and Winther using non-rigorous dynamic functional theory techniques, and we provide rigorous proofs. Our motivating application is a Bayes-AMP algorithm for Principal Components Analysis, when there is prior structure for the principal components (PCs) and possibly non-white noise. For sufficiently large signal strengths and any non-Gaussian prior distributions for the PCs, we show that this algorithm provably achieves higher estimation accuracy than the sample PCs.

Motivation & Objective

• To extend the theoretical foundation of Approximate Message Passing (AMP) algorithms beyond i.i.d. or white noise matrices to general rotationally invariant random matrices.
• To rigorously prove the Onsager corrections and state evolution formulas for AMP that depend on the free cumulants of the matrix's spectral distribution.
• To develop a Bayes-AMP algorithm for structured principal components analysis (PCA) with non-white noise and non-Gaussian priors on principal components.
• To demonstrate provable estimation accuracy gains over classical sample PCA under high signal strength and non-Gaussian priors.
• To establish the validity of state evolution and Onsager corrections in the large-dimensional limit for non-i.i.d. matrix ensembles via free probability theory.

Proposed method

• Derives state evolution and Onsager corrections for AMP using free cumulants (or rectangular free cumulants) of the spectral distribution of rotationally invariant matrices.
• Applies tools from free probability theory, including non-crossing partitions and Möbius inversion, to express cumulants in terms of moments.
• Uses conditional expectations in a rectangular probability space to define and analyze rectangular free cumulants for rectangular matrices.
• Establishes analyticity of the R-transform for the spectral distribution using bounds on cumulants and moment growth.
• Applies the AMP framework to structured PCA by incorporating prior knowledge on principal components and non-Gaussian noise models.
• Employs Wasserstein convergence and empirical spectral distribution analysis to justify the asymptotic Gaussianity of AMP iterates.

Experimental results

Research questions

• RQ1 How can Onsager corrections and state evolution be generalized beyond i.i.d. matrices to rotationally invariant ensembles?
• RQ2 What role do free cumulants play in characterizing the state evolution of AMP for non-white, rotationally invariant matrices?
• RQ3 Can a Bayes-AMP algorithm for PCA achieve higher estimation accuracy than sample PCA when the principal components have non-Gaussian priors and the noise is non-white?
• RQ4 Under what conditions does the AMP iterate distribution converge to a Gaussian limit in the large-dimensional regime for such matrices?
• RQ5 How can the theoretical framework of free probability be used to rigorously justify the state evolution of AMP for general rotationally invariant matrix ensembles?

Key findings

• The Onsager corrections and state evolution for AMP are fully characterized by the free cumulants of the spectral distribution of the rotationally invariant matrix W.
• For rectangular matrices, the state evolution depends on rectangular free cumulants, derived via a conditional expectation structure in a non-commutative probability space.
• The R-transform of the spectral distribution is analytic in a disk of radius proportional to (16M)⁻², where M bounds the moments of the spectral measure.
• The Bayes-AMP algorithm for structured PCA provably achieves higher estimation accuracy than sample PCA when the signal strength is sufficiently large and the prior on principal components is non-Gaussian.
• The convergence of the empirical distribution of AMP iterates to a Gaussian limit is rigorously established under the given spectral and moment conditions.
• The framework generalizes previous results for i.i.d. and block-i.i.d. matrices, providing a unified approach for non-Gaussian, correlated matrix ensembles.