PCA Initialization for Approximate Message Passing in Rotationally Invariant Models
Year: 2021

Abstract:
We study the problem of estimating a rank-1 signal in the presence of rotationally invariant noise—a class of perturbations more general than Gaussian noise. Principal Component Analysis (PCA) provides a natural estimator, and sharp results on its performance have been obtained in the high-dimensional regime. Recently, an Approximate Message Passing (AMP) algorithm has been proposed as an alternative estimator with the potential to improve the accuracy of PCA. However, the existing analysis of AMP requires an initialization that is both correlated with the signal and independent of the noise, which is often unrealistic in practice. In this work, we combine the two methods, and propose to initialize AMP with PCA. Our main result is a rigorous asymptotic characterization of the performance of this estimator. Both the AMP algorithm and its analysis differ from those previously derived in the Gaussian setting: at every iteration, our AMP algorithm requires a specific term to account for PCA initialization, while in the Gaussian case, PCA initialization affects only the first iteration of AMP. The proof is based on a twophase artificial AMP that first approximates the PCA estimator and then mimics the true AMP. Our numerical simulations show an excellent agreement between AMP results and theoretical predictions, and suggest an interesting open direction on achieving Bayes-optimal performance.

heading: 1 Introduction
text: 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 cannot 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).

heading: 2 Preliminaries
text: 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∗.

heading: 3 PCA Initialization for Approximate Message Passing

