Fundamental limits in structured principal component analysis and how to reach them

Our first result is a variational formula for the mutual information via the celebrated replica method (43) outlined in Section 1: If we let τ_* := argmax{f_ρ(τ):τ ∈ ℝ¹³, ∇f_ρ(τ)=0}, then we have the following low-dimensional expression for the mutual information between hidden spike and the data:

The argmax is selected and not the argmin as f_ρ is a free entropy (i.e., minus free energy, the free energy being minimized in physics). f_ρ and its derivation are reported in SI Appendix, section 3.2. The 13 coupled fixed point equations coming from ∇f_ρ = 0 will reduce to only 2 (SI Appendix, Eqs. 79–84) thanks to special symmetries inherent to the Bayes-optimal nature of our analysis. One of the two remaining order parameters, denoted m² and called (squared) “magnetization,” quantifies the asymptotic trace inner product between the minimum mean-square error (MMSE) estimator $\int d{P}_{X\mid Y}(\mathbf{x}\mid \mathbf{Y})\mathbf{x}{\mathbf{x}}^{⊺}$ and the spike ${\mathbf{X}}^{\ast }{\mathbf{X}}^{\ast ⊺}$. It allows us to compute the MMSE as

with m solving the aforementioned system of equations.

The above results hold for a factorized prior P_X^⊗N. Nevertheless, if ${\mathbf{X}}^{\ast }$ is uniformly distributed on the sphere, a variational formula analogous to Eq. 7 can still be derived, as shown in SI Appendix, section 3.6, and the related MMSE computed. Analytical arguments and numerical experiments show that the latter can be achieved using the naive spectral estimator $C\mathit{\nu }{\mathit{\nu }}^{⊺}$ of P^* obtained from the principal eigenvector $\mathit{\nu }=\mathit{\nu }(\mathbf{Y})$ of $\mathbf{Y}$ properly rescaled by a certain factor C(λ, ρ); see ref. 9.

Instead of using an AMP with iterates based on Y, we introduce a preprocessing procedure driven by the AdaTAP formalism (32). The end result is an effective quadratic model (i.e., with only pairwise interactions) which is “equivalent” (in a proper sense described below) to the original one, with coupling matrix

This model being quadratic is now solvable using AdaTAP/AMP and possesses the same thermodynamic properties (free entropy, phase transitions, etc.) as well as the same marginal means and variances as the model in Eq. 4 when N → ∞ (and thus equivalent for our purposes). Therefore, to approximate the MMSE estimator, one can simply “preprocess” $\mathbf{Y}$ by applying $J(\mathbf{Y})$ and then efficiently compute the marginals of the resulting quadratic model by AdaTAP/AMP, see next section. AdaTAP allows us to parametrize the free entropy (i.e., log-partition function) of a model with quadratic Hamiltonian, for a given instance of the interaction matrix, in terms of O(N) order parameters, some of which correspond to the sought marginal means (⟨x_i⟩)_{i ≤ N} and associated variances. The extremization w.r.t. them yields equations that can be solved iteratively and identified with an AMP algorithm. However, the Hamiltonian in Eq. 6 is not quadratic in x but can be made so by fixing certain order parameters as outlined in Section 1. The resulting coupling matrix depends on Y and on the fixed order parameters, whose values are constrained by Bayes optimality (SI Appendix, section 1.4). Using these values, for an initial quartic V(x)=μx²/2 + γx⁴/4, we get the above interaction matrix in Eq. 9 (SI Appendix, sections 5.1 and 5.2).

The “cleaning effect” of $J(\mathbf{Y})$ is illustrated in Fig. 1. In general, for a (K + 1)-order polynomial matrix potential, the preprocessed matrix is a polynomial $J(\mathbf{Y})=\sum _{k\le K}{c}_k{\mathbf{Y}}^k$, with (c_k)_{k ≤ K} depending on V. For example, for V(x)=ξx⁶/6 (with ξ = 27/80 to select unit variance), the preprocessing (derived similarly to the quartic case, see SI Appendix, section 5.3) is J₆(x)=ξλx⁵ − ξλ²x⁴ − ξλ²x²; it has an effect similar to that in Fig. 1. We point out that the statistics of the noise could be only partially known. This issue can be overcome by learning the (c_k) from the data; see SI Appendix, B.

