Robust Heteroskedastic Matrix Factorization: A Generalization of PCA that Flags Outliers and Handles Missing Data

Abstract

We present Robust Heteroskedastic Matrix Factorization (RHMF), a generalization of Principal Component Analysis (PCA) that is robust to outliers, handles per-feature uncertainties and missing data, and automatically flags per-feature and per-object anomalies. RHMF is useful both in recovering a low-dimensional embedding unspoiled by bad data or anomalies, and in identifying those anomalies. It utilises an iterative reweighting algorithm that implicitly maximizes a Student-t likelihood. This admits an equivalent probabilistic interpretation as fitting a hierarchical model with per-data-point latent variances. We deliver a fast JAX implementation, Robusta-HMF, and practical guidance for users. We demonstrate the ability of the model to identify and mitigate outliers of different classes. Identification accuracy is contingent on the choice of hyperparameters, but we show that these can be set reliably by cross-validation. We also apply RHMF to RVS spectra from Gaia DR3 to find main-sequence stars that are strange relative to their neighbors in color-magnitude space. We highlight specific examples, including a known binary hosting a Be star, and M-dwarfs with subtle emission in the Ca II triplet lines, indicative of accretion or magnetic activity, which would not be obvious to identify by eye.