← All posts

Finding structure every data source agrees on: StablePCA

May 2025, revised Mar 2026 · with Molei Liu, Jing Lei, Francis Bach, Zijian Guo · arXiv:2505.00940

A motivating scenario

Picture a single-cell RNA-sequencing study. Gene expression is measured for thousands of cells, but the cells arrive in a dozen experimental batches — different sequencing platforms, reagents, lab protocols. The biology you care about (which cells are B cells, which are monocytes) is shared across batches; the technical noise is not. The same structure appears in electronic health records pooled from many hospitals, or medical images acquired under different protocols. In all of these, the first analysis step is usually the same: run PCA to compress ten thousand features down to a manageable representation.

Everything downstream — the clustering, the visualization, the classifier you finally fit — inherits whatever that representation kept and whatever it discarded. So it is worth asking what, exactly, pooled PCA optimizes.

PCA, and the trap hiding in the word "pooled"

Classical PCA looks for the rank-\(k\) projection \(P = VV^\top\) that captures as much variance as possible. Writing \(\Sigma = \mathbb E[XX^\top]\) for the second-moment matrix, the quantity \(\langle \Sigma, P\rangle\) is exactly the variance explained by \(P\), and PCA maximizes it:

\[ P^\star \;\in\; \arg\max_{P \in \mathcal P^k}\ \langle \Sigma,\, P\rangle, \qquad \mathcal P^k = \{\, P = VV^\top : V^\top V = I_k \,\}. \]

This is optimal for the distribution you trained on — and only that distribution. Here is the trap. Running PCA on the pooled data maximizes \(\langle \bar\Sigma, P\rangle\), where \(\bar\Sigma = \sum_{l} (n_l/n)\,\Sigma^{(l)}\) is a sample-size-weighted average of the per-source second moments. A batch that is larger, or simply has more variance to give, contributes more to that average. Its idiosyncrasies become the leading components, and the small or unusual sources are quietly underserved by the representation everyone downstream relies on.

The idea: maximize the worst case, not the average

StablePCA never commits to a single distribution. Instead it builds an uncertainty set containing every mixture of the \(L\) observed sources \(\mathbb T^{(1)},\dots,\mathbb T^{(L)}\),

\[ \mathcal C \;=\; \Big\{\, \mathbb Q = \textstyle\sum_{l=1}^{L} \omega_l\, \mathbb T^{(l)} \ :\ \omega \in \Delta^L \,\Big\}, \]

where \(\Delta^L\) is the probability simplex, and asks for the projection whose worst-case explained variance over that set is as large as possible:

\[ P^\star \;\in\; \arg\max_{P \in \mathcal P^k}\ \min_{\omega \in \Delta^L}\ \sum_{l=1}^{L} \omega_l\, \langle \Sigma^{(l)},\, P\rangle. \]

Because the inner objective is linear in \(\omega\), the minimizing mixture always sits at a vertex of the simplex: the worst mixture is the worst single source. So StablePCA is really maximizing \(\min_l \langle \Sigma^{(l)}, P\rangle\), the variance captured for whichever source the projection serves least well. No single batch can hijack the components, and when there is no shift at all — \(\mathbb T^{(1)}=\dots=\mathbb T^{(L)}\) — the uncertainty set collapses to a point and StablePCA reduces to ordinary PCA.

Readers of the supervised version of this idea will recognize the shape of the argument: hedge over mixtures of sources rather than trusting their average. The unsupervised setting is in one way easier — with no outcome to model, the worst case is attained at a single source — and in another way harder, because the rank constraint that defines PCA is not convex.

Two dimensions and three sources are enough to see what changes.

Making it computable

The obstacle is the rank constraint \(\mathcal P^k\), which is nonconvex — so the max–min above is a nonconvex minimax problem. We relax the projections to their convex hull, the Fantope,

\[ \mathcal F^k \;=\; \{\, M = M^\top \ :\ 0 \preceq M \preceq I_d,\ \operatorname{Tr}(M) = k \,\}, \]

which leaves the objective bilinear in \((M,\omega)\): convex in \(M\), concave in \(\omega\). We then run Mirror-Prox, a gradient method with an extra-gradient step suited to constrained minimax geometry. Every iteration is closed-form. Writing \(\widehat\Sigma(\omega) = \sum_l \omega_l \widehat\Sigma^{(l)}\), the projection step is a capped matrix exponential and the weight step is a softmax:

\[ M^{t+\frac12} = \sum_{j=1}^d \min\{e^{\lambda_j^t + \nu^t},\,1\}\, U_j^t U_j^{t\top}, \qquad \omega_l^{t+\frac12} \;\propto\; \omega_l^t\, e^{-\eta_\omega \langle \widehat\Sigma^{(l)},\, M^t\rangle}, \]

where \(U^t \mathrm{Diag}(\lambda^t) U^{t\top}\) diagonalizes \(\log M^t + \eta_M \widehat\Sigma(\omega^t)\), and the scalar \(\nu^t\) is chosen so the capped eigenvalues sum to \(k\) — that capping is the projection onto the Fantope. The weight update is exponentiated gradient on the simplex: sources that the current projection serves badly get more weight next round.

Each iteration costs one \(d \times d\) eigendecomposition, so \(O(d^3 T)\) overall. The comparison worth making is against semidefinite programming, the standard solver for the related fair-PCA objective, which costs \(O(d^{6.5})\). Both solvers handle that objective, so the paper races them on it directly:

Seconds per solve, averaged over 100 replications with \(T = 500\) iterations, on the fair-PCA objective. SDP wins at \(d = 30\); the curves cross near \(d = 70\); by \(d = 300\) Mirror-Prox is 39× faster and far more predictable (4.49 ± 0.07 s against 173.6 ± 70.2 s).

What the theory buys you

Three results turn the relaxation from a convenience into something trustworthy.

The optimization converges at the best possible rate. The duality gap of the averaged iterate decays as \(O\big(k\sqrt{k\log(d/k)\log L}\,/\,T\big)\) — that is, \(O(1/T)\), which matches the known lower bound for first-order methods on bilinear minimax problems. The extra-gradient step is exactly what buys \(O(1/T)\) instead of the \(O(1/\sqrt T)\) of plain gradient descent–ascent. Adding sampling error gives a clean split — up to a shared \(\rho_{\max} k\) prefactor, the bound is \(\sqrt{(d+t)/n} + \sqrt{k\log(d/k)\log L}\,/\,T\): statistics in \(n\), optimization in \(T\).

The relaxation is often exactly tight. By Sion's minimax theorem the problem has a dual: the worst-case mixture \(\omega^\star\) minimizes \(\phi(\omega) = \sum_{i=1}^{k}\lambda_i\big(\Sigma(\omega)\big)\), the sum of the top \(k\) eigenvalues. If \(\Sigma(\omega^\star)\) has an eigengap, \(\lambda_k > \lambda_{k+1}\), then the Fantope solution is the projection onto its top-\(k\) eigenspace. The convex relaxation solves the original nonconvex problem, on the nose.

And when you cannot check the eigengap, the data tell you anyway. Rounding the relaxed \(\widehat M_T\) to the nearest rank-\(k\) projection \(\widehat P_T\) costs you

\[ \tau \;=\; \min_{l}\,\langle \widehat\Sigma^{(l)}, \widehat M_T\rangle \;-\; \min_{l}\,\langle \widehat\Sigma^{(l)}, \widehat P_T\rangle, \]

a quantity computable from the data you already have. Added to the optimization gap of the relaxed iterate, it bounds the distance to the true nonconvex optimum — so once \(T\) is large enough that the \(O(1/T)\) term is negligible, \(\tau\) alone certifies the answer: a small \(\tau\) means it is essentially exact. Across every simulated configuration in the paper, \(|\tau| < 0.003\). The rounding is, in practice, free.

Does it work? Twelve batches of single-cell RNA-seq

Back to the motivating scenario, with real data: human bone-marrow scRNA-seq spanning 12 experimental batches, reduced to the top \(d = 1{,}000\) highly variable genes. We fit a rank-\(k = 50\) projection on 8 randomly chosen batches and score it on the 4 held-out batches the method never saw, measuring worst-case explained variance \(\min_{l\,\in\,\text{held-out}} \langle \widehat\Sigma^{(l)}, \widehat P\rangle\). Repeat over 200 random splits.

Aggregated over the 200 splits, StablePCA has the highest worst-case explained variance on both the training and the held-out batches. The gap is widest exactly where it should be — on the hardest split, the one where every method does its worst:

Relative gaps on the most challenging of 200 random 8-train / 4-test batch splits, at \(k = 50\). The ordering is unchanged at \(k = 100\) and \(k = 150\).

The embedding also behaves the way a biologist would want. Projecting the cells onto StablePCA's 50 components and visualizing with t-SNE or UMAP, cells from the 12 batches mix together — the batch effect is suppressed — while B cells, monocytes, T cells and NK cells separate into coherent clusters. The only overlap is between NK and T cells, which is exactly what one would hope for: cytotoxic T cells and NK cells share much of their transcriptional program.

Simulation: \(d = 40\), \(n = 2{,}000\) per source, \(k = 3\), a 3-dimensional shared loading plus source-specific directions. StablePCA's recovery error falls as sources are added; the competitors' does not move.

A surprise: three "robust PCAs" are the same problem

While analyzing the objective we found that StablePCA, SquaredPCA and FairPCA are the identical minimax problem, differing only by a scalar matrix subtracted from each source's second moment before the max–min is taken:

Here \(\lambda_i^{(l)}\) is the \(i\)-th eigenvalue of \(\Sigma^{(l)}\). Subtracting a source-specific constant re-centres what "doing well on source \(l\)" means: fair PCA measures regret against source \(l\)'s own best rank-\(k\) fit, whereas StablePCA measures raw explained variance. One Mirror-Prox implementation solves all three.

Why it matters

Dimension reduction is the first step of countless multi-source pipelines, and it silently inherits the pooled data's biases. StablePCA is a drop-in alternative with the same interface as PCA — pick a rank, get a projection — but its output is fair to every source rather than dominated by the largest. It recovers shared signal in exactly the regimes where pooled PCA locks onto the dominant batch, and it generalizes to sources the model never saw.

The honest caveat: robustness to the worst mixture is not free. If one source is genuinely irrelevant to your target — pure noise, say — StablePCA will still insist on explaining its variance, and a pooled fit tuned to the right source would beat it. It is the right default when you do not know which source your deployment population resembles, which is the usual situation. When you do know something, the paper shows how to shrink the uncertainty set around a prior mixture and recover the lost predictiveness.

© 2026 Zhenyu (Zach) Wang 王振宇