First, we show in SI Appendix, section 4 that existing AMPs (35, 36) do not saturate the MMSE predicted by Eq. 8. We provide a replica-based theory showing that despite these existing AMPs are aware of the noise structure/statistics, they nevertheless make an implicit mismatched assumption of i.i.d. Gaussian noise: The noise structure is “only” exploited to enforce convergence despite the mismatch, rather than as a source of greater statistical accuracy, in contrast to the proposed AMP we explain now.

To cure this issue, we employ the preprocessed $J(\mathbf{Y})$ in AMP, which leads to our (conjectural) Bayes-optimal approximate message-passing (BAMP) algorithm with recursion

with g_t + 1 applied component-wise. For simplicity, we assume to have access to an initialization ${\mathbf{u}}^1\in {\mathbb{R}}^N$ independent of the noise $\mathbf{Z}$ and with a strictly positive correlation with ${\mathbf{X}}^{\ast }$, i.e.,

This requirement is rather standard in the analysis of AMP algorithms (16, 35, 44). However, as having access to such an initialization is often impractical, recent work (18, 36, 45) has designed AMPs initialized with the top eigenvector $\mathit{\nu }(\mathbf{Y})$.

By carefully choosing the Onsager coefficients {c_{t, j}}_{j ∈ [t]}, we rigorously obtain BAMP’s state evolution characterization.

Eq. 12 provides a high-dimensional characterization of our proposed BAMP. A suitable choice of ψ readily gives the MSE of the BAMP iterates. We also note that our result is equivalent to the almost sure convergence in Wasserstein-2 distance of the joint empirical distribution of $({\mathbf{u}}^1,\dots ,{\mathbf{u}}^{t+1},{\mathbf{f}}^1,\dots ,{\mathbf{f}}^t,{\mathbf{X}}^{\ast })$ to (U₁, …, U_t + 1, F₁, …, F_t, X^*), see corollary 7.21 of ref. 44.

We emphasize that our BAMP algorithm is not the usual AMP of ref. 35, where the data matrix $\mathbf{Y}$ are just replaced by the preprocessed matrix $J(\mathbf{Y})$. Indeed, tuning the Onsager coefficients {c_{t, i}} entering BAMP requires a type of “multistage” state evolution recursion which is completely different from the one in ref. 35. The acronym we introduce stresses this crucial distinction. While our replica prediction for the MMSE is nonrigorous, the state evolution analysis of BAMP is rigorous. In Section 2, we show that BAMP improves over the AMP in ref. 35 by comparing their fixed points. This improvement is thus a rigorous conclusion, while the conjecture is that BAMP saturates the Bayes-optimal performance.

Finally, the “multistage” state evolution of BAMP suggests a choice of the denoisers in the AMP of ref. 35, which differs from the greedy strategy of ref. 36 (i.e., picking the full posterior mean denoiser at every iteration). The numerical results of Section 2 also show that this denoiser selection—motivated by BAMP—meets the BAMP performance and, hence, the replica prediction of the Bayes-optimal error.

The Left plot of Fig. 2 considers the quartic ensemble for μ = 0, and the right one refers to the pure power six potential. The signal ${\mathbf{X}}^{\ast }$ has a Rademacher prior $X_i^{\ast }\sim \frac{1}{2}({\delta }_1+{\delta }_{-1})$. The estimators of the spike ${\mathbf{X}}^{\ast }{\mathbf{X}}^{\ast ⊺}$ are compared in terms of the MSE achieved at the fixed point, as a function of the SNR λ. All algorithms are run for N = 8,000, they are initialized with u¹ that satisfies Eq. 11, and the results are averaged over 50 trials; the state evolution recursions and the replica prediction are for N → ∞. In SI Appendix, section 7.2, we provide additional numerical results for a sparse Rademacher prior, which display a similar qualitative behavior.

We observe that all algorithms converge rapidly: 10 iterations are sufficient to reach the corresponding fixed points. A few remarks concerning the results displayed in Fig. 2 are now in order. First, in all settings, the fixed point of the BAMP state evolution (red) matches the replica prediction (black). This is a strong numerical evidence supporting our conjecture that the proposed BAMP algorithm is Bayes-optimal. These theoretical curves for N → ∞ are also remarkably close to the MSE achieved by the BAMP algorithm at N= 8,000.

