PCA Initialization for Approximate Message Passing in Rotationally Invariant Models

1 Introduction . We consider the problem of estimating a rank-1 signal from a noisy data matrix . In the square symmetric case , the data matrix is modeled as X = α n u∗u∗T +W ∈ Rn×n , ( 1.1 ) where u∗ ∈ Rn is the unknown rank-1 signal , W ∈ Rn×n is a symmetric noise matrix , and α > 0 captures the signal-to-noise ratio ( SNR ) . In the rectangular case , we observe the data matrix X = α m u∗v∗T +W ∈ Rm×n , ( 1.2 ) where u∗ ∈ Rm and v∗ ∈ Rn are the unknown signals , andW ∈ Rm×n is a rectangular noise matrix . A natural estimator of the signal in the symmetric case is the principal eigenvector ofX ( singular vectors , in the rectangular case ) . The performance of this principal component analysis ( PCA ) estimator and , more generally , the behavior of eigenvalues and eigenvectors of models like ( 1.1 ) - ( 1.2 ) has been widely studied in statistics [ 27 , 49 ] and random matrix theory [ 2 , 3 , 10 , 11 , 15 , 22 , 30 ] . If u∗ , v∗ are unstructured ( e.g. , they are uniformly distributed on a sphere ) , then it is not generally possible to improve on the PCA estimator . However , in a broad range of applications , the unknown signals have some underlying structure , e.g. , they may be sparse , their entries may belong to a certain 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . set , or they may be modelled using a prior distribution . Examples of structured matrix estimation problems include sparse PCA [ 17 , 28 , 61 ] , non-negative PCA [ 32 , 43 ] , community detection under the stochastic block model [ 1 , 16 , 45 ] , and group synchronization [ 50 ] . Since PCA is ill-equipped to capture the structure of the signal , we aim to improve on it using a family of iterative algorithms known as approximate message passing ( AMP ) . AMP algorithms have two particularly attractive features : ( i ) they can be tailored to take advantage of prior information on the structure of the signal ; and ( ii ) under suitable model assumptions , their performance in the high-dimensional limit is precisely characterized by a succinct deterministic recursion called state evolution [ 7 , 13 , 26 ] . AMP algorithms have been applied to a wide range of inference problems : estimation in linear models [ 8 , 7 , 19 , 31 , 39 ] , generalized linear models [ 5 , 37 , 38 , 42 , 51 , 53 , 55 ] , and low-rank matrix estimation with Gaussian noise [ 6 , 17 , 23 , 29 , 34 , 44 ] . The survey [ 21 ] provides a unified description of AMP for these applications . Using the state evolution analysis , it has been proved that AMP achieves Bayes-optimal performance in some Gaussian models [ 17 , 18 , 44 ] , and a bold conjecture from statistical physics posits that AMP is optimal among polynomial-time algorithms . We study rank-1 matrix estimation in the setting where the noise matrixW is rotationally invariant . This is a much milder assumption than W being Gaussian : it only imposes that the orthogonal matrices in the spectral decomposition ofW are uniformly random , and allows for arbitrary eigenvalues/singular values . Hence , W can capture a more complex correlation structure , which is typical in applications . For the models ( 1.1 ) - ( 1.2 ) with rotationally invariant noise , AMP algorithms were derived in [ 14 , 48 ] and generalized in [ 20 ] . In particular , the AMP algorithm of [ 20 ] for the problem ( 1.1 ) produces estimates ut ∈ Rn as follows : ut = ut ( f t−1 ) , f t = Xut − t∑ i=1 bt , iu i , t ≥ 2 . ( 1.3 ) The iteration is initialized with a pilot estimate u1 . We can interpret ( 1.3 ) as a generalized power method . Recall that the power method approximates the principal eigenvector ofX using the iterative updates ūt = Xūt−1/‖Xūt−1‖ . For each t , the function ut can be chosen to exploit any structural information known about the signal ( e.g. , sparsity ) . The “ memory ” coefficients { bt,1 , . . . , bt , t } have a specific form to ensure that the iterates ( f t , ut+1 ) have desirable statistical properties captured by state evolution . A rigorous state evolution result for the iteration ( 1.3 ) is established in [ 20 ] , but the algorithm and its analysis require an initialization u1 that is correlated with the unknown signal and independent of the noiseW . In practice , one typically does not have access to such an initialization . Main contribution . In this paper , we propose an AMP algorithm initialized via the PCA estimator , namely , the principal eigenvector ofX for the square case ( 1.1 ) and the left singular vector ofX for the rectangular case ( 1.2 ) . Our main technical contribution is a state evolution result for this AMP algorithm , which gives a rigorous characterization of its performance in the high-dimensional limit . The challenge is that , as the PCA initialization depends on the noise matrixW , one can not apply the state evolution machinery of [ 20 ] . To circumvent this issue , our key idea is to construct and analyze a two-phase artificial AMP algorithm . In the first phase , the artificial AMP performs a power method approaching the PCA estimator ; and in the second phase , it mimics the behavior of the true AMP . We remark that the artificial AMP only serves as a proof technique . Thus , we can initialize it with a vector correlated with the signal u∗ and independent of the noise matrix W , which allows us to analyze it using the existing state evolution result . Our analysis is tight in the sense that our AMP algorithm can be initialized with PCA whenever the PCA estimate has strictly positive correlation with the signal . This requires showing that , when PCA is effective , the state evolution of the first phase of the artificial AMP has a unique fixed point . To obtain such a result , we exploit free probability tools developed in [ 10 , 11 ] . The agreement between the practical performance of AMP and the theoretical predictions of state evolution is demonstrated via numerical results for different spectral distributions ofW . Our simulations also show that the performance of AMP—as well as its ability to improve upon the PCA initialization— crucially depends on the choice of the denoising functions ut in the algorithm . Thus , the design of a Bayes-optimal AMP remains an exciting avenue for future research . Related work . The asymptotic Bayes-optimal error for low-rank matrix estimation has been precisely characterized for Gaussian noise [ 4 , 33 ] , but remains an open problem for rotationally invariant noise . An AMP algorithm with PCA initialization was proposed in [ 44 ] for the Gaussian setting , and it was shown to be Bayes-optimal for some signal priors . A recent paper by Zhong et al . [ 60 ] shows how AMP with PCA initialization can be used for estimating the top-k principal components in applications such as high-dimensional genomics datasets . The authors use an empirical Bayes method to determine a joint prior distribution for the k principal components , and assuming a Gaussian noise model , employ an AMP algorithm tailored to the prior to improve the PCA estimates of the principal components . Both our AMP algorithm and proof technique differ significantly from those for Gaussian noise . WhenW is Gaussian , the PCA initialization affects only the first iteration of AMP . In contrast , for more general noise distributions , the AMP algorithm and its associated state evolution require a correction term at every iteration to account for the PCA initialization . This is due to the fact that , while AMP has a single memory term in the Gaussian case , more general noise distributions lead to a more involved memory structure , as in ( 1.3 ) . As regards the proof technique , the argument of [ 44 ] consists of decoupling the PCA estimate from the bulk of the spectrum ofX . In contrast , our approach is based on a two-phase artificial AMP algorithm . This technique has proved successful in the context of generalized linear models [ 41 , 42 ] , albeit for Gaussian measurements . Other extensions of AMP beyond the Gaussian setting include Orthogonal AMP [ 36 , 56 ] , Vector AMP [ 24 , 25 , 52 , 54 ] , convolutional AMP [ 57 ] and Memory AMP [ 35 ] . These algorithms have been derived specifically for linear or generalized linear models , and extending them ( with a practical initialization method ) to low-rank matrix estimation is an interesting research direction . Finally , we mention the recent independent work of Zhong et al . [ 59 ] , which appeared after the original submission of our paper . This work generalizes AMP with PCA initialization to the problem of estimating rank-k matrices in rotationally invariant noise , for k ≥ 1 . We remark that , in order to prove a state evolution result for AMP initialized with PCA , in [ 59 ] it is assumed that the signal strength is sufficiently large . In contrast , our result holds for any signal strength such that the PCA method is effective , but we require the free cumulants of the noise matrix to be non-negative . We also note that , when the signal strength is large , the assumption on the free cumulants can be automatically satisfied ( see the footnote on page 6 ) . 2 Preliminaries . Notation and definitions . Given a ∈ R , we define ( a ) + = max ( a , 0 ) . Given two integers i ≤ j , we define [ i , j ] = { i , . . . , j } . If i > j , then [ i , j ] denotes the empty set ; products over the empty set are taken to be 1 . Given a vector x ∈ Rn , we denote by ‖x‖ its Euclidean norm and by 〈x〉 its empirical mean , i.e. , 〈x〉 = 1n ∑n i=1 xi . The empirical distribution of x = ( x1 , . . . , xn ) T is given by 1 n ∑n i=1 δxi , where δxi denotes a Dirac delta mass on xi . The notation x W−→ X denotes convergence of the empirical distribution of x to the random variable X in Wasserstein distance at all orders . Given a symmetric square matrix A ∈ Rn×n , we denote by λ1 ( A ) ≥ λ2 ( A ) ≥ . . . ≥ λn ( A ) its eigenvalues sorted in decreasing order . Given a rectangular matrix A ∈ Rm×n , with m < n , we denote by σ1 ( A ) ≥ σ2 ( A ) ≥ . . . σm ( A ) its singular values sorted in decreasing order . Rank-1 estimation – Symmetric square matrices . Consider the problem of estimating the signal u∗ ∈ Rn from the data matrix X in ( 1.1 ) . We assume that W is rotationally invariant in law , i.e. , W = OTΛO , where Λ = diag ( λ ) is a diagonal matrix containing the eigenvalues ofW andO is a Haar orthogonal matrix independent of Λ . As n→∞ , we assume that the empirical distributions of λ and u∗ satisfy λ W−→ Λ and u∗ W−→ U∗ , ( 2.1 ) where Λ and U∗ represent the limiting spectral distribution of the noise and the prior on the signal , respectively . We take ‖u‖ = √ n so that E { U2∗ } = limn→∞ 1n‖u ∗‖2 = 1 . We assume that the moment E { U2+ε∗ } < ∞ for some ε > 0 . We also assume that Λ has compact support , and denote by b the supremum of this support . We denote by { κk } k≥1 the free cumulants corresponding to the moments { mk } k≥1 of the empirical eigenvalue distribution of X excluding its largest eigenvalue , i.e. , mk = 1n ∑n i=2 λi ( X ) k ( for details , see ( A.1 ) - ( A.2 ) in Appendix A ) . The assumption ( 2.1 ) implies that , as n→∞ , mk → m∞k = E { Λk } and κk → κ∞k , where { m∞k } k≥1 and { κ∞k } k≥1 are respectively moments and free cumulants of Λ. PCA – Symmetric square matrices . Let uPCA be the principal eigenvector ofX , and define αs = 1/G ( b+ ) , where G ( z ) = E { ( z − Λ ) −1 } is the Cauchy transform of Λ , and G ( b+ ) = limz→b+ G ( z ) . Then , for α > αs , λ1 ( X ) a.s.−→ G−1 ( 1/α ) and λ2 ( X ) a.s.−→ b , where G−1 is the inverse of G ; see Theorem 2.1 in [ 10 ] . Furthermore , Theorem 2.2 in [ 10 ] gives that , for α > αs , 〈uPCA , u∗〉2 n a.s.−→ ρ2α = −1 α2G′ ( G−1 ( 1/α ) ) > 0 . ( 2.2 ) In words , above the spectral threshold αs , the principal eigenvalue of X escapes the bulk of the spectrum and its associated eigenvector becomes strictly correlated with the signal u∗ . Rank-1 estimation – Rectangular matrices . Consider now the problem of estimating the signals u∗ ∈ Rm and v∗ ∈ Rn given the rectangular data matrix X in ( 1.2 ) . Without loss of generality , we assume that m ≤ n ( if m > n , one can just exchange the role of u∗ and v∗ and consider XT in place of X ) . We assume that W is bi-rotationally invariant in law , i.e. , W = OTΛQ , where Λ = diag ( λ ) is a m× n diagonal matrix containing the singular values ofW , andO , Q are Haar orthogonal matrices independent of one another and also of Λ . As n→∞ , we assume that λ W−→ Λ , u∗ W−→ U∗ , v∗ W−→ V∗ and m/n → γ , for some constant γ ∈ ( 0 , 1 ] . We take ‖u‖ = √ m and ‖v‖ = √ n so that E { U2∗ } = E { V 2∗ } = 1 . As before , b < ∞ is the supremum of the compact support of Λ , and U∗ , V∗ are assumed to have finite ( 2 + ε ) -th moment for some ε > 0 . To analyze PCA using the framework in [ 11 ] , we also assume that the entries of u∗ and v∗ are i.i.d. , and their law has zero mean and satisfies a log-Sobolev inequality . We denote by { κ2k } k≥1 the rectangular free cumulants associated to the even moments { m2k } k≥1 , with m2k = 1m ∑m i=2 σi ( X ) 2k ( for details , see ( A.11 ) - ( A.12 ) in Appendix A ) . Furthermore , as n , m→∞ , m2k → m∞2k = E { Λ2k } and κ2k → κ∞2k , where { m∞2k } k≥1 and { κ∞2k } k≥1 are respectively even moments and rectangular free cumulants of Λ. PCA – Rectangular matrices . Denote by uPCA and vPCA the left and right principal singular vectors of X , and define α̃s = 1/ √ D ( b+ ) , where D ( z ) = φ ( z ) φ̄ ( z ) , φ ( z ) = E { z/ ( z2 − Λ2 ) } , φ̄ ( z ) = γφ ( z ) + ( 1−γ ) /z , andD ( b+ ) = limz→b+ D ( z ) . Note that the singular value of the rank-one signal αmu ∗v∗T is α̃ , α/ √ γ . Then , for α̃ > α̃s , σ1 ( X ) a.s.−→ D−1 ( 1/α̃2 ) and σ2 ( X ) a.s.−→ b ; see Theorem 2.8 in [ 11 ] . Furthermore , Theorem 2.9 in [ 11 ] gives that , for α̃ > α̃s , 〈uPCA , u∗〉2 m a.s.−→ ∆PCA = −2φ ( D−1 ( 1/α̃2 ) ) α̃2D′ ( D−1 ( 1/α̃2 ) ) > 0 , ( 2.3 ) 〈vPCA , v∗〉2 n a.s.−→ ΓPCA = −2φ̄ ( D−1 ( 1/α̃2 ) ) α̃2D′ ( D−1 ( 1/α̃2 ) ) > 0 . ( 2.4 ) In words , above the spectral threshold α̃s , the principal singular value escapes from the bulk of the spectrum and the left/right principal singular vectors become correlated with the signal u∗/v∗ .

Summary: The submitted paper addresses estimation of rank-1 matrices from observations (1.1) or (1.2) subject to rotationally invariant noise matrix W. The paper proposes approximate message passing (AMP) with long memory Onsager correction (3.2) in the square case or (3.10) in the rectangular case, inspired by an arXiv preprint [20]. AMP uses the conventional principal component analysis (PCA) estimate as the initial message (3.1) or (3.9). To eliminate dependencies between the initial message and the noise matrix W, non-trivial Onsager correction b_{t,1} in (3.3) or a_{t,1}, b_{t+1,1} given in (3.11) and (3.13) is proposed via state evolution while the other correction b_{t,t-j} in (3.3) or (3.12) and (3.14) is the same as in [20]. Theorems 1 and 2 claim state evolution recursion for the square and rectangular cases, respectively. The theoretical predictions are also verified via numerical simulations in artificial setting.