Semi-supervised Clustering Through Representation Learning of Large-scale EHR Data
Linshanshan Wang, Mengyan Li, Zongqi Xia, Molei Liu, Tianxi Cai
TL;DR
This work tackles learning informative, transferable patient representations from high-dimensional, sparse EHR data with limited labels. It introduces SCORE, a semi-supervised framework that combines a Poisson-adapted latent factor mixture (PALM) with Gaussian variational approximation (GVA) and leverages pre-trained EHR embeddings to learn compact, transferable patient embeddings. The authors establish convergence and error results for a hybrid EM–GVA algorithm in a setting with diverging latent dimensionality, showing that unlabeled data improve accuracy and robustness to label scarcity. Through simulations and an MS disability-phenotyping application, SCORE outperforms existing methods in parameter estimation, embedding quality, and predictive power, and the learned embeddings generalize across related disability measures. Overall, SCORE offers a scalable, theory-backed approach for decision support and stratification in real-world healthcare, enabling better disease subtyping and outcome prediction with partially labeled, multi-domain data.
Abstract
Electronic Health Records (EHR) offer rich real-world data for personalized medicine, providing insights into disease progression, treatment responses, and patient outcomes. However, their sparsity, heterogeneity, and high dimensionality make them difficult to model, while the lack of standardized ground truth further complicates predictive modeling. To address these challenges, we propose SCORE, a semi-supervised representation learning framework that captures multi-domain disease profiles through patient embeddings. SCORE employs a Poisson-Adapted Latent factor Mixture (PALM) Model with pre-trained code embeddings to characterize codified features and extract meaningful patient phenotypes and embeddings. To handle the computational challenges of large-scale data, it introduces a hybrid Expectation-Maximization (EM) and Gaussian Variational Approximation (GVA) algorithm, leveraging limited labeled data to refine estimates on a vast pool of unlabeled samples. We theoretically establish the convergence of this hybrid approach, quantify GVA errors, and derive SCORE's error rate under diverging embedding dimensions. Our analysis shows that incorporating unlabeled data enhances accuracy and reduces sensitivity to label scarcity. Extensive simulations confirm SCORE's superior finite-sample performance over existing methods. Finally, we apply SCORE to predict disability status for patients with multiple sclerosis (MS) using partially labeled EHR data, demonstrating that it produces more informative and predictive patient embeddings for multiple MS-related conditions compared to existing approaches.