Second, there is a clear performance gap between our proposed BAMP (red) and the existing AMP algorithms (35, 36) (single-step denoiser in blue, and multistep in ochre). For V(x)=ξx⁶/6, the gap is even more evident. As predicted by our theory, the gap is reduced when μ approaches 1 with all curves collapsing for μ = 1; see SI Appendix, section 7.2.

Thirdly, we consider a choice of denoisers in the AMP of ref. 35, which is motivated by our BAMP: If the potential has degree K, every K-th nonlinearity is the full memory posterior mean denoiser, and all the other denoisers are chosen to be the identity. The algorithm is dubbed AMP with alternating posteriors (AMP-AP), and its connection to BAMP is discussed at the end of the 1 section. As evident from the smaller plots in the top right corner, AMP-AP (blue) matches the performance of BAMP and of the replica prediction as well.

Last, BAMP is numerically unstable for low SNR. For the quartic potential and λ = 2.3, 5 out of 50 trials do not reach the state evolution fixed point (and are thus discarded). Furthermore, BAMP’s state evolution detaches from the replica prediction as the SNR gets smaller. Considering an initialization closer to the fixed point mitigates the issue. This instability is likely due to the fact that BAMP’s state evolution corresponds to an auxiliary AMP that multiplies the number of iterations (see the 1 section) and which thus amplify errors.

We believe that our results apply beyond the rotational invariance assumption to cases where the eigenbasis of the noise is invariant under more restrictive transformations (such as permutations), or even “quasideterministic.” This intuition comes from recent studies (41, 42) showing that, when AMP or its linearized version are used, the class of rotationally invariant matrices leads to the same performance as a much broader class of matrices (with same spectral density). While the existing literature considers a setting different than ours, this still suggests that our predictions should remain true more generally. To confirm this, we plot in Fig. 2 the performance of BAMP when the uniformly distributed matrix O (i.e., the noise $\mathbf{Z}$ eigenbasis) is replaced by i) the product of the Hadamard–Walsh matrix and a diagonal matrix with i.i.d. Rademacher entries, as in ref. 42 (green squares), or ii) the eigenbasis of the covariance matrix for two popular datasets in computer vision and quantitative genetics, i.e., the CIFAR-10 (46) “plane” class and the “muscle skeletal” GTEx dataset (47, 48) (purple and orange markers, respectively, in the Top-Right plots). The excellent match clearly supports the universality of our predictions. Additional validations are contained in SI Appendix, section 7.2. These results can be understood from the fact that any eigenbasis O is typical w.r.t. the Haar measure, so for a fixed instance, as long as O is sufficiently independent of the eigenvalues, the universality should hold. This suggests that, in practice, our rotational invariance assumption effectively corresponds to assuming decoupling between eigenbasis and eigenvectors.

The starting point of the replica method is the “replica trick” lim_{N → ∞}𝔼ln𝒵(Y)/N = lim_{n → 0}lim_{N → ∞}ln𝔼𝒵ⁿ(Y)/(Nn) that implicitly assumes the commutation of the n, N limits. Another key assumption is to consider n ∈ ℕ in the computation and then assume an analytic continuation to n close to 0₊. The expectation is with respect to $\mathbf{Y}$ or equivalently the independent $\mathbf{O},{\mathbf{X}}^{\ast }$; concerning D, we only need that its empirical eigenvalue distribution converges weakly to ρ and that it has asymptotically no outliers. When computing 𝒵ⁿ, we get multiple integrals over (x_ℓ)_{0 ≤ ℓ ≤ n}, with x₀ ≡ X^*, and a sum of n Hamiltonians as in Eq. 6 in the exponential. Expanding the exponent, we identify some order parameters: For 1 ≤ ℓ ≤ n,

After fixing these using the Fourier representation of the Dirac delta function, the replicated partition function reads

where $H_N(\mathit{\tau },\widehat{\mathit{\tau }},\mathbf{x};{\mathbf{x}}_0,\mathbf{Z}):=Nh(\mathit{\tau },\widehat{\mathit{\tau }})+{\mathbf{x}}^{⊺}{\mathbf{J}}_1(\mathit{\tau },\widehat{\mathit{\tau }},\mathbf{Z})\mathbf{x}+{\mathbf{x}}^{⊺}{\mathbf{J}}_0(\mathit{\tau },\widehat{\mathit{\tau }},\mathbf{Z}){\mathbf{x}}_0$, and τ_ℓ := (v_ℓ, M_(1)ℓ, κ_ℓ, m_ℓ) with ${\widehat{\mathit{\tau }}}_{\ell }$ being the Fourier conjugate. The definitions of (h, J₁, J₀) can be found in SI Appendix, section 3.1. This point is crucial as it allows us to write the n Hamiltonians (one per x_ℓ) as at most quadratic functions of x_ℓ. Due to the quartic nature of the potential, the original H_N would instead have quartic interactions, or higher-order ones for polynomial V of degree greater than four. Yet, by identifying the proper order parameters, a similar reduction to effective quadratic Hamiltonians would still be possible.

