CellAwareGNN: Single-Cell Enhanced Knowledge Graph Foundation Model for Drug Indication Prediction

CellAwareGNN is a novel graph foundation model that integrates single-cell genomics into an expanded biomedical knowledge graph (scPrimeKG) to significantly outperform existing baselines in drug indication prediction, particularly for autoimmune diseases, by capturing cell-type-specific disease mechanisms and enhancing biological interpretability.

The Gist

Imagine you are trying to find the right key to open a specific lock. In the world of medicine, the "lock" is a disease, and the "keys" are drugs. Sometimes, a key designed for one lock (like a headache) might accidentally fit another lock (like an autoimmune disease) perfectly. This is called drug repurposing, and it's a fast way to find new treatments without starting from scratch.

For a long time, scientists have used giant digital maps called Knowledge Graphs to help find these keys. Think of these maps as a massive subway system where every station is a piece of medical information (a drug, a gene, a disease) and the tracks connecting them show how they relate.

The problem with the old maps (like the one used by a previous model called TxGNN) is that they were a bit too "blurry." They knew that a drug affects a gene, but they didn't know where in the body that happened. It was like knowing a train stops at "City," but not knowing if it stops at the "North Station" or the "South Station." For complex diseases like autoimmune disorders, this matters a lot because the trouble often starts in specific neighborhoods of the body (like specific types of immune cells).

Enter CellAwareGNN: The High-Definition Map

The researchers in this paper built a new, super-powered map called CellAwareGNN. Here is how they did it, using some simple analogies:

1. Updating the Map (PrimeKG-U)

First, they took the old map and updated it with the very latest news. Imagine a paper map from 2021; it's missing all the new roads and buildings built in 2025. They refreshed the data so the map reflects the current state of medical science. This alone made the predictions a little better.

2. Adding "Cellular Zoom" (The Secret Sauce)

This is the big innovation. The researchers realized that to really understand autoimmune diseases, you need to know which specific cells are misbehaving.

• The Old Way: The map said, "This gene is linked to Rheumatoid Arthritis."
• The New Way (CellAwareGNN): The map says, "This gene is linked to Rheumatoid Arthritis, but only when it's acting up inside a T-cell or a B-cell."

They did this by integrating data from OneK1K, a massive study that looked at the genetic activity of over a million individual immune cells. It's like upgrading from a blurry satellite photo of a city to a high-definition street view where you can see exactly which house has the lights on.

3. The Result: A Smarter Detective

They trained a computer brain (an AI model) on this new, high-definition map. When they asked it to guess which drugs might treat which diseases, it became much sharper.

The Scorecard:

• The Old Model (TxGNN): Got about 80% of the answers right.
• The Updated Map Model (TxGNN-U): Got about 81.6% right.
• The New Super Model (CellAwareGNN): Got 82.6% right overall, and a massive 86.4% right for autoimmune diseases.

Why Does This Matter? (Real-Life Examples)

The paper shows that this new model didn't just get a few extra points; it found specific, logical connections that the old model missed. Here are two examples:

• Pemphigus (a skin disease): The model suggested Ocrelizumab as a treatment. Why? Because the new map showed that the disease is driven by "B-cells," and Ocrelizumab is a drug known to target B-cells. The old map didn't make that specific cell-level connection as clearly.
• Rheumatoid Arthritis: The model suggested Rosiglitazone. It figured out that this drug turns on a specific switch (PPAR-𝛾) that calms down inflammation in the specific cells causing the arthritis.

The "No-Go" Zone Check

The researchers also made sure they didn't cheat. In the past, AI models sometimes just memorized the answers because the test questions were too easy or the data was mixed up. This team created a strict "exam" where they tested the model on every single disease in the database, not just the easy ones. Even with this harder test, their new model won.

The Bottom Line

Think of CellAwareGNN as upgrading from a general weather forecast ("It will rain in the city") to a hyper-local forecast ("It will rain on the North Side because the clouds are stuck over the North Park").

By adding single-cell data (the "North Park" detail) to the giant medical map, the AI can now see the specific biological mechanisms driving diseases. This leads to:

1. Better predictions: Finding the right drug for the right disease more often.
2. Better explanations: We know why the drug works (e.g., "It targets the specific immune cells causing the problem").
3. New hope: It highlights promising new treatments for difficult diseases like autoimmune disorders that were previously hard to crack.

In short, they gave the AI "eyes" to see the tiny details of our biology, and now it's much better at finding the right keys for our locked-up diseases.

Technical Summary

1. Problem Statement

Current graph foundation models for drug repurposing, such as TxGNN, rely on large-scale biomedical knowledge graphs (KGs) like PrimeKG. While effective, these models face two critical limitations:

1. Lack of Cell-Type Specificity: Existing KGs generally lack fine-grained, cell-type-specific genomic context. They treat gene-disease and drug-target relationships as averaged across tissues, failing to capture mechanisms driven by dysregulated cellular programs (e.g., specific immune cell pathways in autoimmune diseases).
2. Evaluation Bias: Prior evaluations often use randomly selected subsets of drug-disease pairs, leading to test sets dominated by well-studied diseases. This biases performance metrics and fails to assess generalization across the full disease spectrum, particularly for rare or less-characterized conditions.

2. Methodology

A. Knowledge Graph Construction

The authors constructed a hierarchical series of knowledge graphs to address the data limitations:

• PrimeKG-U (Updated Baseline): An updated version of the original PrimeKG, re-curated using the latest versions of 20 biomedical data sources (e.g., DrugBank v6.0+, DisGeNET v8.0, MONDO) as of late 2025. This increased the graph size from ~8.1M edges to ~14.5M edges.
• scPrimeKG (Single-Cell Enhanced KG): The core innovation. The authors integrated high-resolution single-cell data from the OneK1K cohort (1.27 million peripheral blood mononuclear cells from 982 donors).
  • Integration: They mapped 26,597 independent cis-expression quantitative trait loci (cis-eQTLs) to link genes with specific immune cell types (e.g., CD4+ T cells, B cells, Monocytes).
  • Result: The graph expanded to include 22 new cell-type nodes and 54,670 new gene-cell type edges, bringing the total to ~14.52M edges and ~148k nodes. This allows the model to distinguish between global genetic effects and those specific to certain cell populations.

B. Model Architecture: CellAwareGNN

The model utilizes a Graph Foundation Model approach based on a two-layer Heterogeneous Relational Graph Convolutional Network (R-GCN) with a DistMult decoder.

• Encoder: A relation-aware message passing mechanism where node embeddings are updated based on neighbors and relation types.
• Training Strategy (Two-Phase):
  1. Pre-training: The encoder is pre-trained on all relation types in scPrimeKG to learn general biological representations.
  2. Fine-tuning: The model is fine-tuned specifically on drug-disease edges (indications, contraindications, off-label uses) to optimize for therapeutic prediction. The decoder weights are re-initialized during this phase to focus on drug-disease specific relations.
• Loss Function: Binary Cross-Entropy (BCE) loss over positive and negative triples (generated by corrupting source nodes).

C. Evaluation Framework

To address evaluation bias, the authors proposed a Disease-Representation Aware Splitting strategy:

• Stratified Splitting: For every disease in the test set (including rare ones), 30% of its drug-disease edges are held out, ensuring a minimum of one edge per disease.
• Coverage: This guarantees that the model is evaluated on the entire spectrum of diseases in the KG, not just a random subset.
• Metrics: Performance is measured using AUPRC, Recall@k, Average Precision (AP@k), and Mean Reciprocal Rank (MRR@k) across 5 random seeds.

3. Key Contributions

1. scPrimeKG Construction: The creation of the first large-scale, single-cell-enhanced biomedical knowledge graph that integrates cell-type-specific regulatory evidence (cis-eQTLs) with clinical drug-disease data.
2. CellAwareGNN Model: Development of a graph foundation model that leverages this single-cell context to improve drug indication prediction, specifically for complex immune-mediated diseases.
3. Rigorous Evaluation Protocol: Introduction of a disease-stratified evaluation framework that ensures fair assessment across the full disease spectrum, moving beyond random sampling biases.
4. Biological Interpretability: Demonstration that integrating single-cell data not only improves metrics but also aligns model predictions with specific biological mechanisms (e.g., drug effects on specific T-cell or B-cell subsets).

4. Results

Quantitative Performance

• Overall Drug Indication Prediction:
  • CellAwareGNN achieved an AUPRC of 0.826, outperforming the updated baseline TxGNN-U (0.816) and the original TxGNN (0.799).
  • It showed consistent improvements across Recall@5, AP@5, and MRR@5.
• Autoimmune Disease Subset:
  • The improvement was most significant for autoimmune diseases, where cell-type specificity is critical.
  • CellAwareGNN AUPRC: 0.864 vs. TxGNN-U (0.847) and TxGNN (0.815).
  • This represents a 6.0% improvement over the original TxGNN.

Qualitative Analysis & Case Studies

The model successfully prioritized biologically plausible repurposing candidates supported by single-cell evidence:

• Pemphigus: Ranked Ocrelizumab (anti-CD20) and Methotrexate highly. The model correctly linked these drugs to B-cell and T-cell specific gene expression patterns (CD20, DHFR, ATIC) identified in the OneK1K data.
• Rheumatoid Arthritis (RA): Prioritized Rosiglitazone (PPAR-𝛾 activation) and Emapalumab (IFN-𝛾 targeting), aligning with known inflammatory pathways in specific immune cells.
• Multiple Sclerosis (MS): Improved the ranking of Ocrelizumab and Methylprednisolone, correctly capturing their mechanisms regarding B-cell depletion and glucocorticoid receptor expression in pathogenic T/B cells.

Ablation Study

Removing specific edge types (e.g., Drug-Protein, Disease-Phenotype) from CellAwareGNN resulted in only modest performance drops, indicating the model's robustness. However, all variants of CellAwareGNN (even those missing specific edge types) outperformed TxGNN-U, confirming that the inclusion of cell-type edges provides unique, non-redundant value.

5. Significance

• Advancing Precision Medicine: By moving from tissue-averaged data to cell-type-resolved data, the model bridges the gap between genomic risk loci and specific cellular mechanisms, which is crucial for diseases like autoimmunity.
• Interpretability: The model provides not just a prediction, but a biologically grounded rationale (e.g., "Drug X works because it targets Gene Y in Cell Type Z"), enhancing trust in AI-driven drug discovery.
• Scalability: The framework demonstrates that graph foundation models can be effectively scaled with single-cell data to improve performance across thousands of diseases, offering a new paradigm for computational drug repurposing.

In conclusion, CellAwareGNN establishes that integrating single-cell genomics into knowledge graphs significantly enhances the predictive power and biological interpretability of drug repurposing models, particularly for complex, cell-type-dependent diseases.