Scale-Invariant, Robust Sparse PCA for Large Data via Differentiable Penalties

Abstract

Sparse PCA finds low-dimensional structure that loads on few features. Existing methods couple learning and feature selection by applying non-differentiable penalties that force retain-or-zero decisions during optimization. This eliminates features before the optimizer has established which ones matter, limits scalability to serial coordinate-update solvers, and fails when components share support. We introduce DROSS-PCA (Differentiable RObust Scalable Sparse PCA), which decouples learning and feature selection. It learns by optimizing a fully differentiable objective combining robust reconstruction, a smooth sparsity penalty, and an orthogonality term. It selects features post-hoc via cosine-preserving pruning. DROSS-PCA is GPU-accelerated, scaling sparse PCA to large data sets beyond the reach of existing methods. All loss terms are calibrated to be independent of data dimensionality and scale, giving penalty weights consistent meaning across data sets. On synthetic benchmarks, DROSS-PCA matches or outperforms established methods on support recovery. Under outlier contamination it substantially exceeds non-robust methods and remains competitive with dedicated robust methods, which do not scale to large problems. The smooth formulation of DROSS-PCA enables stability analysis via random initializations since no feature is eliminated during training. We find that the method’s feature selection is stable when the generating basis is theoretically identifiable and shows increased uncertainty precisely when it is not. On the Human Lung Cell Atlas (584,944 × 27,402), gene selection is 97% consistent across independent random initializations, and consensus gene modules identified from reproducibility across runs are enriched for known biological pathways. On daily sea-surface temperature fields (ERA5, 23,376 × 484,778), the leading sparse mode reproduces the observed El Niño index and the leading modes localize to individual ocean basins.