Contextual AI models for single-cell protein biology

Abstract

Understanding protein function and developing molecular therapies require deciphering the cell types in which proteins act as well as the interactions between proteins. However, modeling protein interactions across biological contexts remains challenging for existing algorithms. Here, we introduce Pinnacle, a geometric deep learning approach that generates context-aware protein representations. Leveraging a multi-organ single-cell atlas, Pinnacle learns on contextualized protein interaction networks to produce 394,760 protein representations from 156 cell type contexts across 24 tissues. Pinnacle’s embedding space reflects cellular and tissue organization, enabling zero-shot retrieval of the tissue hierarchy. Pretrained protein representations can be adapted for downstream tasks: enhancing 3D structure-based representations for resolving immuno-oncological protein interactions, and investigating drugs’ effects across cell types. Pinnacle outperforms state-of-the-art models in nominating therapeutic targets for rheumatoid arthritis and inflammatory bowel diseases, and pinpoints cell type contexts with higher predictive capability than context-free models. Pinnacle’s ability to adjust its outputs based on the context in which it operates paves way for large-scale context-specific predictions in biology.

Introduction

Proteins are the functional units of cells, and their interactions enable different biological functions. The development of high-throughput methods has facilitated the characterization of large maps of protein interactions. Leveraging these protein interaction networks, computational methods [1, 2] have been developed to improve the understanding of protein structure [3], accurately predict functional annotations [4, 5], and inform the design of therapeutic targets [6, 7]. Among them, representation learning methods have emerged as a leading strategy to model proteins [8–10]. These approaches can resolve protein interaction networks across tissues [11–13] and cell types by integrating molecular cell atlases [14] and extending our understanding of the relationship between protein and function [15]. Protein representation learning methods can predict multicellular functions across human tissues [12], design target-binding proteins [16] and novel protein interactions [17], and predict interactions between transcription factors and genes [15].

Proteins can have distinct roles in different biological contexts [18, 19]. While nearly every cell contains the same genome, the expression of genes and the function of proteins encoded by these genes depend on cellular and tissue contexts [11, 20, 21]. Gene expression and the function of proteins can also differ significantly between healthy and disease states [21, 22]. Methods incorporating biological contexts can improve the characterization of proteins and provide precise, context-specific insights. However, deep learning methods produce protein representations (or embeddings) that are context-free: each protein has only one representation learned from either a single context or an integrated view across many contexts [15, 23]. These methods generate one representation for each protein, providing an integrated summary. Context-free protein representations are not tailored to specific biological contexts, such as cell types and disease states. These representations cannot identify protein functions that vary across different cell types, which in turn hamper the prediction of pleiotropy and protein roles in a cell type-specific manner.

Sequencing technologies that measure gene expression with single-cell resolution pave the way toward addressing this challenge. Single-cell transcriptomic atlases [20, 24–27] measure activated genes across many cellular contexts. Through attention-based deep learning [28, 29], which specify models that can pay attention to large inputs and learn the most important elements to focus on in each context, single-cell atlases can be leveraged to boost the mapping of gene regulatory networks that drive disease progression and reveal treatment targets [30]. However, incorporating the expression of protein-coding genes into protein interaction networks remains a challenge. Existing algorithms, including protein representation learning, cannot contextualize protein representations.

We introduce Pinnacle (Protein Network-based Algorithm for Contextual Learning), a context-specific model for comprehensive protein understanding. Pinnacle is a geometric deep learning model adept at generating protein representations through the analysis of protein interactions within various cellular contexts. Leveraging single-cell transcriptomics combined with networks of protein-protein interactions, cell-type-to-cell-type interactions, and a tissue hierarchy, Pinnacle generates high-resolution protein representations tailored to each cell type. In contrast to existing methods that provide a single representation for each protein, Pinnacle generates a distinct representation for each cell type in which a protein-coding gene is activated. With 394,760 contextualized protein representations produced by Pinnacle, where each protein representation is imbued with cell type specificity, we demonstrate Pinnacle’s capability to integrate protein interactions with the underlying protein-coding gene transcriptomes of 156 cell type contexts. Pinnacle models support a broad array of tasks; they can enhance three-dimensional (3D) structural protein representations, analyze the effects of drugs across cell type contexts, nominate therapeutic targets in a cell type specific manner, retrieve tissue hierarchy in a zero-shot manner, and perform context-specific transfer learning. Pinnacle models dynamically adjust their outputs based on the context in which they operate and can pave the way for the broad use of foundation models tailored to diverse biological contexts.

Results

Constructing context-specific networks.

Generating protein representations embedded with cell type context calls for protein interaction networks that consider the same context. We assembled a dataset of context-sensitive protein interactomes, beginning with a multi-organ single-cell transcriptomic atlas [20] that encompasses 24 tissue and organ samples sourced from 15 human donors (Figure 1a). We compile activated genes for every expert-annotated cell type in this dataset by evaluating the average gene expression in cells from that cell type relative to a designated reference set of cells (Figure 1a; Methods 2). Here, ‘activated genes’ are defined as those demonstrating a higher average expression in cells annotated as a particular type than the remaining cells documented in the dataset. Based on these activated gene lists, we extracted the corresponding proteins from the comprehensive reference protein interaction network and retained the largest connected component (Figure 1a). As a result, we have 156 context-aware protein interaction networks, each with 2, 530 ± 677 proteins, that are maximally similar to the global reference protein interaction network and still highly cell type specific (Supplementary Figures S1–S2). Our context-aware protein interaction networks from 156 cell type contexts span 62 tissues of varying biological scales.

Figure 1: Overview of Pinnacle.

(a) Cell type-specific protein interaction networks and metagraph of cell type and tissue organization are constructed from a multi-organ single-cell transcriptomic atlas of humans, a human reference protein interaction network, and a tissue ontology. (b) Pinnacle has protein-, cell type-, and tissue-level attention mechanisms that enable the algorithm to generate contextualized representations of proteins, cell types, and tissues in a single unified embedding space. (c) Pinnacle is designed such that the nodes (i.e., proteins, cell types, and tissues) that share an edge are embedded closer (decreased embedding distance) to each other than nodes that do not share an edge (increased embedding distance); proteins activated in the same cell type are embedded more closely (decreased embedding distance) than proteins activated in different cell types (increased embedding distance); and cell types are embedded closer to their activated proteins (decreased embedding distance) than other proteins (increased embedding distance). (d) As a result, Pinnacle generates protein representations injected with cell type and tissue context; a unique representation is produced for each protein activated in each cell type. Pinnacle simultaneously generates representations for cell types and tissues. (e) Existing methods, however, are context-free. They generate a single embedding per protein, representing only one condition or context for each protein, without any notion of cell type or tissue context. (f-h) The Pinnacle algorithm and its outputs enable (f) multi-modal deep learning (e.g., single-cell transcriptomic data with interactomes), (g) context-specific transfer learning (e.g., between proteins, cell types, and tissues), and (h) contextualized predictions (e.g., efficacy and safety of therapeutics).

Further, we constructed a network of cell types and tissues (metagraph) to model cellular interactions and the tissue hierarchy (Methods 2). Given the cell type annotations designated by the multi-organ transcriptomic atlas [20], the network consists of 156 cell type nodes. We incorporated edges between pairs of cell types based on the existence of significant ligand-receptor interactions and validated that the proteins correlating to these interactions are enriched in the context-aware protein interaction networks in comparison to a null distribution (Methods 2; Supplementary Figure S1c–d). Leveraging information on tissues in which the cell types were measured, we began with 24 tissue nodes and established edges between cell type nodes and tissue nodes if the cell type was derived from the corresponding tissue. We then identified all ancestor nodes, including the root, of the 24 tissue nodes within the tissue hierarchy (Methods 2) to feature 62 tissue nodes interconnected by parent-child relationships. Our dataset thus comprises 156 context-aware protein interaction networks and a metagraph reflecting cell type and tissue organization.

Overview of Pinnacle model.

Pinnacle is a geometric deep learning model capable of generating protein representations predicated on protein interactions within a spectrum of cell type contexts. Trained on an integrated set of context-aware protein interaction networks, complemented by a network capturing cellular interactions and tissue hierarchy (Figure 1b–c), Pinnacle generates contextualized protein representations that are tailored to cell types in which protein-coding genes are activated (Figure 1d). Unlike context-free models, Pinnacle produces multiple representations for every protein, each contingent on its specific cell type context. Additionally, Pinnacle produces representations of the cell type contexts and representations of the tissue hierarchy (Figure 1d–e). This approach ensures a multifaceted understanding of protein interaction networks, taking into account the myriad of contexts in which proteins act.

Given multi-scale model inputs, Pinnacle learns the topology of proteins, cell types, and tissues by optimizing a unified latent representation space. Pinnacle integrates different context-specific data into one context-aware model (Figure 1f) and transfers knowledge between protein-, cell type- and tissue-level data to contextualize representations (Figure 1g). To infuse cellular and tissue organization into this embedding space, Pinnacle employs protein-, cell type-, and tissue-level attention along with respective objective functions (Figure 1b–c; Methods 3). Conceptually, pairs of proteins that physically interact (i.e., are connected by edges in input networks) are closely embedded. Similarly, proteins are embedded near their respective cell type contexts while maintaining a substantial distance from unrelated ones. This ensures that interacting proteins within the same cell type context are situated proximally within the embedding space yet are separated from proteins from other cell type contexts. This approach yields an embedding space that accurately represents the intricacies of relationships between proteins, cell types, and tissues.

Pinnacle disseminates graph neural network messages between proteins, cell types, and tissues using a series of attention mechanisms tailored to each specific node and edge type (Methods 3). The protein-level pretraining tasks consider self-supervised link prediction on protein interactions and cell type classification on protein nodes. These tasks enable Pinnacle to sculpt an embedding space that encapsulates the topology of the context-aware protein interaction networks and the cell type identity of the proteins. Pinnacle’s cell type- and tissue-specific pretraining tasks rely exclusively on self-supervised link prediction, facilitating the learning of cellular and tissue organization. The topology of cell types and tissues is imparted to the protein representations through an attention bridge mechanism, effectively enforcing tissue and cellular organization onto the protein representations. Pinnacle’s contextualized protein representations capture the structure of context-aware protein interaction networks. The regional arrangement of these contextualized protein representations in the latent space reflects the cellular and tissue organization represented by the metagraph. This leads to a comprehensive and context-specific representation of proteins within a unified cell type- and tissue-specific framework.

