PROTON Model Card

Introduction

Neurological diseases are the leading global cause of disability, yet most lack disease-modifying treatments. To help address this gap, we developed PROTON, a graph AI model that generates hypotheses for neurological disease. PROTON uses a heterogeneous graph transformer contextualized to the adult human brain. PROTON generated predictions across Parkinson's disease (PD), bipolar disorder (BD), and Alzheimer's disease (AD), which we validated using three independent biological systems. In PD, PROTON linked genetic risk loci to genes essential for dopaminergic neuron survival and identified pesticides toxic to patient-derived neurons, including the insecticide Naled ranked within the top 6.75% 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.27, FDR-adjusted p < 1E-4), an ascorbate peroxidase proximity labeling assay (NES = 2.22, FDR < 1E-4), and a high-depth targeted deep sequencing study in 496 synucleinopathy patients (NES = 1.73, FDR < 1.9E-3). In BD, PROTON nominated calcitriol as a candidate drug that reversed proteomic alterations observed in cortical organoids derived from BD patients. In AD, we evaluated PROTON predictions in electronic 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 < 1E-7). PROTON generated and validated mechanistic hypotheses 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.

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.