StablePCA: Distributionally Robust Learning of Shared Representations from Multi-Source Data

When synthesizing multi-source high-dimensional data, a key objective is to extract low-dimensional representations that effectively approximate the original features across different sources. Such representations facilitate the discovery of transferable structures and help mitigate systematic biases such as batch effects. We introduce Stable Principal Component Analysis (StablePCA), a distributionally robust framework for constructing stable latent representations by maximizing the worst-case explained variance over multiple sources. A primary challenge in extending classical PCA to the multi-source setting lies in the nonconvex rank constraint, which renders the StablePCA formulation a nonconvex optimization problem. To overcome this challenge, we conduct a convex relaxation of StablePCA and develop an efficient Mirror-Prox algorithm to solve the relaxed problem, with global convergence guarantees. Since the relaxed problem generally differs from the original formulation, we further introduce a data-dependent certificate to assess how well the algorithm solves the original nonconvex problem and establish the condition under which the relaxation is tight. Finally, we explore alternative distributionally robust formulations of multi-source PCA based on different loss functions.

💡 Research Summary

StablePCA addresses the challenge of learning a shared low‑dimensional representation from multiple heterogeneous high‑dimensional data sources. Classical Principal Component Analysis (PCA) optimizes the explained variance for a single distribution, but when data come from several batches, hospitals, or imaging protocols, the learned subspace often fails to generalize because of batch effects, differing sample sizes, and source‑specific noise. The authors formulate a distributionally robust version of PCA that maximizes the worst‑case explained variance over all possible mixtures of the source distributions. Formally, let Σ⁽ˡ⁾ be the second‑moment matrix of source l (l = 1,…,L). The uncertainty set C consists of all convex combinations Q = Σₗ ωₗ T⁽ˡ⁾ with ω in the simplex Δ_L. The StablePCA problem is

 P* = arg max_{P∈𝒫_k} min_{ω∈Δ_L} ∑ₗ ωₗ ⟨Σ⁽ˡ⁾, P⟩,

where 𝒫_k denotes the set of rank‑k orthogonal projection matrices. The inner minimization selects the most adversarial mixture of sources for a given projection, forcing the solution to be robust across domains.

The rank‑k constraint makes the problem non‑convex. To obtain a tractable formulation, the authors replace 𝒫_k by its convex hull, the Fantope 𝔽_k = {P ∈ ℝ^{d×d} | 0 ≼ P ≼ I, tr(P)=k}. The relaxed problem becomes a convex–concave saddle‑point problem that can be written as

 max_{P∈𝔽_k} min_{ω∈Δ_L} ∑ₗ ωₗ ⟨Σ⁽ˡ⁾, P⟩.

To solve this efficiently, they develop a Mirror‑Prox algorithm, an extra‑gradient method designed for constrained min‑max problems with non‑Euclidean geometry. The algorithm alternates between a gradient step and a mirror step using appropriate Bregman divergences (entropy for the simplex and a trace‑norm‑induced divergence for the Fantope). Crucially, each update admits a closed‑form projection: onto the simplex via soft‑thresholding, and onto the Fantope via eigenvalue clipping, leading to O(d³) per iteration.

Theoretical contributions are threefold. Theorem 2 proves global convergence of Mirror‑Prox with a rate O(1/√T) in the number of iterations T, and provides a statistical error bound O(√(d log n / n)) when empirical second‑moments are used, where n is the total sample size. Because the relaxed problem may differ from the original non‑convex formulation, the authors introduce a data‑dependent certificate (Theorem 3) that upper‑bounds the suboptimality gap between the current (P, ω) pair and the true StablePCA optimum. Theorem 4 gives a sufficient condition for tightness of the relaxation: if all source covariances are identical or if the optimal mixture weights lie in the interior of the simplex (all ω_l > 0), then any optimal solution of the relaxed problem is also optimal for the original problem.

Beyond the explained‑variance loss, the paper explores alternative robust PCA formulations based on squared reconstruction loss and on regret (the “FairPCA” objective). The regret formulation recovers the FairPCA problem studied in prior work, but instead of solving a semidefinite program (SDP) with O(d⁶·⁵) complexity, the authors apply the same Mirror‑Prox framework, achieving O(d³ T) runtime and dramatically better scalability.

Empirical evaluation covers three real‑world multi‑source settings: (1) single‑cell RNA‑seq data across multiple batches, where StablePCA removes batch‑specific artifacts and yields coherent cell‑type clusters; (2) electronic health records from several hospitals, where the method balances differing feature distributions while preserving disease‑relevant signals; (3) medical images acquired under varying acquisition protocols, where StablePCA aligns anatomical structures across sources. Quantitative metrics—explained variance, silhouette scores, downstream classification accuracy—show consistent improvements over naïve pooled PCA and existing robust methods. In a timing experiment with d = 300, the Mirror‑Prox implementation is roughly 40× faster than the SDP baseline.

In summary, the paper makes three major contributions: (i) a principled distributionally robust PCA formulation that explicitly guards against worst‑case source mixtures; (ii) a convex relaxation combined with a provably convergent Mirror‑Prox algorithm that scales to moderate‑dimensional data; and (iii) a certification framework that quantifies the gap to the original non‑convex problem and identifies conditions under which the relaxation is exact. These advances provide a practical and theoretically sound tool for extracting stable shared representations from heterogeneous multi‑source datasets.