Pinnacle captures cellular and tissue organization.

Pinnacle generates protein representations for each of the 156 cell type contexts spanning 62 tissues of varying hierarchical scales. In total, Pinnacle’s unified multi-scale embedding space comprises 394,760 protein representations, 156 cell type representations, and 62 tissue representations (Figure 1a). We show that Pinnacle learns an embedding space where proteins are positioned based on cell type context. We first quantify the spatial enrichment of Pinnacle’s protein embedding regions using a systematic method, SAFE [31] (Methods 6.3). Pinnacle’s contextualized protein representations self-organize in Pinnacle’s embedding space as evidenced by the enrichment of spatial embedding regions for protein representations that originate from the same cell type context (significance cutoff $\alpha =0.05$; Figure 2; Supplementary Figures S3–S4).

Figure 2: Enrichment of Pinnacle’s protein embedding regions.

(a-f) Two-dimensional UMAP plots of contextualized protein representations generated by Pinnacle from six different contexts: (a) medullary thymic epithelial cell, (b) bronchial vessel endothelial cell, (c) mesenchymal stem cell, (d) lung microvascular endothelial cell, (e) kidney epithelial cell, and (f) fibroblast of breast. Each dot is a protein representation. Gray dots are representations of proteins from other cell types, and nongray colors indicate the cell type context. Each protein embedding region is expected to be enriched neighborhoods that are spatially localized according to cell type context. To quantify this, we compute spatial enrichment of each protein embedding region using SAFE [31], and provide the mean and max neighborhood enrichment scores (NES) and the number of enriched neighborhoods output by the tool (Methods 6 and Supplementary Figure S3–S4). (g-h) Distribution of (g) the maximum SAFE NES and (h) the number of enriched neighborhoods for 156 cell type contexts (each context has $p$-value < 0.05; hypergeometric test, adjusted using the Benjamin-Hochberg false discovery rate correction with significance cutoff $\alpha =0.05$). 10 randomly sampled cell type contexts are annotated, with their maximum SAFE NES or number of enriched neighborhoods in parentheses.

Next, we evaluate embedding regions to confirm that they are separated by cell type and tissue identity by calculating the similarities between protein representations across cell type contexts. Protein representations from the same cell type are more similar than those from different cell types (Figure 3a). In contrast, a model without cellular or tissue context fails to capture any differences between protein representations across cell type contexts (Figure 3b). Further, we expect the representations of proteins that act on multiple cell types to be highly dissimilar, reflecting specialized cell type-specific protein functions (Supplementary Note S1). We calculate the similarities of protein representations (i.e., cosine similarities of a protein’s representations across cell type contexts) based on the number of cell types in which the protein is active (Supplementary Figure S5a–b). Representational similarities of proteins negatively correlate with the number of cell types in which they act (Spearman’s $\rho =-0.9798;p$-value < 0.001), and the correlation is weaker in the ablated model with cellular and tissue metagraph turned off (Spearman’s $\rho =-0.6334;p$-value < 0.001).

Figure 3: Evaluation of Pinnacle’s contextual representations.

(a-b) Gap between embedding similarities using (a) PINNACLE’s protein representations and (b) a non-contextualized model’s protein representations on $n=394$, 760 samples (i.e., cell type specific protein representations). Similarities are calculated between pairs of proteins in the same cell type (dark shade of color) or different cell types (light shade of color), and stratified by the compartment from which the cell types are derived. We use the two-sided two-sample Kolmogorov-Smirnov test for goodness of fit. Annotations indicate median values. The non-contextualized model is an ablated version of Pinnacle without any notion of tissue or cell type organization (i.e., remove cell type and tissue network and all cell type- and tissue-related components of Pinnacle’s architecture and objective function). The bounds of the box show the quartiles of the data, the center indicates the median value of the data, and the whiskers represent the farthest data point within 1.5× IQR. (c) Embedding distance of Pinnacle’s 62 tissue representations as a function of tissue ontology distance. Gray bars indicate a null distribution (refer to Methods 6 for more details). Both the Spearman correlation ($p$-value = 1.85×10⁻¹¹⁹) and Kolmogorov-Smirnov ($p$-value < 0.001) statistical tests are two-sided. Data are represented as mean values with error bars indicating a 95% confidence interval. (d) Prediction task in which protein representations are optimized to maximize the gap between binding and non-binding proteins. (e) Cell type context (provided by Pinnacle) is injected into context-free structure-based protein representations (provided by MaSIF [3], which learns a protein representation from the protein’s 3D structure) via concatenation to generate contextualized protein representations. Lack of cell type context is defined by an average of Pinnacle’s protein representations. (f) Comparison of context-free and contextualized representations in differentiating between binding and non-binding proteins. Scores are computed using cosine similarity on $n=22$ unique protein pairs (2 binding and 20 non-binding); since Pinnacle generates multiple representations per protein based on context, there are $n=\mathrm{7,956}$ pairwise computations (180 binding and 7,776 non-binding) for the contextualized representations. The binding proteins evaluated are PD-1/PD-L1 and B7–1/CTLA-4. Pairwise scores also are calculated for each of these four proteins and proteins that they do not bind with (i.e., RalB, RalBP1, EPO, EPOR, C3, and CFH). The gap between the average scores of binding and non-binding proteins is annotated for context-free and contextualized representations. The significance of the score gaps between binding and non-binding proteins is measured using a one-sided non-parametric permutation test. Data are represented as mean values with error bars indicating a 95% confidence interval.

We additionally examine whether protein embedding regions are organized by the tissue hierarchy. We leverage Pinnacle’s tissue representations to perform zero-shot retrieval of the tissue hierarchy and then compare tissue ontology distance to tissue embedding distance. Tissue ontology distance is defined as the sum of the shortest path lengths from two tissue nodes to the lowest common ancestor node in the tissue hierarchy, and tissue embedding distance is the cosine distance between the corresponding tissue representations. We expect a positive correlation: the farther apart the nodes are according to the tissue hierarchy, the more dissimilar the tissue representations are. As hypothesized, embedding distances in the latent space and the corresponding distances in the tissue ontology of the same tissues are positively correlated (Spearman’s $\rho =0.36;p$-value = 1.85 × 10⁻¹¹⁹; Figure 3c), and the distribution of tissue embedding distances cannot be attributed to random effects (Kolmogorov-Smirnov two-sided test = 0.50; $p$-value < 0.001). When the tissue ontology is randomly shuffled, the correlation with distances in the embedding space diminishes significantly (Spearman’s $\rho =0.005;p$-value = 0.349; Figure 3c). Since Pinnacle uses the metagraph to systematically integrate tissue organization into both cell type and protein representations, it follows that all of Pinnacle’s representations inherently reflect this tissue organization (Methods 3; Supplementary Figure S6).

Pinnacle enhances 3D structural representations of PPIs.

Protein-protein interactions (PPIs) depend on both 3D structure conformations of the proteins [32, 33] and cell type contexts within which the proteins act [34]. However, protein representations produced by existing artificial intelligence (AI) models based on 3D molecular structures lack cell type context information. We hypothesize that incorporating cellular context information can better differentiate binding from non-binding proteins (Figure 3d). Because 3D structures of molecules (containing precise atom or residue level contact information) provide complementary knowledge to protein-protein interaction networks (summarizing binary interactions between proteins), we expect that context-aware protein interaction networks can improve the ability to differentiate between binding and non-binding proteins across different cell types [35]. As no large-scale dataset with matched structural biology and genomic readouts currently exists to perform systematic analyses, we focus on PD-1/PD-L1 and B7–1/CTLA-4 interacting proteins, important immune checkpoint protein interactors involved in cancer immunotherapies [36].

We compare contextualized and context-free protein representations for binding proteins (i.e., PD-1/PD-L1 and B7–1/CTLA-4) and non-binding proteins (i.e., one of the four binding proteins paired with RalB, RalBP1, EPO, EPOR, C3, or CFH). Cell type context is incorporated into 3D structure-based protein representations [3, 17] by concatenating them with Pinnacle’s protein representation (Figure 3e; Methods 4). Context-free protein representations are generated by concatenating 3D structure-based representations [3, 17] with an average of Pinnacle’s protein representations across all cell type contexts (Methods 4). Contextualized representations, resulting from a combination of protein representations based on 3D structure and context-aware PPI networks, give scores (via cosine similarity) for binding and non-binding proteins of 0.9690 ± 0.0049 and 0.9571 ± 0.0127, respectively. Using Pinnacle’s context-specific protein representations, which have no 3D structure information, binding and non-binding proteins are scored 0.0385 ± 0.1531 and 0.0218 ± 0.1081, respectively. In contrast, using context-free representations, binding and non-binding proteins are scored at 0.9789 ± 0.0004 and 0.9742 ± 0.0078, respectively. Further, comparative analysis of the gap in scores between interacting vs. non-interacting proteins yields gaps of 0.011 (PD-1/PD-L1) and 0.015 (B7–1/CTLA-4) for Pinnacle’s contextualized representations ($p$-value = 0.0299; Supplementary Figure S7), yet only 0.003 (PD-1/PD-L1) and 0.006 (B7–1/CTLA-4) for context-free representations (Figure 3f and Supplementary Figure S7). Incorporating information about biological contexts can help better distinguish protein interactions from non-interacting proteins in specific cell types, suggesting that Pinnacle’s contextualized representations can enhance protein representations derived from 3D protein structure modality. Modeling context-dependent interactions involving immune checkpoint proteins can deepen our understanding of how these proteins are used in cancer immunotherapies. Our benchmarking results further suggest that incorporating context can improve 3D structure prediction of protein interactions (Supplementary Note S2).

Contextual models outperform context-free target prediction.