In 𝔼𝒵ⁿ, the replicas are coupled in the system only through the expectation over the quenched noise, that can be rewritten as an expectation over the Haar distributed noise eigenbasis 𝔼_O. The entire computation then boils down to the evaluation of an inhomogeneous log-spherical integral that we introduced and defined as follows: Let the matrices C_ℓℓ′ = diag((C_{i, ℓℓ′})_{i ≤ N}), C_i = (C_{i, ℓℓ′})_{ℓ, ℓ′≤n}, and vectors h_ℓ = (h_{i, ℓ})_{i ≤ N}, h_i = (h_{i, ℓ})_{ℓ ≤ n} all having bounded entries uniformly in N. The sequence (h_i ∈ ℝⁿ, C_i ∈ ℝ^n × n)_{i ≤ N} is assumed to have an empirical law tending to that of the random variable (h ∈ ℝⁿ, C ∈ ℝ^n × n). The inhomogeneous log-spherical integral is defined as

Its limit depends only on the law of (C, h) and on the overlaps $q_{\ell \ell \prime }:=\frac{1}{N}{\mathbf{x}}_{\ell }^{⊺}{\mathbf{x}}_{\ell \prime }$, ℓ ≤ ℓ′, that we need to fix with additional Dirac deltas in addition to the previous order parameters. We find that lim_{N → ∞}ℐ_N is expressed by a variational formula; see SI Appendix, section 2.1. This integral is a natural generalization of the standard spherical integral (49) and thus may have an interest beyond the present model, in particular, in random matrix theory or spin glasses.

The final ingredient is a replica symmetric ansatz, justified by the strong concentration-of-measure effects taking place in the Bayes-optimal setting (50, 51). It amounts to assume that all order parameters entering the model are independent of the replica index ℓ. Finally, a saddle point yields an extremization over ℝ¹³ of an effective action. Eqs. 7 and 8 follow directly.

Concerning the reduction from 13 to 2 order parameters (saddle point equations): This is possible thanks to a symmetry arising as a consequence of the Bayes rule which is specific to the Bayes-optimal setting and often called Nishimori identity. It allows to “interchange” the ground-truth signal ${\mathbf{X}}^{\ast }$ with a sample $\mathbf{x}$ from the posterior Eq. 4 inside joint expectations over the posterior and data (see, e.g., ref. 51) and as a consequence to automatically fix the value of most order parameters.

The Onsager coefficients {c_{t, i}}_{i ∈ [t],t ≥ 1} are designed so that, conditioned on the signal, the empirical distribution of the iterate f^t is Gaussian, namely $({\mathbf{f}}^1,\dots ,{\mathbf{f}}^t)\stackrel{W_2}{⟶}(F_1,\dots ,F_t):={\mathit{\mu }}_t{X}^{\ast }+{\mathbf{W}}_t$, with ${\mathbf{W}}_t\sim \mathcal{N}(0,{\mathbf{\Sigma }}_t)$ for some mean vector μ_t and covariance matrix Σ_t. For the AMP in ref. 35, this condition is enforced via the reduction to an auxiliary AMP, which also allows to track the iterates of the original algorithm and yields the state evolution parameters, such as μ_t and Σ_t above. This reduction crucially relies on splitting the matrix Y that multiplies the iterate u^t, into the rank-one signal plus the noise matrix. In contrast, in Eq. 10, the iterate is multiplied by the preprocessed matrix $J(\mathbf{Y})$, which cannot be directly split in a similar fashion. Hence, we track all the contributions (Y^ku^t)_{k ≤ K}, so that we can split them as ${\mathbf{Y}}^k{\mathbf{u}}^t=\mathbf{Y}{\mathbf{Y}}^{k-1}{\mathbf{u}}^t=\frac{\lambda }{N}{\mathbf{X}}^{\ast }⟨{\mathbf{X}}^{\ast },{\mathbf{Y}}^{k-1}{\mathbf{u}}^t⟩+\mathbf{Z}{\mathbf{Y}}^{k-1}{\mathbf{u}}^t$.

