🧬 Heterogeneous Graph Neural Networks for Drug Discovery (CX4803-ML-in-Computational-Biology)

This repository contains the final project for CX4803 ML in Computational Biology.
We aim to predict drug-disease associations using a heterogeneous graph built from Hetionet.

---

📌 Task Summary

• Objective: Predict missing links between Compounds and Diseases.

• Graph Construction:

  • Node types:
    • Compound
    • Gene
    • Disease
  • Meta-edges:
    • CrC: Compound resembles Compound
    • CpD: Compound palliates Disease
    • CtD: Compound treats Disease
    • CbG: Compound binds Gene
    • CdG: Compound downregulates Gene
    • CuG: Compound upregulates Gene
    • DrD: Drug resembles Drug
    • DaG: Drug associates Gene
    • DdG: Drug downregulates Gene
    • DuG: Drug upregulates Gene
    • GcG: Gene covaries Gene
    • GiG: Gene interacts Gene
    • GrG: Gene regulates Gene

---

⚙️ Jupyter Cluster Environment Configuration

Setting	Value
Python Environment	Anaconda3 2023.03
Jupyter Interface	Jupyter Notebook
Node Type	NVIDIA GPU (first avail)
Number of Nodes	2
Cores per Node	2
GPUs per Node	1
Memory per Core	8 GB
Total Memory per Node	16 GB
Total GPUs	2
Job Duration	7 hours

🧪 Notes

• PyTorch Geometric (PyG), Sentence Transformers, and other libraries were used.
• GPU memory was critical for Node2Vec and GNN model training.

---

🔬 Methodology

Node Feature Initialization

• Node2Vec embeddings trained using random walks.
• Sentence Transformer embeddings (all-MiniLM-L6-v2) for node name text embeddings.
• Nodes missing embeddings are assigned small random vectors.

Model Variants

• Node2Vec + GCN:
  Basic homogeneous GCN using Node2Vec embeddings.

• Sentence Transformer + HeteroConv:
  Heterogeneous GNN using semantic embeddings.

• Node2Vec + RGCNConv:
  Relational GCN specialized for multi-relational data.

• Node2Vec + HGT Link Predictor:
  Heterogeneous Graph Transformer (HGT) model.

---

📈 Evaluation Metrics

Model	AUC	F1-score
Node2Vec + GCN	0.9883	0.9578
Sentence Transformer + HeteroConv	0.8929	0.7794
Node2Vec + RGCNConv	0.9590	0.8951
Node2Vec + HGT Link Predictor	0.7709	0.8547

✅ Final Model Selected:
Sentence Transformer + HeteroConv for better generalization, avoiding overfitting to top-n trivial compounds.

---

🔍 Observations & Challenges

• High performance models sometimes recommended the same few compounds across different diseases (top-n problem).
• Node representation was critical; thus, richer semantic embeddings were preferred.
• HGT had persistent errors due to PyG internal handling of metadata, requiring model architecture revision.

---

🚀 Future Work

To further enhance model performance:

• Advanced embeddings such as:
  • SMILES embeddings (molecular graph)
  • Gene ontology embeddings
  • Protein sequence BERT embeddings
• Large-scale heterogeneous pretraining.
• Multi-task learning for better generalization across drug discovery tasks.

---

🛠 Reproducibility

Please ensure the following:

• Python ≥ 3.8
• PyTorch ≥ 1.12
• torch-geometric ≥ 2.2
• sentence-transformers
• scikit-learn

Datasets and pretrained embeddings should be available in the working directory.

---

👨‍💻 Team

Contributors:

• Kyungbeom Kim
• Eunsu Hwang
• Lasya Pasumarthy
• Malar Paavai Muthukumaran

---

🔗 Project Code

The full implementation is available on our private GitHub repository:
GitHub Link (GT internal)