With the representations from Pinnacle infused with cellular and tissue context, we can fine-tune them for downstream tasks (Figure 1f–h). We hypothesize that Pinnacle’s contextualized latent space can better differentiate between therapeutic targets and proteins with no therapeutic potential than a context-free latent space. Here, we focus on modeling the therapeutic potential of proteins across cell types for therapeutic areas with cell type-specific mechanisms of action (MoA) (Figure 4). Certain cell types are known to play crucial and distinct roles in the disease pathogenesis of rheumatoid arthritis (RA) and inflammatory bowel disease (IBD) therapeutic areas [24, 37–40]. There is currently no cure for either type of condition, and the medications prescribed to mitigate the symptoms can lead to undesired side effects [41]. The new generation of therapeutics in development for RA and IBD conditions is designed to target specific cell types so that the drugs maximize efficacy and minimize adverse events (e.g., by directly impacting the affected/responsible cells and avoiding off-target effects on other cells) [41, 42]. We adopt Pinnacle models to predict the therapeutic potential of proteins in a cell type-specific manner.

Figure 4: Fine-tuning contextualized protein representations for therapeutic target prioritization.

(a) Workflow to curate positive training examples for rheumatoid arthritis (left) and inflammatory bowel disease (right) therapeutic areas. (b) We construct positive examples by selecting proteins from our protein-protein interaction network (PPIN) that are targeted by compounds that have at least completed phase 2 for treating the therapeutic area of interest. These proteins are deemed safe and potentially efficacious for humans with the disease. We construct negative examples by selecting proteins from our PPIN that do not have associations with the therapeutic area yet have been targeted by at least one existing drug/compound. (c) Cell type-specific protein interaction networks are embedded by Pinnacle, and finetuned for a downstream task. Here, the predictor module (i.e., multi-layer perceptron) finetunes the (pretrained) contextualized protein representations for predicting whether a given protein is a strong candidate for the therapeutic area of interest. Additional insights of our setup include hypothesizing highly predictive cell types for examining candidate therapeutic targets. (d-e) Benchmarking of context-aware and context-free approaches for (d) RA and (e) IBD therapeutic areas. Each dot is the performance (averaged across 10 random seeds) of protein representations from a given context (i.e., cell type context for Pinnacle, context-free global reference protein interaction network for GAT and random walk, and context-free multi-modal protein interaction network for BIONIC).

We fine-tune Pinnacle to predict therapeutic targets for RA and IBD diseases. Specifically, we perform binary classification on each contextualized protein representation, where $y=1$ indicates that the protein is a therapeutic candidate for the given therapeutic area and $y=0$ otherwise. The ground truth positive examples (where $y=1$) are proteins targeted by drugs that have at least completed one clinical trial of phase two or higher for indications under the therapeutic area of interest, indicating that the drugs are safe and potentially efficacious in an initial cohort of humans (Figure 4a–b). The negative examples (where $y=0$) are druggable proteins that have not been studied for the therapeutic area (Figure 4b; Methods 5). The binary classification model can be of any architecture; our results for nominating RA and IBD therapeutic targets are generated by a multi-layer perceptron trained for each therapeutic area (Figure 4c).

To evaluate Pinnacle’s contextualized protein representations, we compare Pinnacle’s fine-tuned models against three context-free models. We apply a random walk algorithm [43] and a graph attention network (GAT) [44] on the context-free reference protein interaction network. The BIONIC model is a graph convolutional neural network designed for (context-free) multi-modal network integration [15].

We find that Pinnacle’s protein representations for all cell type contexts outperform the random walk model for both RA (Figure 4d) and IBD (Figure 4e) diseases. Protein representations from 44.9% (70 out of 156) and 37.5% (57 out of 152) cell types outperform the GAT model for RA (Figure 4d) and IBD (Figure 4e) diseases, respectively. Although both Pinnacle and BIONIC can integrate the 156 cell type-specific protein interaction networks, Pinnacle’s protein representations outperform BIONIC [15] in 18.6% of cell types (29 out of 156) and 8.6% of cell types (13 out of 152) for RA (Figure 4d) and IBD diseases (Figure 4e), respectively, highlighting the utility of contextualizing protein representations. Pinnacle outperforms these three context-free models via other metrics for both RA and IBD therapeutic areas (Supplementary Figure S8). We have confirmed no significant correlation between the node degree of proteins in cell type-specific PPI networks and performance in RA and IBD models (Supplementary Figure S9a). Additionally, there is only a moderate correlation between Pinnacle’s performance and the enrichment of positive targets in these cell type-specific PPI networks (Supplementary Figure S9b–c). These findings underscore that Pinnacle’s predictions cannot be solely ascribed to the characteristics of the cell type-specific PPI networks. Benchmarking results indicate combining global reference networks with advanced deep graph representation learning techniques, such as GAT, can yield better predictors than network-based random walk methods alone. Integrative approaches, exemplified by methods like BIONIC, enhance performance, a finding consistent with the established benefits of data integration. Contextualized learning approaches, like Pinnacle, have the potential to enhance model performance and enable predictions tailored to specific contexts.

Pinnacle can nominate targets across cell type contexts.

There is existing evidence that drug effects vary with cell type depending on where therapeutic targets are expressed and where proteins act [45–49]. For instance, CD19-targeting chimeric antigen receptor T (CAR-T) cell therapy has been highly effective in treating B cell malignancies, yet causes a high incidence of neurotoxicity [47]. A recent study shows that CAR-T cells induce off-target effects by targeting the CD19 expressed in brain mural cells, likely causing the brain barrier leakiness responsible for neurotoxicity [47]. We hypothesize that the predicted protein druggability varies across cell types, and such variations can provide insights into the cell types’ relevance for a therapeutic area.

Among the 156 biological contexts modeled by Pinnacle’s protein representations, we examine the most predictive cell type contexts for nominating therapeutic targets of RA. We find that the most predictive contexts consist of CD4+ helper T cells, CD4+ $\alpha \beta$ memory T cells, CD1c+ myeloid dendritic cells, gut endothelial cells, and pancreatic acinar cells (Figure 5a). Immune cells play a significant role in the disease pathogenesis of RA [37, 38]. Since CD4+ helper T cells (Pinnacle-predicted rank = 1), CD4+ $\alpha \beta$ memory T cells (Pinnacle-predicted rank = 2), and CD1c+ myeloid dendritic cells (Pinnacle-predicted rank = 3) are immune cells, it is expected that Pinnacle’s protein representations in these contexts achieve high performance in our prediction task. Also, patients with RA often have gastrointestinal (GI) manifestations, whether concomitant GI autoimmune diseases or GI side effects of RA treatment [50]. Pancreatic acinar cells (Pinnacle-predicted rank = 5) can behave like inflammatory cells during acute pancreatitis [51], one of the accompanying GI manifestations of RA [50]. In addition to GI manifestations, endothelial dysfunction is commonly detected in patients with RA [52]. While rare, rheumatoid vasculitis, which affects endothelial cells and is a serious complication of RA, has been found to manifest in the large and small intestines (gut endothelial cell context has Pinnacle-predicted rank = 4), liver, and gallbladder [50, 53]. Further, many of the implicated cell types for RA patients (e.g., T cells, B cells, natural killer cells, monocytes, myeloid cells, and dendritic cells) are highly ranked by Pinnacle [24, 25, 39] (Supplementary Table S1). Our results suggest that injecting cell type context to protein representations can significantly improve performance in nominating therapeutic targets for RA diseases while potentially revealing the cell types underlying disease processes.

Figure 5: Performance of contextualized target prioritization for RA and IBD therapeutic areas.

(a,d) Model performance (measured by APR@5) for RA and IBD therapeutic areas, respectively. APR@K (or Average Precision and Recall at K) is a combination of Precision@K and Recall@K (refer to Methods 6 for more details). Each dot is the performance (averaged across 10 random seeds) of Pinnacle’s protein representations from a specific cell type context. The gray and dark orange lines are the performance of the GAT and BIONIC models, respectively. For each therapeutic area, 22 cell types are annotated and colored by their compartment category. Supplementary Figure S8 contains model performance measured by APR@10, APR@15, and APR@20 for RA and IBD therapeutic areas. (b-c, e-f) Selected proteins for RA and IBD therapeutic areas. Dotted line separates the top and bottom 5 cell types. (b-c) Two selected proteins, JAK3 and IL6R, that are targeted by drugs that have completed Phase IV of clinical trials for treating RA therapeutic area. (e-f) Two selected proteins, ITGA4 and PPARG, that are targeted by drugs that have completed Phase IV for treating IBD therapeutic area.

The most predictive cell type contexts for nominating therapeutic targets of IBD are CD4+ $\alpha \beta$ memory T cells, enterocytes of epithelium of large intestine, T follicular helper cells, plasmablasts, and myeloid dendritic cells (Figure 5d). The intestinal barrier comprises a thick mucus layer with antimicrobial products, a layer of intestinal epithelial cells, and a layer of mesenchymal cells, dendritic cells, lymphocytes, and macrophages [54]. As such, these five cell types are expected to yield high predictive ability. Moreover, many of the implicated cell types for IBD (e.g., T cells, fibroblasts, goblet cells, enterocytes, monocytes, natural killer cells, B cells, and glial cells) are highly ranked by Pinnacle [26, 27, 55] (Supplementary Table S2). For example, CD4+ T cells are known to be the main drivers of IBD [56]. They have been found in the peripheral blood and intestinal mucosa of adult and pediatric IBD patients [57]. Patients with IBD tend to develop uncontrolled inflammatory CD4+ T cell responses, resulting in tissue damage and chronic intestinal inflammation [58, 59]. Due to the heterogeneity of CD4+ T cells in patients, treatment efficacy can depend on the patient’s subtype of CD4+ T cells [58, 59]. Thus, the highly predictive cell type contexts according to Pinnacle should be further investigated to design safe and efficacious therapies for RA and IBD diseases.

Conversely, we hypothesize that the cell type contexts of protein representations that yield worse performance than the cell type-agnostic protein representations may not have the predictive power (given the current list of targets from drugs that have at least completed phase 2 of clinical trials) for studying the therapeutic effects of candidate targets for RA and IBD therapeutic areas.