heading: 3.1 Symmetric Square Matrices
text: We consider a family of Approximate Message Passing (AMP) algorithms to estimate u∗ from X = αnu ∗u∗T +W . We initialize using the PCA estimate uPCA:
u1 = √ nuPCA, f 1 = Xu1 − b1,1u1, (3.1)
with b1,1 = ∑∞ i=0 κi+1α −i. Then, for t ≥ 2, the algorithm computes
ut = ut(f t−1), f t = Xut − t∑ i=1 bt,iu i, (3.2)
where the memory coefficients {bt,i}i∈[1,t] are given by bt,t = κ1, and
bt,1 = t∏ `=2 〈u′`(f `−1)〉 ∞∑ i=0 κi+tα −i, bt,t−j = κj+1 t∏ i=t−j+1 〈u′i(f i−1)〉, for (t− j) ∈ [2, t− 1].
(3.3)
Here, the function ut : R→ R is continuously differentiable and Lipschitz, it is applied componentwise to vectors, i.e., ut(f t−1) = (ut(f t−11 ), . . . , ut(f t−1 n )), and u ′ t denotes its derivative. The AMP algorithm in (3.1)-(3.3) is similar to the one in [20, Sec. 3.1] (and the ones in [14, 48]), with the main differences being the initialization u1 and the formula for the memory term bt,1. We highlight that the algorithm does not require the knowledge of α or of the noise distribution. In fact, α can be consistently estimated from the principal eigenvalue of X via α̂ = (G(λ1(X)))−1. Furthermore, one can compute the moments {mk}k≥1 of the empirical eigenvalue distribution ofX (excluding its largest one) and, from these, deduce the free cumulants {κk}k≥1.
The asymptotic empirical distribution of the iterates ut,f t, for t ≥ 1, can be succinctly characterized via a deterministic recursion, called state evolution, and expressed via a sequence of mean vectors µK = (µt)t∈[1,K] and covariance matrices ΣK = (σs,t)s,t∈[1,K]. For K = 1, set µ1 = αρα and σ11 = α 2(1− ρ2α), with ρα given in (2.2). Then define µK+1,ΣK+1 from µK ,ΣK as follows. Let
(F1, . . . , FK) = µKU∗ + (Z1, . . . , ZK), where (Z1, . . . , ZK) ∼ N (0,ΣK), (3.4)
Ut = ut(Ft−1) for 2 ≤ t ≤ K + 1, and Ut = F1 α for −∞ < t ≤ 1. (3.5)
Then, the entries of µK+1 are given by µt = αE{UtU∗} for t ∈ [1,K + 1]. Furthermore, the entries of ΣK+1 can be expressed via the following formula, for s, t ∈ [1,K + 1]:
σs,t = ∞∑ j=0 ∞∑ k=0 κ∞j+k+2 ( 1 α )(k−t+1)++(j−s+1)+ · E{Us−jUt−k} · ( s∏ i=max(2,s+1−j) E{u′i(Fi−1)} ) · ( t∏ i=max(2,t+1−k) E{u′i(Fi−1)} ) .
(3.6)
Our main result, Theorem 1, shows that for t ≥ 1, the empirical joint distribution of the entries of (u∗,f1, . . . ,f t) converges in Wasserstein distance W2 to the law of the random vector (U∗, F1, . . . , Ft). We provide a proof sketch in Section 5, and the complete proof is deferred to Appendix B. This result is stated in terms of pseudo-Lipschitz test functions. A function ψ : Rm → R is pseudo-Lipschitz of order 2, i.e., ψ ∈ PL(2), if there is a constant C > 0 such that
‖ψ(x)− ψ(y)‖ ≤ C(1 + ‖x‖+ ‖y‖) ‖x− y‖ . (3.7)
The equivalence between convergence in terms of PL(2) functions and convergence in W2 distance follows from [58, Definition 6.7 and Theorem 6.8]. Theorem 1. In the square symmetric model (1.1), assume that α > αs, and that the free cumulants of order 2 and higher are non-negative, i.e., κ∞k ≥ 0 for k ≥ 2. Consider the AMP algorithm with PCA initialization in (3.1)-(3.2), with continuously differentiable and Lipschitz functions ut : R→ R. (Without loss of generality, assume that 〈u∗,uPCA〉 ≥ 0.) Then, for t ≥ 1 and any PL(2) function ψ : R2t+2 → R, we almost surely have:
lim n→∞
1
n n∑ i=1 ψ(u∗i , u 1 i , . . . , u t+1 i , f 1 i , . . . f t i ) = E {ψ(U∗, U1, . . . , Ut+1, F1, . . . , Ft)} , (3.8)
where U1, . . . , Ut+1 and F1, . . . , Ft are defined in (3.4).
Assumptions of the theorem. The basic assumption that the noise matrix is rotationally invariant is rather mild as it allows for arbitrary eigenvalue distributions. The assumption α > αs ensures that the PCA initialization is correlated with the signal. This condition is necessary and sufficient for PCA to be effective: under the additional requirement that G′(b) = −∞, we have that, if α < αs, then the normalized correlation between uPCA and u∗ vanishes almost surely; see Theorem 2.3 of [10]. Conversely for α > αs, the asymptotic correlation is strictly non-zero and given by (2.2).
Non-negativity of free cumulants: The assumption that κ∞k ≥ 0 for k ≥ 2 appears to be an artifact of the proof technique. As detailed in the proof sketch in Section 5, this assumption is needed to show that the state evolution of the artificial AMP in the first phase has a unique fixed point. We expect our approach to generalize to any limiting noise distribution Λ with compact support, and defer such a generalization to future work. In support of this view, the simulations of Section 4 verify the
claim of Theorem 1 in a setting where the free cumulants of Λ have alternating signs (corresponding to an eigenvalue distribution Λ ∼ Uniform[−1/2, 1/2]; see Figs. 1b-1d and 2b–2d). Finally, we remark that, ifW follows a Marcenko-Pastur distribution (W = AAT, whereA has i.i.d. Gaussian entries), then the free cumulants of Λ are all equal and strictly positive; see [40, Chap. 2, Exercise 11]. Thus, the assumption of Theorem 1 holds for noise distributions that are sufficiently close to the Marcenko-Pastur one, or for sufficiently large values of the signal-to-noise ratio α.1
Continuous differentiability and other technical assumptions: The assumption that ut is continuously differentiable can be weakened to: (i) ut being differentiable almost everywhere, and (ii) satisfying a mild non-degeneracy condition (Assumption 4.2(e) in [20]). In this way, we can cover most practically relevant choices of ut such as soft thresholding and ReLU. Theorem 1 also requires the technical assumptions in (2.1) and the text below it: convergence of the empirical distributions of the signal and of the eigenvalues of the noise matrix; boundedness of the (2 + ε)-moment of the signal; and compact support of the spectrum of the noise matrix. We regard these technical assumptions as minor, and remark that they are quite standard in the literature. For the rectangular case, we also need the additional assumption that the law of the signal is zero mean and satisfies a log-Sobolev inequality, which is necessary to apply the framework in [10].
How PCA initialization influences AMP. The form of the memory coefficient bt,1 in (3.3) reflects the PCA initialization of the AMP iteration. PCA initialization can be interpreted as the result of a first AMP phase with linear denoisers (see the proof sketch in Sec. 5). The coefficient bt,1 multiplying the initialization u1 represents the cumulative effect of this first AMP phase leading to the PCA estimate. The main differences from the AMP algorithm in [20] (where the initialization is independent of W ) are the expressions for the coefficient bt,1 and the state evolution parameters σs,t (compare (3.3) and (3.6) in this paper with (1.15) and (1.17) in [20]). One can interpret the new form of bt,1 and σs,t as a memory of the PCA initialization. For the special case of Gaussian noise, the spectral initialization only affects the first iteration of AMP [44]. This is due to the fact that, while in a rotationally invariant model the AMP iterate at step t depends on all previous iterates, in the Gaussian case it depends only on the iterate at step t− 1.
Choice of ut(·). Theorem 1 holds for any choice of denoisers {ut} that are Lipschitz and continuously differentiable. Indeed, our analysis shows that by picking ut(f) = f/α, AMP just returns the PCA estimate; see the proof sketch in Section 5. If some structural information about the signal is available (e.g., sparsity), denoisers that take advantage of this structure can give substantial improvements over PCA. Thus, a key question is how to optimally select the ut’s. Theorem 1 tells us that the empirical distribution of f t converges to the law of µtU∗ + √ σt,tZ, for Z ∼ N (0, 1) and independent of U∗. Hence, the quality of the estimate at each iteration t is governed by the SNR ρt := µ 2 t/σt,t. Consider running the algorithm for t̄ iterations, and let u t̄+1 = ut̄+1(f t̄) be the final estimate. Then, for each t ∈ [2, t̄ ], the Bayes-optimal choice for ut is the one that maximizes ρt, i.e., the SNR for the next iteration. In the case of Gaussian noise [44], the maximum is achieved by the posterior mean ut(f) = E{U∗ |µtU∗ + √ σt,tZ = f}. For rotationally invariant noise, this choice minimizes the mean-squared error 1n‖u t − u∗‖2 (for fixed u1, . . . , ut−1), but it does not necessarily maximize the SNR ρt. We provide an example of this behavior in the simulations reported in Section 4. Therefore, the optimal strategy would be to choose functions u2, . . . , ut̄ to maximize the SNRs ρ2, . . . , ρt̄, and then in the final iteration, to pick ut̄+1 to minimize the desired loss. Note that ut depends on the previously chosen functions u1, . . . , ut−1 in a complicated way, due to the definition of σt,t in (3.6). Thus, finding ut that maximizes the SNR ρt remains an outstanding challenge. Finally, we remark that though we only consider one-step denoisers in this paper, Theorem 1 can be readily extended to cover denoisers with memory, i.e., those of the form ut(f1, . . . ,f t−1).

heading: 3.2 Rectangular Matrices
text: We now present an AMP algorithm to estimate u∗ and v∗ from the m × n data matrix X = α mu ∗v∗T +W . We initialize the algorithm using the PCA estimate uPCA:
1One can add an independent artificial noise matrix with Marcenko-Pastur distribution to the data in order to make the required free cumulants non-negative, and the result would hold for α greater than the new spectral threshold.
u1 = √ muPCA, g 1 = ( 1 + γ
∞∑ i=1 κ2i ( γ α2 )i)−1 XTu1, v1 = v1(g 1) = γ α g1. (3.9)
Then, for t ≥ 1, we iteratively compute: f t=Xvt− t∑ i=1 at,iu i, ut+1 =ut+1(f t), gt+1 =Xut+1− t∑ i=1 bt+1,iv i, vt+1 =vt+1(g t+1).
(3.10)
Here, ut+1, vt+1 : R → R are continuously differentiable Lipschitz functions that act componentwise on vectors. We define a1,1 = α ∑∞ i=1 κ2i ( γ α2 )i , and for t ≥ 2:
at,1 = 〈v′t(gt)〉 t∏ i=2 〈u′i(f i−1)〉〈v′i−1(gi−1)〉 ( ∞∑ i=0 κ2(i+t) ( γ α2 )i) , (3.11)
at,t−j = 〈v′t(gt)〉 t∏
i=t−j+1 〈u′i(f i−1)〉〈v′i−1(gi−1)〉κ2(j+1), for (t− j) ∈ [2, t]. (3.12)
Furthermore, for t ≥ 1,
bt+1,1 = γ〈u′t+1(f t)〉 t∏ i=2 〈v′i(gi)〉〈u′i(f i−1)〉
( κ2t +
∞∑ i=1 κ2(i+t) ( γ α2 )i) , (3.13)
bt+1,t+1−j = γ〈u′t+1(f t)〉 t∏ i=t+2−j 〈v′i(gi)〉〈u′i(f i−1)〉κ2j , for (t+ 1− j) ∈ [2, t]. (3.14)
Similarly to the square case, α can be consistently estimated from the largest singular value of X via α = √ γ(D(σ1(X)))−1, and the rectangular free cumulants {κ2k}k≥1 can be obtained from the even moments of the empirical distribution of the singular values ofX (excluding its largest one).
The asymptotic empirical distributions of the iterates (f t , gt) can be characterized via a state evolution recursion, which specifies a sequence of mean vectors µK = (µt)t∈[0,K], νK = (νt)t∈[1,K] and covariance matrices ΣK = (σs,t)s,t∈[0,K],ΩK = (ωs,t)s,t∈[1,K]. These are iteratively defined, starting with the initialization µ0 = α √ ∆PCA and σ0,0 = α2(1−∆PCA), where ∆PCA is given by (2.3). Having defined µK ,ΣK ,νK ,ΩK , let
(F0, . . . , FK) = µKU∗ + (Y0, . . . , YK), where (Y0, . . . , YK) ∼ N (0,ΣK), (3.15)
Ut = ut(Ft−1) for 2 ≤ t ≤ K + 1, and Ut = F0 α for −∞ < t ≤ 1, (3.16)
(G1, . . . , GK) = νKV∗ + (Z1, . . . , ZK), where (Z1, . . . , ZK) ∼ N (0,ΩK), (3.17) Vt = vt(Gt) for 2 ≤ t ≤ K + 1, and Vt = γ
α G1 for −∞ < t ≤ 1. (3.18)
Given µK and ΣK , the entries of νK+1 are given by νt = αE{UtU∗} (for t ∈ [1,K + 1]), and the entries of ΩK+1 (for s+ 1, t+ 1 ∈ [1,K + 1]) are given by
ωs+1,t+1 = ∞∑ j=0 ∞∑ k=0 γ ( γ α2 )(j−s)++(k−t)+ ( s+1∏ i=max(2,s+2−j) xi · yi−1 ) · ( t+1∏ i=max(2,t+2−k) xi · yi−1 )
· [ κ∞2(j+k+1)E{Us+1−jUt+1−k} + κ ∞ 2(j+k+2)E{Vs−jVt−k}xs+1−j · xt+1−k ] . (3.19)
Here, we define xi = E{u′i(Fi−1)} if i ≥ 2, and xi = 1/α otherwise; yi = E{v′i(Gi)} if i ≥ 2, and yi = γ/α otherwise. We note that ω11 is computed by solving the linear equation obtained by setting s = t = 0 in (3.19) (see (C.96)). Next, given νK+1 and ΩK+1 for some K ≥ 1, the entries of µK+1 are µt = αγE{VtV∗} (for t ∈ [0,K + 1]), and the entries of ΣK+1 (for s, t ∈ [0,K + 1]) are
σs,t = ∞∑ j=0 ∞∑ k=0 ( γ α2 )(j−s+1)+ + (k−t+1)+ ( s∏ i=max(2,s+1−j) xi · yi ) · ( t∏ i=max(2,t+1−k) xi · yi )
· [ κ∞2(j+k+1)E{Vs−jVt−k} + κ ∞ 2(j+k+2)E{Us−jUt−k}ys−j · yt−k ] . (3.20)
Our main result for the rectangular case, Theorem 2, shows that for t ≥ 1, the empirical joint distribution of the entries of (u∗,f1, . . . ,f t) converges in Wasserstein distance W2 to the law of the random vector (U∗, F1, . . . , Ft). Similarly, the empirical joint distribution of the entries of (v∗, g1, . . . , gt) converges to the law of (V∗, G1, . . . , Gt). The proof is given in Appendix C. As in the square case, we state this result in terms of pseudo-Lipschitz test functions. Theorem 2. In the rectangular model (1.2), assume that α̃ > α̃s and that κ∞2k ≥ 0 for k ≥ 1. Consider the AMP algorithm with PCA initialization in (3.9)-(3.10), with continuously differentiable and Lipschitz functions ut, vt : R→ R. (Assume without loss of generality that 〈u∗,uPCA〉 ≥ 0.) Then, for t ≥ 1 and any PL(2) functions ψ : R2t+2 → R and ϕ : R2t+1 → R, we almost surely have:
lim m→∞
1
m m∑ i=1 ψ(u∗i , u 1 i , . . . , u t+1 i , f 1 i , . . . f t i ) = E {ψ(U∗, U1, . . . , Ut+1, F1, . . . , Ft)} , (3.21)
lim n→∞
1
n n∑ i=1 ϕ(v∗i , v 1 i , . . . , v t i , g 1 i , . . . g t i) = E {ϕ(V∗, V1, . . . , Vt, G1, . . . , Gt)} , (3.22)
where (U1, . . . , Ut+1), (F1, . . . , Ft), (V1, . . . , Vt) and (G1, . . . , Gt) are defined as in (3.15)-(3.18).
The condition α̃ > α̃s is necessary and sufficient for PCA to be effective: under the additional requirement that φ′(b+) = −∞, if α̃ < α̃s, then the normalized correlation between uPCA and u∗ vanishes almost surely, see [11, Theorem 2.10]. Comments similar to those at the end of Section 3.1 can be made about (i) the requirement that the rectangular free cumulants are non-negative, (ii) the effect of the PCA initialization on AMP, and (iii) the choice of the denoisers ut, vt.

heading: 4 Numerical Simulations
text: We consider the following settings: (i) square model (1.1) with Marcenko-Pastur noise, i.e.,W = 1 nAA
T ∈ Rn×n, where the entries ofA ∈ Rn×p are i.i.d. standard Gaussian, see (a) in the figures; (ii) square model (1.1) with uniform noise, i.e.,W = OTΛO ∈ Rn×n, whereO is a Haar orthogonal matrix and the entries of Λ are i.i.d. and uniformly distributed in the interval [−1/2, 1/2], see (b) in the figures; (iii) rectangular model (1.2) with uniform noise, i.e.,W = OTΛQ ∈ Rm×n, where O,Q are Haar orthogonal matrices and the entries of Λ2 are i.i.d. and uniformly distributed in the interval [0, 1], see (c)-(d) in the figures.
In the simulations, α is estimated from the largest eigenvalue/singular value ofX . Furthermore, the free cumulants κk (κ2k in the rectangular case) are replaced by their limits κ∞k (κ ∞ 2k resp.), which are obtained as follows. For (a), all the free cumulants of Λ are equal to c , p/n, i.e., κ∞k = c for k ≥ 1, see [40, Chap. 2, Exercise 11]. For (b), the odd free cumulants of Λ are 0 and the even ones are given by κ∞2n = B2n/(2n!), where B2n denotes the 2n-th Bernoulli number. For details, see the derivation of (A.10) in Appendix A. For (c)-(d), the even moments of Λ are given bym∞2k = 1/(k+1) and, from these, we numerically compute the rectangular free cumulants via (A.12) in Appendix A. Furthermore, the spectral threshold for the setting in (a) is αs = 1 + √ c ; for (b), αs = 0 ; and for (c)-(d), α̃s = 0. In (a), we set n = 8000 and c = 2; in (b), we set n = 4000; and in (c)-(d), we set n = 8000 and γ = 1/2. The signal u∗ has a Rademacher prior, i.e., its entries are i.i.d. and uniform in {−1, 1}. In the rectangular case, the signal v∗ has a Gaussian prior, i.e., it is uniformly distributed on the sphere of radius √ n. Given these priors, ut is chosen to be the single-iterate posterior mean denoiser given by ut(x) = tanh(µt x/σt,t), where µt and σt,t are the state evolution parameters; these are replaced by consistent estimates in the simulations. For the rectangular case, we choose vt(x) = x. Each experiment is repeated for ntrials = 100 independent runs. We report the average and error bars at 1 standard deviation.
Figure 1 compares the performance between the proposed AMP algorithm with PCA initialization (PCA+AMP) and the theoretical predictions of state evolution (SE), for two different values of α. On the x-axis, we have the number of iterations of AMP, and on the y-axis the normalized squared correlation between the iterate and the signal. As a reference, we also plot the performance of PCA as a horizontal line. We observe an excellent agreement of AMP with state evolution, even in the settings (b)-(c)-(d) where the free cumulants (resp. rectangular free cumulants) are alternating in sign. This supports our conjecture that Theorems 1-2 hold for more general noise distributions.
In Figure 2, we run PCA+AMP until the algorithm converges, and we compare the results with (i) the AMP with PCA initialization developed in [44] which assumes that the noise matrix is Gaussian (with the correct variance), and (ii) the PCA method alone, as a function of the SNR α. For MarchenkoPastur noise (setting (a)), PCA+AMP always improves upon the PCA initialization. However, this is not the case when the eigenvalues/singular values of the noise are uniformly distributed (settings (b), (c) and (d)). In fact, we observe a phase transition phenomenon: below a certain critical α, AMP converges to a trivial fixed point at 0, while PCA shows positive correlation with the signal; above the critical α, PCA+AMP is no worse than PCA. This is due to the sub-optimal choice of ut; recall the discussion on p.6. We observe no improvement for the estimation of the right singular vector (setting (d)), as the prior of v∗ is Gaussian, in which case we expect the PCA estimate to be optimal. The interesting behavior demonstrated in Figure 2 motivates the study of the optimal choice for ut, vt in future work. We also note that, in settings (c)-(d), α̃s = 0, which means that the PCA estimator has non-zero correlation with the signal for all α > 0. However, for α < 0.1, this correlation remains rather small. Finally, we highlight that our proposed rotationally invariant PCA+AMP always improves upon the Gaussian PCA+AMP. In general, this performance gap will be significant unless the sequence of free cumulants κ∞k (κ ∞ 2k in the rectangular case) decays quickly. For Marchenko-Pastur noise, the free cumulants are all equal, and thus the performance gap is significant. If the eigenvalues/singular values of the noise are uniform, then the sequence of free cumulants decays rapidly and the performance gap is small.

heading: 5 Proof Sketch: Symmetric Square Matrices
text: We consider the following artificial AMP algorithm, whose iterates are denoted by ũt, f̃ t
for t ≥ 1. We initialize with ũ1 = ραu∗ + √ 1− ρ2α n and f̃ 1 = Xũ1 − κ1ũ1. Here, n is standard Gaussian and ρα is the normalized (limit) correlation of the PCA estimate given in (2.2). We note that this initialization is impractical, as it requires the knowledge of the unknown signal u∗. However, this is not an issue since the artificial AMP serves only as a proof technique. (The true AMP (3.2) used for estimation uses the PCA initialization in (3.1).) The subsequent iterates of the artificial AMP are defined in two phases. In the first phase, which lasts up to iteration (T + 1), the functions defining
the artificial AMP are chosen so that ũT+1 is closely aligned with the eigenvector uPCA as T →∞. In the second phase, the functions are chosen to match those in the true AMP.
The artificial AMP initialization ũ1 is chosen such that it has non-zero asymptotic correlation with the signal u∗. Indeed, when the signal prior has zero mean, a random initialization (independent of u∗) would be asymptotically uncorrelated with the signal; consequently, the first phase of the artificial AMP would get stuck at a trivial fixed point and the iterates would not be guaranteed to converge to the principal eigenvector. We ensure that this does not happen by defining the initialization ũ1 to be a linear combination of the signal and Gaussian noise.
First phase. For 2 ≤ t ≤ (T + 1), the artificial AMP iterates are
ũt = f̃ t−1 /α, f̃ t = Xũt − t∑ i=1 b̃t,iũ i, (5.1)
where b̃t,t−j = κj+1α−j , for (t − j) ∈ [1, t]. We claim that, for sufficiently large T , ũT+1 approaches the PCA estimate uPCA, that is, limT→∞ limn→∞ 1√n‖ũ T+1 − √ nuPCA‖ = 0. This result is proved in Lemma B.7 in Appendix B.3. We give a heuristic sanity check here. Assume that the iterate ũT+1 converges to a limit ũ∞ in the sense that limT→∞ limn→∞ 1√n‖ũ
T+1 − ũ∞‖ = 0. Then, from (5.1), the limit ũ∞ satisfies
ũ∞ = 1
α Xũ∞ − ∞∑ i=1 κi ( 1 α )i ũ∞ ⇐⇒ ( α+ ∞∑ i=1 κi ( 1 α )i−1) ũ∞ = Xũ∞, (5.2)
which means that ũ∞ is an eigenvector of X . Furthermore, by using known identities in free probability (see (A.4) and (A.6)), the eigenvalue α+ ∑∞ i=1 κi ( 1 α )i−1 can be re-written asG−1(1/α). Recall that, for α > αs,X exhibits a spectral gap and its largest eigenvalue converges to G−1(1/α). Thus, u∞ must be aligned with the principal eigenvector ofX , as desired.
A key step in our analysis is to show that, as T →∞, the state evolution of the artificial AMP in the first phase has the unique fixed point (µ̃ = αρα, σ̃ = α2(1− ρ2α)). This is established in Lemma B.2 proved in Appendix B.2. The proof follows the approach developed in Section 7 of [20]. However, the analysis of [20] requires that α is sufficiently large, while our result holds for all α > αs. Our idea is to exploit the expression of the limit correlation between the PCA estimate and the signal. In particular, we prove that, when the PCA estimate is correlated with the signal, state evolution is close to a limit map which is a contraction. The price to pay for this approach is the requirement that the free cumulants are non-negative.
Second phase. The second phase of the artificial AMP is designed so that its iterates (ũT+k, f̃ T+k ) are close to (uk,fk), for k ≥ 2. For t ≥ (T + 2), the artificial AMP computes:
ũt = ut−T (f̃ t−1 ), f̃ t = Xũt − t∑ i=1 b̃t,iũ i. (5.3)
Here, the functions {uk}k≥2, are the ones used in the true AMP (3.2). The coefficients {b̃t,i} for t ≥ (T + 2) are given by:
b̃tt = κ1, b̃t,t−j = κj+1
( 1
α )(T+1−(t−j))+ t∏ i=max{t−j+1,T+2} 〈u′i−T (f̃ i−1 )〉, (t− j) ∈ [1, t− 1].
(5.4)
Since the artificial AMP is initialized with ũ1 that is correlated with u∗ and independent of the noise matrixW , a state evolution result for it can be obtained directly from [20, Theorem 1.1]. We then show in Lemma B.8 in Appendix B.4 that the second phase iterates in (5.3) are close to the true AMP iterates in (3.2), and that their state evolution parameters are also close. This result yields Theorem 1, as shown in Appendix B.5. The complete proof of Theorem 2 (rectangular case) is given in Appendix C. We describe the artificial AMP for this case along with a proof sketch in Appendix C.1.

heading: Acknowledgements
text: M. Mondelli would like to thank László Erdös for helpful discussions. M. Mondelli was partially supported by the 2019 Lopez-Loreta Prize. R. Venkataramanan was partially supported by the Alan Turing Institute under the EPSRC grant EP/N510129/1.

Review summaries:
This paper considers the problem of estimating a signal from its noisy rank-1 measurement. PCA is a traditional method that can be used for this task. AMP is an alternative that can additionally leverage statistical priors on the unknown signal. AMP algorithms are traditionally analyzed using a scalar iteration known as state evolution (SE).  The focus of this paper is to theoretically establish SE for AMP under a more general noise matrix starting from an initialization using PCA. Having SE enables the analysis of the error behaviour of AMP in different settings. Empirical results focus on comparing AMP+PCA against PCA on simulated data.  Theoretical contributions: The paper develops a new AMP algorithm initialized using PCA. Theorem 1 and 2 are the main theoretical results in the paper that establish SE for the proposed algorithm for square-symmetric and rectangular matrices.  Empirical contribution: Validation of the superiority of AMP on simulated data over PCA. 

The authors consider the problem of estimating a structured rank-1 signal in the presence of rotationally invariant noise, which is a class of perturbations that is more general than Gaussian noise. For this problem, they propose an AMP algorithm that is initialized by PCA, and they derivate a state-evolution characterization that holds in the infinite-dimensional limit. The rationale is that PCA is a near-optimal estimator in the absence of structure, whereas AMP facilitates the exploitation of structure yet allows a rigorous asymptotic analysis. Their work builds upon previous work, such as PCA-initialized AMP with Gaussian noise, but is different in that, to handle rotationally invariant noise, their algorithm uses a PCA correction at every iteration, not just the first one. Their state-evolution analysis uses a two-phase artificial AMP formulation that has previously been used in the context of generalized linear models.  Although their analysis assumes a noise distribution with non-negative free cumulants, they provide empirical simulation results that suggest that derived state evolution is correct with any noise distribution with compact support.

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. 

This paper considers the rank-1 signal recovery for rotationally invariant noise matrices. It uses PCA to generate an initialization of AMP that is correlated with the signal but independent of the noise. Different from the IID Gaussian noise matrix case, rotationally invariant noise matrix makes the construction and the state evolution of AMP be more difficult. For example, the Onsager term in AMP not only involves the memory in last iteration, but also contains all the memories in preceding iterations. Therefore, the PCA initialization will change the expression of AMP (and its SE) in all the iterations. In a word, this work focuses on a hard rank-1 signal recovery problem with general rotationally invariant noise matrices. The authors propose a modified AMP to solve this problem. Furthermore, the state evolution of the proposed AMP is also rigorously proved.