The key idea is to map the first T iterations of Eq. 10 to the first K × T iterations of an auxiliary AMP with iterates ${({\tilde{\mathbf{z}}}^t,{\tilde{\mathbf{u}}}^t)}_{t\in [KT]}$ and denoisers ${\{{\tilde{h}}_{t+1}\}}_{t\in [KT]}$,

whose state evolution can instead be deduced from ref. 35. The denoisers ${\{{\tilde{h}}_{t+1}\}}_{t\in [KT]}$ of this multistage auxiliary AMP are chosen so that, for t ∈ [T] and ℓ ∈ [K],

More specifically, for t ∈ [T] and ℓ ∈ {2, …, K}, the denoiser ${\tilde{h}}_{K(t-1)+\ell }$ giving ${\tilde{\mathbf{u}}}^{K(t-1)+\ell }$ is a linear combination of past iterates ${\tilde{\mathbf{u}}}^1,\dots ,{\tilde{\mathbf{u}}}^{K(t-1)+\ell -1}$ and of ${\tilde{\mathbf{z}}}^{K(t-1)+\ell -1}$; furthermore, the coefficients of these linear combinations are chosen to ensure that ${\tilde{\mathbf{u}}}^{K(t-1)+\ell }\approx {\mathbf{Y}}^{\ell -1}{\mathbf{u}}^t$. Hence, by using Eq. 13 with Kt in place of t, one gets ${({\mathbf{Y}}^{\ell }{\mathbf{u}}^t)}_{\ell \in [K]}$ from ${\tilde{\mathbf{z}}}^{\mathit{Kt}}$ and ${({\tilde{\mathbf{u}}}^{K(t-1)+\ell })}_{\ell \in \{2,\dots ,K\}}$ (up to an o_N(1)). Thus, $J(\mathbf{Y}){\mathbf{u}}^t$ can be expressed as a linear combination of $({\tilde{\mathbf{u}}}^1,\dots ,{\tilde{\mathbf{u}}}^{\mathit{Kt}},{\tilde{\mathbf{z}}}^{\mathit{Kt}})$, which in turn is a linear combination of i) the past iterates {uⁱ}_{i ∈ [t]}, ii) the signal ${\mathbf{X}}^{\ast }$, plus iii) independent Gaussian noise. By inspecting the coefficients of this linear combination, one deduces a) the Onsager coefficients {c_{t, i}}_{i ∈ [t],t ≥ 1} (as the coefficients multiplying the past iterates {uⁱ}_{i ∈ [t]}), b) the mean μ_t (as the coefficient multiplying the signal ${\mathbf{X}}^{\ast }$), and c) the covariance matrix Σ_t (as the covariance matrix of the remaining noise terms). Finally, by making ${\tilde{h}}_{Kt+1}$ depend on g_t + 1, we enforce that ${\tilde{\mathbf{u}}}^{Kt+1}\approx {\mathbf{u}}^{t+1}$. The description of the auxiliary AMP is deferred to SI Appendix, C.1, and its state evolution follows in SI Appendix, C.2.

In summary, the derivation of BAMP’s Onsager coefficients involves approximating ${\{{\mathbf{Y}}^k{\mathbf{u}}^t\}}_{k\le K-1}$. This suggests an alternative choice of denoisers leading to the algorithm dubbed AMP-AP: For each batch of K iterations, we pick linear denoisers in the first K − 1 of them, as this allows to construct ${\{{\mathbf{Y}}^k{\mathbf{u}}^t\}}_{k\le K-1}$; then, at the K-th iteration, we pick the posterior mean using all the past iterates, as this—in principle—allows one to assemble the vectors ${\{{\mathbf{Y}}^k{\mathbf{u}}^t\}}_{k\le K-1}$ to obtain $J(\mathbf{Y}){\mathbf{u}}^t$ as in BAMP. We note that AMP-AP does not require the coefficients of the polynomial $J(\mathbf{Y})$, but it rather leaves to the posterior mean denoiser to learn them from the data. As such, it provides an efficient alternative to our proposed BAMP.

J.B., F.C., M.M., and M.S. designed research; performed research; analyzed data; and wrote the paper.

The authors declare no competing interest.