In the context-aware model trained to nominate therapeutic targets for RA diseases, the protein representations of duodenum glandular cells, endothelial cells of hepatic sinusoid, myometrial cells, and hepatocytes performance worse than the cell type-agnostic protein representations (Figure 5a). The RA therapeutic area is a group of inflammatory diseases in which immune cells attack the synovial lining cells of joints [37]. Since duodenum glandular cells (Pinnacle-predicted rank = 153), endothelial cells of hepatic sinusoid (Pinnacle-predicted rank = 126), myometrial cells (Pinnacle-predicted rank = 119), and hepatocytes (Pinnacle-predicted rank = 116) are neither immune cells nor found in the synovium, these cell type contexts’ protein representations expectedly perform poorly. For IBD diseases, the protein representations of the limbal stem cells, melanocytes, fibroblasts of cardiac tissue, and radial glial cells have worse performance than the cell type-agnostic protein representations (Figure 5d). The IBD therapeutic area is a group of inflammatory diseases in which immune cells attack tissues in the digestive tract [40]. As limbal stem cells (Pinnacle-predicted rank = 152), melanocytes (Pinnacle-predicted rank = 147), fibroblasts of cardiac tissue (Pinnacle-predicted rank = 135), and radial glial cells (Pinnacle-predicted rank = 107) are neither immune cells nor found in the digestive tract, these cell type contexts’ protein representations should also perform worse than context-free representations.

The least predictive cellular contexts in Pinnacle’s models for RA and IBD have no known role in disease, indicating that protein representations from these cell type contexts are poor predictors of RA and IBD therapeutic targets. Pinnacle’s overall improved predictive ability compared to context-free models indicates the importance of understanding cell type contexts where therapeutic targets are expressed and act.

Predictive cell type contexts reflect MoAs in RA therapies.

Recognizing and leveraging the most predictive cell type context for examining a therapeutic area can be beneficial for predicting candidate therapeutic targets [45–49]. We find that considering only the most predictive cell type contexts can yield significant performance improvements compared to context-free models (Supplementary Figure S10). We examine cell type contexts selected by Pinnacle as the most predictive for JAK3 and IL6R, two protein targets of RA drugs.

Disease-modifying anti-rheumatic drugs (DMARDs), such as Janus kinase (JAK) inhibitors (i.e., tofacitinib, upadacitinib, and baricitinib), are commonly prescribed to patients with RA [60, 61]. For JAK3, Pinnacle’s five most predictive cell type contexts are T follicular helper cells, microglial cells, DN3 thymocytes, CD4+ $\alpha \beta$ memory T cells, and hematopoietic stem cells (Figure 5b). Since the expression of JAK3 is limited to hematopoietic cells, mutations or deletions in JAK3 tend to cause defects in T cells, B cells, and NK cells [62–65]. For instance, patients with JAK3 mutations tend to be depleted of T cells [63], and the abundance of T follicular helper cells is highly correlated with RA severity and progression [66]. JAK3 is also highly expressed in double negative (DN) T cells (early stage of thymocyte differentiation) [67], and the levels of DN T cells are higher in synovial fluid than peripheral blood, suggesting a possible role of DN T cell subsets in RA pathogenesis [68]. Lastly, dysregulation of the JAK/STAT pathway, which JAK3 participates in, has pathological implications for neuroinflammatory diseases, a significant component of disease pathophysiology in RA [69, 70].

Tocilizumab and sarilumab are approved by the Food and Drug Administration (FDA) for treating RA, and target the interleukin six receptor, IL6R [61]. For IL6R, Pinnacle’s five most predictive cellular contexts are classical monocytes, NAMPT neutrophils, intermediate monocytes, mesenchymal stem cells, and regulatory T cells (Figure 5c). IL6R is predominantly expressed on neutrophils, monocytes, hepatocytes, macrophages, and some lymphocytes [71]. IL6R simulates the movement of T cells and other immune cells to the site of infection or inflammation [72] and affects T cell and B cell differentiation [71, 73]. IL6 acts directly on neutrophils, essential mediators of inflammation and joint destruction in RA, through membrane-bound IL6R [71]. Experiments on fibroblasts isolated from the synovium of RA patients show that anti-IL6 antibodies prevented neutrophil adhesion, indicating a promising therapeutic direction for IL6R on neutrophils [71]. Lastly, mice studies have shown that pre-treatment of mesenchymal stem/stromal cells with soluble IL6R can enhance the therapeutic effects of mesenchymal stem/stromal cells in arthritis inflammation [74].

Pinnacle’s hypotheses to examine JAK3 and IL6R in the highly predictive cell type contexts, according to Pinnacle, to maximize therapeutic efficacy seem to be consistent with their roles in the cell types. It seems that targeting these proteins may directly impact the pathways contributing to the pathophysiology of RA therapeutic areas. Further, our results for IL6R suggest that Pinnacle’s contextualized representations could be leveraged to evaluate potential enhancement in efficacy (e.g., targeting multiple points in a pathway of interest).

Predictive cell type contexts elucidate MoAs in IBD therapies.

Like RA, we must understand the cells in which therapeutic targets are expressed and act to maximize the efficacy of targeted IBD therapies [75]. To support our hypothesis, we evaluate PIN-NACLE’s predictions for two protein targets of commonly prescribed treatments for IBD diseases: ITGA4 and PPARG.

Vedolizumab and natalizumab target the integrin subunit alpha 4, ITGA4, to treat the symptoms of IBD therapeutic area [61]. Pinnacle’s five most predictive cell type contexts for ITGA4 are regulatory T cells, dendritic cells, myeloid dendritic cells, granulocytes, and CD8+ $\alpha \beta$ cytotoxic T cells (Figure 5e). Integrins mediate the trafficking and retention of immune cells to the gastrointestinal tract; immune activation of integrin genes increases the risk of IBD [76]. For instance, ITGA4 is involved in homing memory and effector T cells to inflamed tissues, including intestinal and non-intestinal tissues, and imbalances in regulatory and effector T cells may lead to inflammation [77]. Circulating dendritic cells express the gut homing marker encoded by ITGA4; the migration of blood dendritic cells to the intestine allows these dendritic cells to become mature, which leads to gut inflammation and tissue damage, indicating that future studies are warranted to elucidate the functional properties of blood dendritic cells in IBD [78].

Balsalazide and mesalamine are aminosalicylate drugs (DMARDs) commonly used to treat ulcerative colitis by targeting peroxisome proliferator activated receptor gamma (PPARG) [61, 79]. Pinnacle’s five most predictive cell types for PPARG are paneth cells of the epithelium of large intestines, endothelial cells of the vascular tree, classic monocytes, goblet cells of small intestines, and serous cells of epithelium of bronchus (Figure 5f). PPARG is highly expressed in the gastrointestinal tract, higher in the large intestine (e.g., colonic epithelial cells) than the small intestine [80–82]. In patients with ulcerative colitis, PPARG is often substantially downregulated in their colonic epithelial cells [82]. PPARG promotes enterocyte development [83] and intestinal mucus integrity by increasing the abundance of goblet cells [82]. Further, PPARG activation can inhibit endothelial inflammation in vascular endothelial cells [84, 85], which is significant due to the importance of vascular involvement in IBD [86]. Additionally, PPARG agonists have been shown to act as negative regulators of monocytes and macrophages, which can inhibit the production of proinflammatory cytokines [87]. Intestinal mononuclear phagocytes, such as monocytes, play a major role in maintaining epithelial barrier integrity and fine-tuning mucosal immune system responsiveness [88]. Studies show that newly recruited monocytes in inflamed intestinal mucosa drive the immunopathogenesis of IBD, suggesting that blocking monocyte recruitment to the intestine could be one avenue for therapeutic development [88]. Lastly, PPARG is found to regulate mucin and inflammatory factors in bronchial epithelial cells [89]. Given the pulmonary complications of IBD, PPARG could be a promising target to investigate for treating IBD and pulmonary symptoms [90]. The predictive power of cell type contexts to examine ITGA4 and PPARG, according to Pinnacle, for IBD therapeutic development are thus well-supported.

Discussion

Pinnacle is a flexible geometric deep learning approach for contextualized prediction in user-defined biological contexts. Integrating single-cell transcriptomic atlases with the protein interactome, cell type interactions, and tissue hierarchy, Pinnacle produces latent protein representations specialized to biological contexts. Pinnacle’s protein representations capture cellular and tissue organization spanning 156 cell types and 62 tissues of varying hierarchical scales. In addition to multi-modal data integration, a pretrained Pinnacle model generates protein representations that can be used for downstream prediction on tasks where cell type dependencies and cell type-specific mechanisms are relevant.

One limitation of the study is the use of the human protein interactome, which is not measured in a cell type-specific manner [91]. No systematic measurements of protein interactions across cell types exist. We create cell type-specific protein interaction networks by overlaying single-cell measurements on the protein interaction network, leveraging previously validated techniques for the reconstruction of cell-type-specific interactomes at single-cell resolution [14] and conducting sensitivity network analyses to confirm the validity of the networks used to train Pinnacle models (Supplementary Figures S2–S3). This approach enriches networks for cell type-relevant proteins (Supplementary Figure S2). The resulting networks may contain false-positive protein interactions (e.g., proteins that interact in the reference protein interaction network but do not interact in a specific cell type) and false-negative protein interactions (e.g., proteins that interact only within a particular cell type context that has not yet been measured). Pinnacle does not currently model proteins that may play a role in the cell type yet are unaffected by cell type specificity. Nevertheless, strong performance gains of Pinnacle over context-free models indicate the importance of contextualized prediction and suggest a direction to enhance existing analyses on protein interaction networks [4, 6, 7].

