PyTDC: A multimodal machine learning training, evaluation, and inference platform for biomedical foundation models

TL;DR

PyTDC addresses the lack of end-to-end infrastructure for multimodal biomedical foundation models in therapeutics by delivering an API-first platform that unifies heterogeneous single-cell data, model weights, and standardized benchmarking. The platform combines a multimodal data retrieval layer, a model server for access to SoTA models, and context-specific benchmarking to enable end-to-end training, evaluation, and inference across therapeutic tasks. Its case study on single-cell drug-target nomination demonstrates that context-aware methods outperform context-agnostic baselines but still struggle with unseen cell types and integrating additional modalities, motivating ongoing development of multimodal, context-aware foundation models. By providing open-source tooling, datasets, benchmarks, and model-hub integration, PyTDC aims to accelerate therapeutically relevant advances in biomedical AI and establish reusable standards for future work.

Abstract

Existing biomedical benchmarks do not provide end-to-end infrastructure for training, evaluation, and inference of models that integrate multimodal biological data and a broad range of machine learning tasks in therapeutics. We present PyTDC, an open-source machine-learning platform providing streamlined training, evaluation, and inference software for multimodal biological AI models. PyTDC unifies distributed, heterogeneous, continuously updated data sources and model weights and standardizes benchmarking and inference endpoints. This paper discusses the components of PyTDC's architecture and, to our knowledge, the first-of-its-kind case study on the introduced single-cell drug-target nomination ML task. We find state-of-the-art methods in graph representation learning and domain-specific methods from graph theory perform poorly on this task. Though we find a context-aware geometric deep learning method that outperforms the evaluated SoTA and domain-specific baseline methods, the model is unable to generalize to unseen cell types or incorporate additional modalities, highlighting PyTDC's capacity to facilitate an exciting avenue of research developing multimodal, context-aware, foundation models for open problems in biomedical AI.

Figures (9)

• Figure 1: We present PyTDC, a machine-learning platform providing streamlined training, evaluation, and inference software for single-cell biological foundation models to accelerate research in transfer learning method development in therapeutics Theodoris2024perspectivebenchmarks. PyTDC introduces an API-first beaulieu2022api architecture (sections \ref{['sec:multimodalsinglecellretreival-apifirst']} and \ref{['sec:architecture-appendix']}) that unifies heterogeneous, continuously updated data sources. The platform introduces a model server, which provides unified access to model weights across distributed repositories and standardized inference endpoints (sections \ref{['sec:modelserver']}, \ref{['sec:modelhubserver-appendix']}, and \ref{['sec:model-server-section-appendix']}). The model server accelerates research workflows Kahn2022flashlighticmlplatform by exposing state-of-the-art, research-ready models and training setups for biomedical representation learning models across modalities. Building upon Therapeutic Data Commons huang2021therapeuticshuang2022artificialvelez-arce2024signals, we present single-cell therapeutics tasks, datasets, and benchmarks for model development and evaluation (sections \ref{['sec:benchmarks']} and \ref{['sec:tasks-appendix']}).
• Figure 2: AI inferencing and model evaluation components. The PyTDC model server (sections \ref{['sec:modelserver']} and \ref{['sec:model-server-section-appendix']}) streamlines retrieval, inferencing, and training setup for an array of context-aware biological foundation models and models spanning multiple modalities. A model store retrieval API provides unified access to model weights stored in the Hugging Face Model Hub (https://huggingface.co/apliko), Chan-Zuckerberg CELLxGENE Census fine-tuned models, and TDC huang2021therapeuticshuang2022artificialvelez-arce2024signals storage. The model server also provides access to model classes, tokenizer functions, and inference endpoints supporting PyTorch paszke2019pytorch and Hugging Face Transformers Wolf2020huggingfacetransformers. Extracted embeddings, from either model server inference or pre-computed embedding storage, are ready for downstream use by task-specific benchmarking modules.
• Figure 3: The below code illustrates the integration of PyTDC's single-cell datasets with the model server (sections \ref{['sec:modelserver']}, \ref{['sec:model-server-section-appendix']}, and \ref{['sec:modelhubserver-appendix']}). Here we retrieve the scperturbAissa2021scperturb chemical perturbation response dataset (sections \ref{['sec:chemperturb-appendix']}, \ref{['sec:perturboutcome-appendix']}, and \ref{['sec:dataset-scperturb-appendix']}) and extract gene expression vector embeddings from geneformer. Such a workflow would take a user hundreds-to-thousands of lines of code to develop. PyTDC allows the user to extract single-cell foundation model embeddings from complex and customized gene expression datasets with less than 30 lines of code.
• Figure 4: The below code illustrates how to retrieve the train, test, and validation splits and run model evaluation for the tdc_ml.scDTN benchmark.
• Figure 5: The below code illustrates how to retrieve the PPI network and cell type labels pinnacle used for training and benchmarking of models in tdc_ml.scDTN. Retrieval of cell type labels and embeddings from pinnacle are also supported.