Table of Contents
Fetching ...

StablePCA: Learning Shared Representations across Multiple Sources via Minimax Optimization

Zhenyu Wang, Molei Liu, Jing Lei, Francis Bach, Zijian Guo

TL;DR

StablePCA tackles the challenge of learning shared, low‑dimensional representations from multisource high‑dimensional data under distributional shifts. It casts unsupervised multi‑source PCA as a minimax problem over mixtures of source distributions and relaxes the nonconvex rank constraint with a Fantope relaxation, yielding a convex program that maximizes worst‑case explained variance across sources. An optimistic‑gradient Mirror Prox algorithm with closed‑form updates solves the relaxed problem efficiently and is proven to converge globally at a rate of $O(1/T)$, with a spectral projection step to recover a rank‑$k$ solution that remains near‑optimal for the original nonconvex problem. Empirical results on synthetic data show that StablePCA delivers higher worst‑case explained variance and robustness to sample‑size imbalance and source heterogeneity, while maintaining practical computational efficiency. The framework also opens avenues for extensions to stable CCA and other cross‑domain unsupervised problems, offering a principled approach to robust cross‑domain representation learning.

Abstract

When synthesizing multisource high-dimensional data, a key objective is to extract low-dimensional feature representations that effectively approximate the original features across different sources. Such general feature extraction facilitates the discovery of transferable knowledge, mitigates systematic biases such as batch effects, and promotes fairness. In this paper, we propose Stable Principal Component Analysis (StablePCA), a novel method for group distributionally robust learning of latent representations from high-dimensional multi-source data. A primary challenge in generalizing PCA to the multi-source regime lies in the nonconvexity of the fixed rank constraint, rendering the minimax optimization nonconvex. To address this challenge, we employ the Fantope relaxation, reformulating the problem as a convex minimax optimization, with the objective defined as the maximum loss across sources. To solve the relaxed formulation, we devise an optimistic-gradient Mirror Prox algorithm with explicit closed-form updates. Theoretically, we establish the global convergence of the Mirror Prox algorithm, with the convergence rate provided from the optimization perspective. Furthermore, we offer practical criteria to assess how closely the solution approximates the original nonconvex formulation. Through extensive numerical experiments, we demonstrate StablePCA's high accuracy and efficiency in extracting robust low-dimensional representations across various finite-sample scenarios.

StablePCA: Learning Shared Representations across Multiple Sources via Minimax Optimization

TL;DR

StablePCA tackles the challenge of learning shared, low‑dimensional representations from multisource high‑dimensional data under distributional shifts. It casts unsupervised multi‑source PCA as a minimax problem over mixtures of source distributions and relaxes the nonconvex rank constraint with a Fantope relaxation, yielding a convex program that maximizes worst‑case explained variance across sources. An optimistic‑gradient Mirror Prox algorithm with closed‑form updates solves the relaxed problem efficiently and is proven to converge globally at a rate of , with a spectral projection step to recover a rank‑ solution that remains near‑optimal for the original nonconvex problem. Empirical results on synthetic data show that StablePCA delivers higher worst‑case explained variance and robustness to sample‑size imbalance and source heterogeneity, while maintaining practical computational efficiency. The framework also opens avenues for extensions to stable CCA and other cross‑domain unsupervised problems, offering a principled approach to robust cross‑domain representation learning.

Abstract

When synthesizing multisource high-dimensional data, a key objective is to extract low-dimensional feature representations that effectively approximate the original features across different sources. Such general feature extraction facilitates the discovery of transferable knowledge, mitigates systematic biases such as batch effects, and promotes fairness. In this paper, we propose Stable Principal Component Analysis (StablePCA), a novel method for group distributionally robust learning of latent representations from high-dimensional multi-source data. A primary challenge in generalizing PCA to the multi-source regime lies in the nonconvexity of the fixed rank constraint, rendering the minimax optimization nonconvex. To address this challenge, we employ the Fantope relaxation, reformulating the problem as a convex minimax optimization, with the objective defined as the maximum loss across sources. To solve the relaxed formulation, we devise an optimistic-gradient Mirror Prox algorithm with explicit closed-form updates. Theoretically, we establish the global convergence of the Mirror Prox algorithm, with the convergence rate provided from the optimization perspective. Furthermore, we offer practical criteria to assess how closely the solution approximates the original nonconvex formulation. Through extensive numerical experiments, we demonstrate StablePCA's high accuracy and efficiency in extracting robust low-dimensional representations across various finite-sample scenarios.
Paper Structure (15 sections, 4 theorems, 38 equations, 4 figures, 1 algorithm)

This paper contains 15 sections, 4 theorems, 38 equations, 4 figures, 1 algorithm.

Key Result

Theorem 1

Suppose that for each $\omega\in \mathcal{S}_\omega$, the problem $\min_{M\in \mathcal{F}^k}-\sum_{l=1}^L \omega_l \left\langle \Sigma^{(l)}, M\right\rangle$ has a unique minimizer. Then we have where $\mathcal{P}^k$ and $\mathcal{F}^k$ are defined in eq: original problem and eq: fantope, respectively.

Figures (4)

  • Figure 1: Comparison of the first principal component estimated by PooledPCA (dashed black) and StablePCA (solid red) across settings 1, 2, and 3. PooledPCA directions vary with the sample sizes and the source-specific relationships, while StablePCA consistently identifies the shared variation along $X_1$.
  • Figure 2: Comparative performance of StablePCA (blue squares) and PooledPCA (green circles) with dimensions $d\in \{20,30,...,100\}$. Leftmost: In-distribution worst-case explained variance. Middle: Out-of-distribution worst-case explained variance. Rightmost: Computation time (seconds) versus dimension.
  • Figure 3: Finite-sample error $\|\widehat{M} - M^*\|_{\rm op}$ across samples sizes $n\in \{500, 1000,...,5000\}$ and $d\in \{10,20,30\}$. $\widehat{M}$ is the empirical StablePCA implemented by Algorithm \ref{['algo: mp']} without applying the final spectral truncation. $M^*$ represents the population-level StablePCA, computed with $n=200,000$ samples using Algorithm \ref{['algo: mp']}.
  • Figure 4: The optimality gap $\tau(\widehat{M}^{\rm Proj},\widehat{M})\coloneqq \max_{\omega\in \Delta^L}f(\widehat{M}^{\rm Proj}, \omega)- \max_{\omega\in \Delta^L}f(\widehat{M}, \omega)$, across sample sizes $n\in \{500,1000,1500,...,5000\}$ and $d\in \{10,20,30\}$. $\widehat{M}^{\rm Proj}, \widehat{M}$ are the outputs of Algorithm \ref{['algo: mp']} with and without the last projection step, respectively.

Theorems & Definitions (4)

  • Theorem 1
  • Proposition 1
  • Theorem 2
  • Theorem 3