We can leverage and extend Pinnacle in many ways. Pinnacle can accommodate and supplement diverse data modalities. We developed Pinnacle models using Tabula Sapiens [20], a molecular reference atlas comprising almost 500,000 cells from 24 distinct tissues and organs. However, since the tissues and cell types associated with specific diseases may not be entirely represented in the atlas of healthy human subjects, we anticipate that our predictive power may be limited. Tabula Sapiens does not include synovial tissues associated with RA disease progression [25, 39], but these can be found in synovial RA atlases [92] and stromal cells obtained from individuals with chronic inflammatory diseases [93]. To enhance the predictive ability of Pinnacle models, they can be trained on disease-specific or perturbation-specific networks. In this study, Pinnacle representations capture physical interactions between proteins at the cell type level (Supplementary Note S3); Pinnacle can also be applied to cell type-specific protein networks created from other modalities, such as cell type-specific gene expression networks [94]. We show that Pinnacle’s representations can supplement protein representations generated from other data modalities, including protein 3D structure surfaces [3, 17]. While this study focuses on protein-coding genes, information on protein isoforms and differential information, such as alternative splicing or allosteric changes, can be used with Pinnacle when such data are broadly available. In addition to prioritizing candidate therapeutic targets, Pinnacle’s representations can be fine-tuned to identify populations of cells with specific characteristics, such as drug resistance [95], adverse drug events [96], or disease progression biomarkers [97]. Lastly, to move towards a “labin-the-loop” framework, where computational and experimental scientists can iteratively refine the machine learning model and validate hypotheses via experiments, recent techniques on conformal prediction [98] and evidential layers can be integrated with Pinnacle to quantify the uncertainty of model outputs.

Protein representation learning models are context-free and are limited in analyzing protein phenotypes that are resolved by contexts and vary with cell types and tissues. To address this limitation, we introduce Pinnacle that produces protein representations tailored to cell type contexts. We demonstrate that contextual learning can provide a more comprehensive understanding of protein roles across cell type contexts [99]. As experimental technologies advance, it is becoming feasible to generate adaptive protein representations across cell type contexts and leverage contextualized representations to predict cell type specific protein functions and nominate therapeutic candidates at the cell type level. Looking to the future, understanding protein functions and developing molecular therapies will require a comprehensive understanding of the roles that proteins have in different cell types and the interactions between proteins across diverse cell type contexts [100]. Approaches like Pinnacle can help realize this potential by generating contextualized protein representations, which can then be used to predict cell type specific protein functions and identify therapeutic targets at the cellular level.

Online Methods

The Methods describe (1) the curation of datasets, (2) the construction and representation of multi-scale single-cell networks, (3) Pinnacle, a multi-scale graph neural network, (4) the finetuning of Pinnacle for target prioritization, and (5) the metrics and statistical analyses used.

1. Datasets

Reference human physical protein interaction network. Our reference protein-protein interaction (PPI) network is the union of physical multi-validated interactions from BioGRID [101, 102], the Human Reference Interactome (HuRI) [91], and Menche et al. [103] with 15,461 nodes and 207,641 edges. Different sources of PPI have their own methods of curating and validating physical interactions between proteins. BioGRID, HuRI, and Menche et al. are PPI networks from three well-cited publications and databases regarding human protein interactions. By joining the three networks, we construct a comprehensive global PPI network for our analysis.

Multi-organ, single-cell transcriptomic atlas of humans.

We leverage Tabula Sapiens [20] data source as our multi-organ, single-cell transcriptomic atlas of humans. The data consists of 15 donors, with 59 specimens total. There are 483,152 cells after quality control, of which 264,824 are immune cells, 104,148 are epithelial cells, 31,691 are endothelial cells, and 82,478 are stromal cells. The cells correspond to 177 unique cell ontology classes.

2. Construction of multi-scale networks

Our multi-scale networks comprises protein-protein physical interactions, cell type to cell type communication, cell type to tissue relationships, and tissue-tissue hierarchy.

Cell type-specific protein interaction networks.

For each cell type, we create a cell type specific network that represents the physical interactions between proteins (or genes) that are likely expressed in the cell type. Intuitively, our approach identifies genes significantly expressed in a given cell type with respect to the rest of the cells in the dataset. Concretely, we use the processed Tabula Sapiens count matrix to calculate the average expression of each gene in a cell type of interest and the average expression of the corresponding gene in all other cells. Then, we use the Wilcoxon rank-sum test on the two sets of average gene expression. From the resulting ranked list of genes based on activation, we filter for the top K most activated genes. We repeat these two steps N times and filter for genes that appear in at least 90% of iterations. Finally, we extract these genes’ corresponding proteins from the global protein interaction network and take only the largest connected component. To ensure high-quality representations of cell types in our networks, we keep networks with at least 1,000 proteins. We do not perform subsampling of cells (i.e., sample the same number of cells per cell type) to minimize information loss for constructing protein interaction networks (Supplementary Figure S2). Limitations are described in the Discussion section.

Cell type and tissue relationships in the metagraph.

We identify interactions between cell types based on ligand-receptor (LR) expression using the CellPhoneDB [104] tool and database. An edge between a pair of cell types indicates that CellphoneDB predicts at least one significantly expressed LR pair (with a $p$-value of less than 0.001) between them. As recommended by CellPhoneDB, cells are subsampled prior to running the algorithm, which uses geometric sketching [105] to efficiently sample a small representative subset of cells from massive datasets while preserving biological complexity. We choose to subsample 25% of cells and run CellPhoneDB for 100 iterations. We determine cell type-tissue relationships and extract tissue-tissue relationships using Tabula Sapiens meta-data. For relationships between cell types and tissues, we draw edges between cell types and the tissue that the cells were taken from. For tissue-tissue relationships, we select the nodes corresponding to the tissues where samples were taken from and all parent nodes up to the root of the BRENDA tissue ontology [106]. We perform sensitivity and ablation analyses on different components of the metagraph (Supplementary Table S3–S5).

Final dataset.

We have 156 cell type specific protein interaction networks, which have, on average, 2, 530±677 proteins per network. The number of unique proteins across all cell type specific protein interaction networks is 13, 643 of the 15, 461 proteins in the global reference protein interaction network. In the metagraph, we have 62 tissues (nodes), and 24 are directly connected to cell types. There are 3,567 cell-cell interactions, 372 cell-tissue edges, and 79 tissue-tissue edges.

3. Multi-scale graph neural network

Overview.

Pinnacle performs biologically-informed message passing through proteins, cell types, and tissues to learn cell type specific protein representations, cell type representations, and tissue representations in a unified multi-scale embedding space. Specifically, Pinnacle traverses through protein-protein physical interactions in each cell type specific PPI network, cell type-cell type communication, cell type-tissue relationships, and tissue-tissue hierarchy with an attention mechanism over individual nodes and edge types. Its objective function is designed and optimized for learning the topology across biological scales, from proteins to cell types to tissues. The resulting embeddings from Pinnacle can be visualized and manipulated for hypothesis-driven interrogation and finetuned for diverse downstream biomedical prediction tasks.

Problem formulation.

Let $𝒢=\left\{G_1,\dots ,G_{\left|𝒞\right|}\right\}$ be a set of cell type specific protein-protein interaction networks, where $𝒞$ is a set of unique cell types. Each $G_{c_i}=\left(V_{c_i},E_{c_i}\right)$ consists of a set of nodes $V_{c_i}$ and edges $E_{c_i}$ for a given cell type $c_i\in 𝒞$ specific protein-protein interaction network. Their nodes $u,v\in V_{c_i}$ are proteins, and edges $e_{u,v}^{\text{PP}}\in E_{c_i}$ are physical protein-protein interactions (denoted with pp in the superscript). Cell types and tissues form a network, referred to as a metagraph. The metagraph’s set of nodes comprises cell types $c_i\in 𝒞$ and tissues $t_i\in 𝒯$. The types of edges are cell type-cell type interactions (denoted with cc in the superscript) $e_{c_i,c_j}^{\text{CC}}$ between any pair of cell types $c_i,c_j\in 𝒞$; cell type-tissue associations (denoted with ct in the superscript) $e_{c_i,t_i}^{\mathrm{C}\mathrm{T}}$ between any pair of cell type $c_i\in 𝒞$ and tissue $t_i\in 𝒯$; and tissue-tissue relationships (denoted with tt in the superscript) $e_{t_i,t_j}^{\mathrm{T}\mathrm{T}}$ between any pair of tissues $t_i,t_j\in 𝒯$.

3.1. Protein-level attention with cell type specificity

For each cell type specific protein-protein interaction network ${𝒢}_{c_i}$, we leverage protein-level attention to learn cell type specific embeddings of proteins. Intuitively, protein-level attention learns which neighboring nodes are likely most important for characterizing a particular cell type’s protein. As such, each cell type specific protein interaction network has its own cell type specific set of learnable parameters. Concretely, at each layer-wise update of layer $l$, the node-level attention learns the importance ${\alpha }_{u,v}$ of protein $u$ to its neighboring protein $v$ in a given cell type $c_i\in 𝒞$ :

$$h_u^{\text{PP}}\leftarrow \text{AGG}\left(\sigma \left(\sum _{v\in {𝒩}_u}{\alpha }_{u,v}{W}^{\text{PP}}{h}_v^{\text{PP}}\right)\right)$$ (1)

where AGG is an aggregation function (i.e., concatenation across $K$ attention heads), $\sigma$ is the nonlinear activation function (i.e., ReLU), ${𝒩}_u$ is the set of neighbors for $u$ (including itself via self-attention), ${\alpha }_{u,v}$ is an attention mechanism defined as ${\alpha }_{u,v}=\frac{\text{exp}\left(\sigma \left({\mathbf{a}}^T\cdot \left[{\mathbf{h}}_u\parallel {\mathbf{h}}_v\right]\right)\right)}{{\sum }_{v\in {𝒩}_u}\text{exp}\left(\sigma \left({\mathbf{a}}^T\cdot \left[{\mathbf{h}}_u\parallel {\mathbf{h}}_v\right]\right)\right)}$ between a pair of interacting proteins from a specific cell type, ${\mathbf{W}}^{\text{PP}}$ is a pp-specific transformation matrix to project the features of protein $u$ in its cell type specific protein interaction network, and ${\mathbf{h}}_v^{\text{PP}}$ is the previous layer’s cell type specific embedding for protein $v$. Practically, we leverage the attention function in graph attention neural networks (i.e., GATv2) [44]. Proteins of the same identity are initialized with the same random Gaussian vector to maintain their identity during training.

3.2. Metagraph-level attention on cellular interactions and tissue hierarchy

