Building Foundation Models to Characterize Cellular Interactions via Geometric Self-Supervised Learning on Spatial Genomics

Cellular interactions form the fundamental/core circuits that drive development, physiology, and disease within tissues. Advances in spatial genomics (SG) and artificial intelligence (AI) offer unprecedented opportunities to computationally analyze and predict the behavior of cell intricate networks, and to identify interactions that drive disease states. However, challenges arise in both methodology and scalability: (i) how to computationally characterize complicated cellular interactions of multi-scale nature where chemical genes/circuits in individual cells process information and drive interactions among large numbers of diverse cell types, and (ii) how to scale up the pipeline to accommodate the increasing volumes of SG data that map transcriptome-scale gene expression and spatial proximity across millions of cells. In this paper, we introduce the Cellular Interaction Foundation Model (CI-FM), an AI foundation model functioning to analyze and simulate cellular interactions within living tissues. In the CI-FM pipeline, we explicitly capture and embed cellular interactions within microenvironments by leveraging the powerful and scalable geometric graph neural network model, and optimize the characterization of cellular interactions with a novel self-supervised learning objective - we train it to infer gene expressions of cells based upon their interacting microenvironment. As a result, we construct CI-FM with 100 million parameters by consuming SG data of 23 million cells. Our benchmarking experiments show CI-FM effectively infers gene expressions conditional on the microenvironmental contexts: we achieve a high correlation and a low mismatch error (MSE of 1.1% relative to the square median expression), with 79.4% of cells on average being annotated as the similar cell type based on their predicted and actual expressions. We demonstrate the downstream utility of CI-FM by: (i) applying CI-FM to embed tumor samples to capture cellular interactions within tumor microenvironments (ROC-AUC score of 0.76 on classifying sample conditions via linear probing on embeddings), and identifying shared signatures across samples; and (ii) using CI-FM to simulate changes in microenvironmental composition in response to T cell infiltration, which highlights how CI-FM can be leveraged to model cellular responses to tissue perturbations - an essential step toward constructing "AI virtual tissues". Our model is open source and publicly accessible at https://huggingface.co/ynyou/CIFM.