mims-harvard/PROTON

Model Card

Introduction

Neurological diseases are the leading global cause of disability, yet most lack disease-modifying treatments. We present PROTON, a heterogeneous graph transformer that generates testable hypotheses across molecular, organoid, and clinical systems. To evaluate PROTON, we apply it to Parkinson's disease (PD), bipolar disorder (BD), and Alzheimer's disease (AD). In PD, PROTON linked genetic risk loci to genes essential for dopaminergic neuron survival and predicted pesticides toxic to patient-derived neurons, including the insecticide endosulfan, which ranked within the top 1.29% of predictions. In silico PROTON screens reproduced six genome-wide $\alpha$-synuclein experiments, including a split-ubiquitin yeast two-hybrid system (normalized enrichment score [NES] = 2.30, FDR-adjusted $p < 1 \times 10^{-4}$), an ascorbate peroxidase proximity labeling assay (NES = 2.16, FDR $< 1 \times 10^{-4}$), and a high-depth targeted exome sequencing study in 496 synucleinopathy patients (NES = 2.13, FDR $< 1 \times 10^{-4}$). In BD, PROTON predicted calcitriol as a candidate drug that reversed proteomic alterations observed in cortical organoids derived from BD patients. In AD, we evaluated PROTON predictions in health records from $n$ = 610,524 patients at Mass General Brigham, confirming that five PROTON-predicted drugs were associated with reduced seven-year dementia risk (minimum hazard ratio = 0.63, 95% CI: 0.53–0.75, $p < 1 \times 10^{-7}$). PROTON generated neurological hypotheses that were evaluated across molecular, organoid, and clinical systems, defining a path for AI-driven discovery in neurological disease.

Training Data

PROTON was trained on NeuroKG, a heterogeneous, undirected biomedical knowledge graph contextualized to the human brain. NeuroKG unifies 36 human datasets and ontologies, and integrates single-nucleus RNA-sequencing atlases comprising 3,756,702 cells from the adult human brain. The knowledge graph contains 147,020 nodes across 16 entity types and 7,366,745 edges across 47 relation types. NeuroKG is available via Harvard Dataverse at DOI: 10.7910/DVN/ZDLS3K. For more details, please refer to our project website.

Model Architecture

PROTON is a a 578-million-parameter heterogeneous graph transformer for neurological disease. It was trained on NeuroKG using a self-supervised link prediction objective. Through Bayesian hyperparameter optimization, we selected a model architecture that achieved high link-prediction performance (AUROC = 0.9145; accuracy = 82.23%) on an independent test set.

Model Hyperparameters

• num_feat: 1024
• num_heads: 4
• hidden_dim: 256
• output_dim: 128
• num_layers: 3
• dropout_prob: 0.4546844003628963
• pred_threshold: 0.5

Files Included

• model.ckpt: PyTorch Lightning checkpoint containing model weights.
• decoder.pt: Decoder weights for link prediction (shape [94, 512]).
• edge_types.pt: Ordered list of 47 edge types in NeuroKG to create edge type IDs.
• embeddings.pt: Store of learned embeddings for all 147,020 nodes in NeuroKG (shape [147020, 512]).
• embeddings.csv: Embedding store as a CSV file.
• disease_splits/: Directory containing embeddings of PROTON trained on disease-centric splits.

Files within the disease_splits/ directory follow the naming convention {node_id}_{artifact}, where {node_id} represents the unique identifier for the disease node in NeuroKG. For more details, please refer to our project website.

Usage Instructions

To use PROTON, please clone the GitHub repository and follow the instructions in the README.md. For example, after downloading the model weights and modifying the conf/default.config.yaml file appropriately, you can load the model with the following code:

import pytorch_lightning as pl
from src.config import conf
from src.constants import TORCH_DEVICE
from src.dataloaders import load_graph
from src.models import HGT

pl.seed_everything(conf.seed, workers=True)
kg = load_graph(nodes, edges)
pretrain_model = HGT.load_from_checkpoint(
    checkpoint_path=str(conf.hgt.checkpoint_path),
    kg=kg,
    strict=False,
)
pretrain_model.eval()
pretrain_model.cache_graph(kg, overwrite=False, degree_threshold=conf.neurokg.hparams.degree_threshold)
pretrain_model = pretrain_model.to(TORCH_DEVICE)

License

PROTON is released under the MIT License.

Citation

If you use PROTON, please cite:

@article{noori_graph_2025,
  title={Graph AI generates neurological hypotheses validated in molecular, organoid, and clinical systems},
  author={Noori, Ayush and Polonuer, Joaquin and Meyer, Katharina and Budnik, Bogdan and Morton, Shad and Wang, Xinyuan and Nazeem, Sumaiya and He, Yingnan and Arango, Iñaki and Vittor, Lucas and Woodworth, Matthew and Krolewski, Richard C. and Li, Michelle M. and Liu, Ninning and Kamath, Tushar and Macosko, Evan and Ritter, Dylan and Afroz, Jalwa and Henderson, Alexander B. H. and Studer, Lorenz and Rodriques, Samuel G. and White, Andrew and Dagan, Noa and Clifton, David A. and Church, George M. and Das, Sudeshna and Tam, Jenny M. and Khurana, Vikram and Zitnik, Marinka},
  journal={arXiv preprint},
  note={arXiv:XXXX.XXXXX (placeholder)},
  year={2025}
}

Contact

For any questions or feedback, please open an issue in the GitHub repository or contact Ayush Noori and Marinka Zitnik.