For the metagraph, we use node-level and edge-level attention to learn which neighboring nodes and edge types are likely most important for characterizing the target node (i.e., the node of interest). Intuitively, to learn an embedding for a specific cell type or tissue, we evaluate the informativeness of each direct cell type or tissue neighbor, as well as the relationship type between the cell type or tissue and their neighbors (e.g., parent-child tissue relationship, tissue from which a cell type is found, cell type with which the cell type of interest communicates with).

Concretely, at each layer $l$ of Pinnacle, the embeddings of a cell type $c_i\in 𝒞$ are the result of aggregating (via function AGG) the embeddings (${\mathbf{h}}_c^{CC}$ and ${\mathbf{h}}_t^{CT}$) of its direct cell type neighbor $c$ and tissue neighbor $t$ that are projected via edge-type-specific transformation matrices (${\mathbf{W}}^{\text{CC}}$ and ${\mathbf{W}}^{\mathrm{C}\mathrm{T}}$) and weighted by learned attention weights (${\alpha }_{c_i,c}$ and ${\alpha }_{c_i,t}$ respectively):

$$h_{c_i}^{\text{CC}}\leftarrow \text{AGG}\left(\sigma \left(\sum _{c\in {𝒩}_{c_i}}{\alpha }_{c_i,c}{W}^{\text{CC}}{h}_c^{\text{CC}}\right)\right)$$ (2)

$$h_{c_i}^{\text{CT}}\leftarrow \text{AGG}\left(\sigma \left(\sum _{t\in {𝒩}_{c_i}}{\alpha }_{c_i,t}{W}^{\text{CT}}{h}_t^{\text{CT}}\right)\right)$$ (3)

The embeddings generated from separately propagating messages through cell type-cell type edges or cell type-tissue edges are combined using learned attention weights ${\beta }^{\text{CC}}$ and ${\beta }^{\mathrm{C}\mathrm{T}}$, respectively:

$${\mathbf{h}}_{c_i}={\beta }^{\text{CC}}{\mathbf{h}}_{c_i}^{\text{CC}}+{\beta }^{\mathrm{C}\mathrm{T}}{\mathbf{h}}_{c_i}^{\mathrm{C}\mathrm{T}}$$ (4)

Similarly, the embeddings of a tissue $t_i\in 𝒯$ are the result of aggregating (via function AGG) the embeddings (${\mathbf{h}}_t^{\mathrm{T}\mathrm{T}}$ and ${\mathbf{h}}_c^{\mathrm{T}\mathrm{C}}$) of its direct tissue neighbor $t$ and cell type neighbor $c$ that are projected via edge-type-specific transformation matrices $\left({\mathbf{W}}^{\mathrm{T}\mathrm{T}}\right.$ and ${\mathbf{W}}^{\mathrm{T}\mathrm{C}}$) and weighted by learned attention weights (${\alpha }_{t_i,t}$ and ${\alpha }_{t_i,c}$ respectively):

$$h_{t_i}^{\text{TT}}\leftarrow \text{AGG}\left(\sigma \left(\sum _{t\in {𝒩}_{t_i}}{\alpha }_{t_i,t}{W}^{\text{TT}}{h}_t^{\text{TT}}\right)\right)$$ (5)

$$h_{t_i}^{\text{TC}}\leftarrow \text{AGG}\left(\sigma \left(\sum _{c\in {𝒩}_{t_i}}{\alpha }_{t_i,c}{W}^{\text{TC}}{h}_c^{\text{TC}}\right)\right)$$ (6)

The embeddings generated from separately propagating messages through tissue-tissue edges or tissue-cell type edges are combined using learned attention weights ${\beta }^{\mathrm{T}\mathrm{T}}$ and ${\beta }^{\mathrm{T}\mathrm{C}}$, respectively:

$${\mathbf{h}}_{t_i}={\beta }^{\mathrm{T}\mathrm{T}}{\mathbf{h}}_{t_i}^{\mathrm{T}\mathrm{T}}+{\beta }^{\mathrm{T}\mathrm{C}}{\mathbf{h}}_{t_i}^{\mathrm{T}\mathrm{C}}$$ (7)

For the node-level attention mechanisms (Equations 2, 3, 5, and 6), AGG is an aggregation function (i.e., concatenation across $K$ attention heads), $\sigma$ is the nonlinear activation function (i.e., ReLU), ${𝒩}_{c_i}$ and ${𝒩}_{t_i}$ are the sets of neighbors for $c_i$ and $t_i$ respectively (includes itself via self-attention), ${\mathbf{W}}^{\text{CC}},{\mathbf{W}}^{\mathrm{C}\mathrm{T}},{\mathbf{W}}^{\mathrm{T}\mathrm{C}}$, and ${\mathbf{W}}^{\mathrm{T}\mathrm{T}}$ are edge-type-specific transformation matrices to project the features of a given target node, ${\mathbf{h}}_c^{\text{CC}},{\mathbf{h}}_t^{\mathrm{C}\mathrm{T}},{\mathbf{h}}_t^{\mathrm{T}\mathrm{T}}$, and ${\mathbf{h}}_c^{\mathrm{T}\mathrm{C}}$ are the previous layer’s embedding for $c$ given the edge type cc $,$ $t$ given the edge type ct, $t$ given the edge type tt, and $c$ given the edge type tc, respectively. Practically, we leverage the attention function in graph attention neural networks (i.e., GATv2) [44]. Finally, the node-level attention mechanism for a given source node $u$ and edge type $r$ is ${\alpha }_{u,v}^r=\frac{\text{exp}\left(\sigma \left(a_r^T\cdot \left[h_u\parallel h_v\right]\right)\right)}{{\sum }_{v\in {𝒩}_u}\text{exp}\left(\sigma \left(a_r^T\cdot \left[h_u\parallel h_v\right]\right)\right)}$. For the attention mechanisms over edge types (Equations 4 and 7), ${\beta }^r=\frac{\text{exp}\left(m_r\right)}{{\sum }_{r\in R}\text{exp}\left(m_r\right)}$ such that $m_r={\sum }_{u\in V_q}{\mathbf{s}}^T\cdot \mathrm{t}\mathrm{a}\mathrm{n}\mathrm{h}\left(\mathbf{M}\cdot {\mathbf{h}}_u^r+\mathbf{b}\right)$ where $V_q$ is the set of nodes in the metagraph, $\mathbf{s}$ is the attention vector, $\mathbf{M}$ is the weight matrix, and $\mathbf{b}$ is the bias vector. These parameters are shared for all edge types in the metagraph.

3.3. Bridge between protein and cell type embeddings

Using a pooling mechanism, we bridge cell type specific protein embeddings with their corresponding cell type embeddings. We initialize cell type embeddings by taking the average of their proteins’ embeddings: ${\mathbf{h}}_{c_i}=\frac{1}{\left|V_{c_i}\right|}{\sum }_{u\in V_{c_i}}{\mathbf{h}}_u$ where ${\mathbf{h}}_u$ is the embedding of protein node $u\in V_{c_i}$ in the PPI subnetwork for cell type $c_i$. Similarly, we initialize tissue embeddings by taking the average of their neighbors: ${\mathbf{h}}_{t_i}=\frac{1}{\left|{𝒩}_{t_i}\right|}\left({\sum }_{t\in {𝒩}_{t_i}}{\mathbf{h}}_t+{\sum }_{c\in {𝒩}_{t_i}}{\mathbf{h}}_c\right)$, where $h_t$ and $h_c$ are the embeddings of tissue node $t$ and cell type node $c$, respectively, in the immediate neighborhood of source tissue node $t_i$. At each layer $l>0$, we learn the importance ${\gamma }_{c_i,u}$ of node $u\in V_{c_i}$ to cell type $c_i$ such that

$$h_{c_i}\leftarrow h_{c_i}+\text{AGG}\left(\sigma \left(\sum _{u\in V_{c_i}}{\gamma }_{c_i,u}{h}_u\right)\right)$$ (8)

After propagating cell type and tissue information in the metagraph (namely, Equations 2–6), we apply ${\gamma }_{c_i,u}$ to the cell type embedding of $c_i$ such that

$${\mathbf{h}}_u\leftarrow {\mathbf{h}}_u+{\gamma }_{c_i,u}{\mathbf{h}}_{c_i}$$ (9)

Intuitively, we are imposing the structure of the metagraph onto the PPI subnetworks based on a protein’s importance to its corresponding cell type’s identity.

3.4. Pinnacle: Overall objective function

PInNACLE is optimized for three biological scales: protein-, cell type-, and tissue-level. Concretely, the loss function $ℒ$ has three components corresponding to each biological scale:

$$ℒ={ℒ}_{\mathit{\text{protein}}}+(1-\theta )\left({ℒ}_{\mathit{\text{celltype}}}+{ℒ}_{\mathit{\text{tissue}}}\right)$$ (10)

where ${ℒ}_{\mathit{\text{protein}}},{ℒ}_{\mathit{\text{celltype}}}$, and ${ℒ}_{\mathit{\text{tissue}}}$ minimize the loss from protein-level predictions, cell type-level predictions, and tissue-level predictions, respectively. $\theta$ is a tunable parameter with a range of 0 and 1 that scales the contribution of the link prediction loss of the metagraph relative to that of the protein-protein interactions. At the protein level, we consider two aspects: prediction of protein-protein interactions at each cell type specific PPI network $\left({ℒ}_{\mathit{\text{ppi}}}\right.$) and prediction of cell type identity of each protein $\left({ℒ}_{\mathit{\text{celltypeid}}}\right)$. The contribution of the latter is scaled by $\lambda$, which is a tunable parameter with a range of 0 and 1 :

$${ℒ}_{\mathit{\text{protein}}}=\theta {ℒ}_{\mathit{\text{ppi}}}+\lambda {ℒ}_{\mathit{\text{celltypeid}}}$$ (11)

Intuitively, we aim to simultaneously learn the topology of each cell type specific PPI network (i.e., ${ℒ}_{\mathit{\text{ppi}}}$) and the nuanced differences between proteins activated in different cell types. Specifically, we use binary cross-entropy to minimize the error of predicting positive and negative protein-protein interactions in each cell type specific PPI network:

