POME: Graph-based embeddings for partially observed mixed-type data

Abstract

Partially observed mixed-type (POM) data, as often encountered in clinical, epidemiological, and phenotypic datasets, are very common in biomedical research. Yet, advanced data analysis and machine learning based on POM data is complicated by their heterogeneous nature and often substantial fractions of missing values. While one promising way to overcome these issues is to compute vector-valued embeddings for POM datasets which can then be used for downstream analyses, existing embedding methods are mostly not designed for POM data. To address this gap, we developed POME (partially observed mixed-type data embeddings), a self-supervised model that yields low-dimensional representations of both samples and variables, using shared concept learning and a bipartite graph representation of the underlying POM data. We validated POME through extensive experiments on three real-world biomedical datasets, with diverse downstream tasks and objectives: POME achieves state-of-the-art imputation performance and produces high-quality patient representations that support not only unsupervised discovery of well-separated and clinically meaningful patients subgroups but also supervised predictive modeling and zero-shot representation space mining for use cases such as adjuvant therapy modality recommendation. POME is available as a Python package on GitHub (https://github.com/bionetslab/POME) and PyPI (https://pypi.org/project/pome-py).