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.