$${ℒ}_{ppi}=\sum _{c_i\in 𝒞}\sum _{u,v\in V_{c_i}}{y}_{u,v}\text{log}\left({\hat{y}}_{u,v}\right)+\left(1-y_{u,v}\right)\text{log}\left(1-{\hat{y}}_{u,v}\right)$$ (12)

and center loss [107] for discriminating between protein embeddings ${\mathbf{z}}_u$ from different cell types, represented by embeddings denoted as ${\mathbf{z}}_{c_i}$ :

$${ℒ}_{\mathit{\text{celltypeid}}}=\sum _{c_i\in 𝒞}\sum _{u\in V_{c_i}}{∥{\mathbf{z}}_u-{\mathbf{z}}_{c_i}∥}_2^2$$ (13)

At the cell type level, we use binary cross-entropy to minimize the error of predicting cell type-cell type interactions and cell type-tissue relationships:

$${ℒ}_{\mathit{\text{celltype}}}={ℒ}_{\mathit{\text{celltype}}}^{\text{CC}}+{ℒ}_{\mathit{\text{celltype}}}^{\mathrm{C}\mathrm{T}}$$ (14)

such that

$${ℒ}_{\mathit{\text{celltype}}}^{\text{CC}}=\sum _{c_i,c_j\in 𝒞}{y}_{c_i,c_j}\text{log}\left({\hat{y}}_{c_i,c_j}\right)+\left(1-y_{c_i,c_j}\right)\text{log}\left(1-{\hat{y}}_{c_i,c_j}\right)$$ (15)

$${ℒ}_{\mathit{\text{celltype}}}^{\mathrm{C}\mathrm{T}}=\sum _{c_i\in 𝒞}\sum _{t_k\in 𝒯}{y}_{c_i,t_k}\text{log}\left({\hat{y}}_{c_i,t_k}\right)+\left(1-y_{c_i,t_k}\right)\text{log}\left(1-{\hat{y}}_{c_i,t_k}\right)$$ (16)

Similarly, at the tissue level, we use binary cross-entropy to minimize the error of predicting tissue-tissue and tissue-cell type relationships:

$${ℒ}_{\mathit{\text{tissue}}}={ℒ}_{\mathit{\text{tissue}}}^{\mathrm{T}\mathrm{T}}+{ℒ}_{\mathit{\text{tissue}}}^{\mathrm{T}\mathrm{C}}$$ (17)

such that

$${ℒ}_{\mathit{\text{tissue}}}^{\mathrm{T}\mathrm{T}}=\sum _{t_k,t_q\in 𝒯}{y}_{t_k,t_q}\text{log}\left({\hat{y}}_{t_k,t_q}\right)+\left(1-y_{t_k,t_q}\right)\text{log}\left(1-{\hat{y}}_{t_k,t_q}\right)$$ (18)

$${ℒ}_{\mathit{\text{tissue}}}^{\mathrm{T}\mathrm{C}}=\sum _{t_k\in 𝒯}\sum _{c_i\in 𝒞}{y}_{t_k,c_i}\text{log}\left({\hat{y}}_{t_k,c_i}\right)+\left(1-y_{t_k,c_i}\right)\text{log}\left(1-{\hat{y}}_{t_k,c_i}\right)$$ (19)

The probability of an edge of type $i$ between nodes $u$ and $v$ is calculated using a bilinear decoder:

$$y_{u,v}={\mathbf{z}}_u\cdot {\mathbf{r}}_i\cdot {\mathbf{z}}_v$$ (20)

where ${\mathbf{z}}_u$ and ${\mathbf{z}}_v$ are embeddings of nodes $u$ and $v$, and ${\mathbf{r}}_i$ is the embedding for edge type $i$. Note that any decoder can be used for link prediction in Pinnacle.

3.5. Training details for Pinnacle

Overview.

Pinnacle is trained using the cell type identity of the protein interaction networks and the graph connectivity of the cell type specific protein interaction networks and metagraph. To learn cell type identity, Pinnacle predicts the cell type(s) that the node(s) corresponding to each protein are activated in. For capturing graph connectivity, Pinnacle performs self-supervised link prediction; it predicts whether an edge (and its type) exists between a pair of nodes. For link prediction, a randomly selected subset of edges is masked (or hidden) from the model, and the model must be able to predict that such edges exist (and that the randomly generated false edges do not exist). Practically, this means that the graphs being fed as input into Pinnacle during train, validation, or test do not contain the masked edges.

Data split.

Protein-protein edges are randomly split into train (80%), validation (10%), and test (10%) sets. The metagraph edges are not split into train, validation, and test sets because there are relatively few of them, and they are all critical for injecting cell type and tissue organization to the model. The proteins involved in the train edges are considered in the cell type identification term of the loss function $\left({ℒ}_{\mathit{\text{celltypeid}}}\right)$.

Sampling negative edges.

For link prediction, false (or negative) edges have the label of 0 and are randomly generated (via structured_negative_sampling function in Pytorch Geometric). The ratio of positive to negative edges is 1:1.

Hyperparameter tuning.

We leverage Weights and Biases [108] to select optimal hyperparameters via a random search over the hyperparameter space. The best-performing hyperparameters for Pinnacle are selected by optimizing the ROC and Calinski-Harabasz score [109] on the validation set. The hyperparameter space on which we perform a random search to choose the optimal set of hyperparameters is: the dimension of the nodes’ feature matrix ∈ [1024, 2048], dimension of the output layer ∈ [4, 8, 16, 32], lambda ∈ [0.1, 0.01, 0.001], learning rate for link prediction task ∈ [0.01, 0.001], learning rate for protein’s cell type classification task ∈ [0.1, 0.01, 0.001], number of attention heads ∈ [4, 8], weight decay rate ∈ [0.0001, 0.00001], dropout rate ∈ [0.3, 0.4, 0.5, 0.6, 0.7], and normalization layer ∈ llayernorm, batchnorm, graphnorm, none]. The best hyperparameters are as follows: the dimension of the nodes’ feature matrix = 1024, dimension of the output layer = 16, lambda = 0.1, learning rate for link prediction task = 0.01, learning rate for protein’s cell type classification task = 0.1, number of attention heads = 8, weight decay rate = 0.00001, dropout rate = 0.6, and normalization layers are layernorm and batchnorm. Further, Pinnacle consists of two custom graph attention neural network layers (Section 3) per cell type specific PPI network and metagraph, and is trained for 250 epochs.

Implementation.

We implement PinNACLE using Pytorch (Version 1.12.1) [110] and Pytorch Geometric (Version 2.1.0) [111]. We leverage Weights and Biases [108] for hyperparameter tuning and model training visualization, and we create interactive demos of the model using Gradio [112]. Models are trained on a single NVIDIA Tesla V100-SXM2–16GB GPU.

4. Generating contextualized 3D protein representations

After pre-training Pinnacle, we can leverage the output protein representations for diverse downstream tasks. Here, we demonstrate PinnACLE’s ability to improve the prediction of proteinprotein interactions by injecting context into 3D molecular structures of proteins.

Overview.

Given a protein of interest, we generate both the context-free structure-based representation via MaSIF [3, 17] and a contextualized PPI network-based representation via Pinnacle. We calculate the binding score of a pair of proteins based on either context-free representations or contextualized representations of the proteins. To quantify the added value, if any, provided by contextualizing protein representations with cell type context, we compare the size of the gap between the average binding scores of binding and non-binding proteins in the two approaches.

Dataset.

The proteins being compared are PD-1, PD-L1, B7–1, CTLA-4, RalB, RalBP1, EPO, EPOR, C3, and CFH. The pairs of binding proteins are PD-1/PD-L1 (PDB ID: 4ZQK) and B7–1/CTLA-4 (PDB ID: 1I8L). The non-binding proteins are any of the four proteins paired with any of the remaining six proteins (e.g., PD-1/RalB, PD-1/RalBP1, PD-L1/RalBP1). The PDB IDs for the other six proteins are 2KWI for RalB/RalBP1, 1CN4 for EPO/EPOR, and 3OXU for C3/CFH.

Structure-based protein representation learning.

We directly apply the pretrained model for MaSIF [3, 17] to generate the 3D structure-based protein representations. We use the model pretrained for MaSIF-site task, named all_feat_3l_seed_benchmark. The output of the pretrained model for a given protein is $P\times d$, where $P$ is the number of patches (precomputed by the authors of MaSIF [3,17]) and $d=4$ is the dimension of the pretrained model’s output layer. As proteins vary in size (i.e., the number of patches to cover the surface of the protein), we select a fixed $k$ number of patches that are most likely to be part of the binding site (according to the pretrained MaSIF model). For each protein, we select $k=200$ patches, which is the average number of patches for PD-1, PD-L1, B7–1, and CTLA-4, resulting in a matrix of size 200 × 4. Finally, we take the element-wise median on the 200 × 4 matrix to transform it into a vector of length 200. This vector becomes the structure-based protein representation for a given protein.

Experimental setup.

For each cell type context of a given protein, we concatenate the 3D structure-based protein representation (from MaSIF) with the cell type specific protein representation (from PinNACLE) to generate a contextualized structure-based protein representation. To create the context-free protein representation, we concatenate the structure-based protein representation with an element-wise average of Pinnacle’s protein representations. This is to maintain consistent dimensionality and latent space between context-free and contextualized protein representations. Given a pair of proteins, we calculate a score via cosine similarity (a function provided by sklearn [113]) using the context-free or contextualized protein representations. Lastly, we quantify the gap between the scores of binding and non-binding proteins using context-free or contextualized protein representations to evaluate the added value (if any) of contextual AI.

5. Fine-tuning PinNACLE for target prioritization

After pre-training PinnACLE, we can fine-tune the output protein representations for diverse biomedical downstream tasks. Here, we demonstrate Pinnacle’s ability to enhance the performance of predicting a protein’s therapeutic potential for a specific therapeutic area.

Overview.

For each protein of interest, we feed its Pinnacle-generated embedding into a multi-layer perceptron (MLP). The model outputs a score between 0 and 1, where 1 indicates strong candidacy to target (i.e., by a compound/drug) for treating the therapeutic area and 0 otherwise. Since a protein has multiple representations corresponding to the cell types it is activated in, the MLP model generates a score for each of the protein’s cell type-specific representations (Figure 4a). For example, Protein 1’s representation from Cell type 1 is scored independently of its representation from Cell type 2. The output scores can be examined to identify the most predictive cell types and the strongest candidates for therapeutic targets in any specific cell type.

5.1. Therapeutic targets dataset

We obtain labels for therapeutic targets from the Open Targets Platform [61].

Therapeutic area selection.

To curate target information for a therapeutic area, we examine the drugs indicated for the therapeutic area of interest and its descendants. The two therapeutic areas examined are rheumatoid arthritis (RA) and inflammatory bowel disease. For rheumatoid arthritis, we collected therapeutic data (i.e., targets of drugs indicated for the therapeutic area) from Open-Targets [61] for rheumatoid arthritis (EFO_0000685), ankylosing spondylitis (EFO_0003898), and psoriatic arthritis (EFO_0003778). For inflammatory bowel disease, we collected therapeutic data for ulcerative colitis (EFO_0000729), collagenous colitis (EFO_1001293), colitis (EFO_0003872), proctitis (EFO_0005628), Crohn’s colitis (EFO_0005622), lymphocytic colitis (EFO_1001294), Crohn’s disease (EFO_0000384), microscopic colitis (EFO_1001295), inflammatory bowel disease (EFO_0003767), appendicitis (EFO_0007149), ulcerative proctosigmoiditis (EFO_1001223), and small bowel Crohn’s disease (EFO_0005629).

Positive training examples.

We define positive examples (i.e., where the label $y=1$) as proteins targeted by drugs that have at least completed phase 2 of clinical trials for treating a certain therapeutic area. As such, a protein is a promising candidate if a compound that targets the protein is safe for humans and effective for treating the disease. We retain positive training examples that are activated in at least one cell type specific protein interaction network. The final number of positive training examples for RA and IBD are 152 and 114, respectively.

Negative training examples.

We define negative examples (i.e., where the label $y=0$) as druggable proteins that do not have any known association with the therapeutic area of interest according to OpenTargets. A protein is deemed druggable if it is targeted by at least one existing drug [114]. We extract drugs and their nominal targets from DrugBank [79]. We retain negative training examples that are activated in at least one cell type specific protein interaction network. The final number of negative training examples for RA and IBD are 1,465 and 1,377, respectively.

Data processing workflow.

For a therapeutic area of interest, we identify its descendants. With the list of disease terms for the therapeutic area, we curate its positive and negative training examples. We split the dataset such that about 60%, 20%, and 20% of the proteins are in the train, validation, and test sets, respectively. We additionally apply two criteria to avoid data leakage and ensure that all cell types are represented during training/inference: Proteins are assigned to train (60%), validation (20%), and test (20%) datasets based on their identity; this is to prevent data leakage where cell type specific representations of a single protein are observed in multiple data splits. We also ensure that there are sufficient numbers of train, validation, and test positive samples per cell type; proteins may be reassigned to a different data split so that each cell type is represented during training, validating, and testing stages. With these criteria, the train, validation, and test dataset splits may not necessarily consist of approximately 60%, 20%, and 20% of the total protein representations (Supplementary Table S6).

5.2. Finetuning model details

Model architecture.

Our multi-layer perceptron (MLP) comprises an input feedforward neural network, one hidden feedforward neural network layer, and an output feedforward neural network layer. In between each layer, we have a non-linear activation layer. In addition, we use dropout and normalization layers between the input and hidden layer (see the Implementation section for more information). Our objective function is binary cross-entropy loss.

Hyperparameter tuning.

We leverage Weights and Biases [108] to select optimal hyperparameters via a random search over the hyperparameter space. The best-performing hyperparameters are selected by optimizing the AUPRC on the validation set. The hyperparameter space on which we perform a random search to choose the optimal set of hyperparameters is the dimension of the first hidden layer ∈ [8, 16, 32], dimension of the second hidden layer ∈ [8, 16, 32], learning rate ∈ [0.01, 0.001, 0.0001], weight decay rate ∈ [0.001, 0.0001, 0.00001, 0.000001], dropout rate ∈ [0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8], normalization layer ∈ [layernorm, batchnorm, none], and the ordering of dropout and normalization layer (i.e., normalization before dropout, or vice versa).

Implementation.

We implement the MLP using Pytorch (Version 1.12.1) [110]. In addition, we use Weights and Biases [108] for hyperparameter tuning and model training visualization. Models are trained on a single NVIDIA Tesla M40 GPU.

6. Metrics and statistical analyses

Here, we describe metrics, visualization methods, and statistical tests used in our analysis.

6.1. Visualization of embeddings

We visualize Pinnacle’s embeddings using a uniform manifold approximation and projection for dimension reduction (UMAP) [115] and seaborn. Using the Python package, umap, we transform Pinnacle’s embeddings to two-dimensional vectors via the parameters: n_neighbors = 10, min_dist = 0.9, n_components = 2, and the euclidean distance metric. The plots are created using the seaborn package’s scatterplot function.

6.2. Visualization of cell type embedding similarity

The pairwise similarity of Pinnacle’s cell type embeddings is calculated using cosine similarity (a function provided by sklearn [113]). Then, these similarity scores are visualized using the seaborn package’s clustermap function. For visualization purposes, similarity scores are mapped to colors after being raised to the 20th power.

6.3. Spatial enrichment analysis of Pinnacle’s protein embeddings

To quantify the spatial enrichment for Pinnacle’s protein embedding regions, we apply a systematic approach, SAFE [31], that identifies regions that are over-represented for a feature of interest (Supplementary Figures S3–S4). The required input data for SAFE are networks and annotations of each node. We first construct an unweighted similarity network on Pinnacle protein embeddings: (1) calculate pairwise cosine similarity, (2) apply a similarity threshold on the similarity matrix to generate an adjacency matrix, and (3) extract the largest connected component. The protein nodes are labeled as 1 if they belong to a given cell type context and 0 otherwise. We then apply SAFE to each network using the recommended settings: neighborhoods are defined using the shortpath_weighted_layout metric for node distance and neighborhood radius of 0.15, and $p$-values are computed using the hypergeometric test, adjusted using the Benjamin-Hochberg false discovery rate correction (significance cutoff $\alpha =0.05$).

Due to computation and memory constraints, we sample 50 protein embeddings from a cell type context of interest and 10 protein embeddings from each of the other 155 cell type contexts. We use a threshold of 0.3 in our evaluation of Pinnacle’s protein embedding regions (Figure 2; Supplementary Figures S3). We also evaluate the spatial enrichment analysis on networks constructed from different thresholds to ensure that the enrichment is not sensitive to our network construction method: [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9] (Supplementary Figures S4). We use the Python implementation of SAFE: https://github.com/baryshnikova-lab/safepy.

6.4. Statistical significance of tissue embedding distance

Tissue embedding distance between a given pair of tissue nodes is calculated using cosine distance (a function provided by sklearn [113]). Tissue ontology distance between a given pair of tissue nodes is calculated by taking the sum of the nodes’ shortest path lengths to the lowest common ancestor (functions provided by networkx [116]. We use the two-sample Kolmogorov-Smirnov test (a function provided by scipy [117]) to compare PinNACLE embedding distances against randomly generated vectors (via the randn function in numpy to sample an equal number of vectors from a standard normal distribution). We also use the Spearman correlation (a function provided by scipy [117]) to correlate Pinnacle embedding distance to tissue ontology distance. We additionally generate a null distribution of tissue ontology distance by calculating tissue ontology distance on a shuffled tissue hierarchy (repeated 10 times). Concretely, we shuffle the node identities of the Brenda Tissue Ontology [106] and compute the pairwise tissue ontology distances.

6.5. Statistical significance of binding and non-binding proteins’ score gaps

We perform a one-sided non-parametric permutation test. First, we concatenate the scores for the $N$ binding pairs and $M$ non-binding pairs. Next, for 100,000 iterations, we randomly sample $N$ scores as the new set of binding protein scores and $M$ scores as the new set of non-binding protein scores, calculate the mean ${\mu }_N$ of the $N$ binding protein scores and the mean ${\mu }_M$ of the $M$ non-binding protein scores, calculate the score gap by taking the difference of the means as ${\mu }_N-{\mu }_M$, and keep track of the score gaps that are greater than or equal to the true score gap calculated from the real data. Lastly, we calculate the $p$-value, defined as the fraction of 100,000 iterations in which the permuted score gap is greater than or equal to the true score gap (i.e., one-sided non-parametric permutation test).

6.6. Performance metric for therapeutic target prioritization

For our downstream therapeutic target prioritization task (Methods 5), we use a metric called Average Precision and Recall at K (APR@K) to evaluate model performance. APR@K leverages a combination of Precision@K and Recall@K to measure the ability to rank the most relevant items (in our case, proteins) among the top K predictions. In essence, APR@K calculates Precision@K for each $k\in [1,\dots ,K]$, multiplying each Precision@ $k$ by whether the $k$ th item is relevant, and divides by the total number of relevant items $r$ at $K$:

$$\text{APR}@K=\frac{1}{r}\sum _{k=1}^K\text{Precision}@k\times \text{rel}\left(k\right)$$

where

$$rel\left(k\right)=\left\{\begin{array}{ll}1,& \text{if}\text{item}\text{at}k\text{is}\text{relevant}\\ 0,& \text{otherwise}\end{array}\right.$$

Given the nature of our target prioritization task, some key advantages of using APR@K include robustness to (1) varied numbers of protein targets activated across cell type-specific protein interaction networks and (2) varied sizes of cell type specific protein interaction networks.

Data availability.

All data used in the paper, including the cell type specific protein interaction networks, the metagraph of cell type and tissue relationships, Pinnacle’s contextualized representations, the therapeutic targets of rheumatoid arthritis and inflammatory bowel diseases, and the final and intermediate results of the analyses, are shared via the project website at https://zitniklab.hms.harvard.edu/projects/PINNACLE. Datasets are available via Figshare at https://doi.org/10.6084/m9.figshare.22708126.