lighteternal/psychgnn-masked-edge-imputation-model

PsychGNN Masked-Edge Imputation Model

Model summary

This repository contains a heterogeneous graph neural network trained to recover masked SNP-disorder links within a psychiatric cross-disorder graph.

The model is intended for variant-level research use. It does not predict patient-level diagnosis, prognosis, or treatment response.

Intended task

The training and evaluation task is disorder-conditional masked-edge imputation:

• —start from a graph containing SNP, gene, and disorder nodes
• —remove a subset of observed SNP-disorder links from message passing
• —retain the same SNP's remaining cross-disorder context where available
• —predict whether the hidden SNP-disorder link should exist

This task answers a specific question:

given a psychiatric variant already represented in the graph, can the model recover missing cross-disorder links from the rest of the graph structure?

Data provenance

The checkpoint was trained on:

• —harmonized dataset: lighteternal/pgc-psychiatric-gwas-harmonized
• —graph artifact: lighteternal/psychgnn-psychiatric-graph

The harmonized dataset was derived from public OpenMed / PGC Hugging Face repositories, including:

• —OpenMed/pgc-adhd
• —OpenMed/pgc-anxiety
• —OpenMed/pgc-autism
• —OpenMed/pgc-bipolar
• —OpenMed/pgc-eating-disorders
• —OpenMed/pgc-mdd
• —OpenMed/pgc-ocd-tourette
• —OpenMed/pgc-other
• —OpenMed/pgc-ptsd
• —OpenMed/pgc-schizophrenia
• —OpenMed/pgc-substance-use

Scope

The underlying graph contains 11 modeled disorder groups:

• —ADHD
• —Anxiety
• —Autism
• —Bipolar disorder
• —Borderline personality disorder
• —Eating disorders
• —Major depressive disorder
• —Obsessive-compulsive disorder
• —Post-traumatic stress disorder
• —Schizophrenia
• —Substance use

At inference time the checkpoint can score any of these disorders for variants present in the graph artifact.

Architecture

The model is a custom heterogeneous GraphSAGE-style network over three node types:

• —SNP nodes
• —gene nodes
• —disorder nodes

Message passing uses:

• —SNP self-updates
• —SNP-to-gene aggregation
• —gene self-updates
• —gene-to-SNP aggregation
• —disorder self-updates
• —SNP-to-disorder aggregation

This release uses no disorder-disorder edges.

Decoder heads

The checkpoint contains two task heads:

• —a bilinear link decoder for SNP-disorder link scoring
• —an effect-size regression head for predicting a normalized effect estimate on positive edges

Training configuration

Best hyperparameters:

• —hidden dimension: 192
• —layers: 3
• —dropout: 0.2
• —learning rate: 7.5e-4
• —weight decay: 1e-5
• —negative sampling ratio: 1

Checkpoint metadata:

• —SNP feature dimension: 7
• —gene feature dimension: 4
• —disorder feature dimension: 5
• —effect normalization mean: -0.001849
• —effect normalization std: 0.094672

Graph context

Graph metadata for this release:

• —variants: 18,979
• —genes: 1,205
• —disorders: 11
• —SNP-disorder edges: 22,687
• —SNP-gene edges: 65,634
• —disorder-disorder edges: 0
• —GWS threshold for graph edge construction: 5e-8
• —SNP-gene positional window: 100,000 bp

Evaluation

Primary masked-edge benchmark:

• —test AUROC: 0.9805
• —test AP: 0.9818
• —effect-size Pearson r: 0.9379
• —best validation AUROC: 0.9759

Per-disorder benchmark coverage:

Disorders not included in the masked-edge metric table because the current graph did not provide enough stable maskable positives for this benchmark:

• —Autism
• —Borderline personality disorder
• —Eating disorders
• —Obsessive-compulsive disorder
• —Post-traumatic stress disorder
• —Substance use

Baseline comparison

Inputs and outputs

Inputs

The checkpoint expects:

• —SNP feature matrix
• —gene feature matrix
• —disorder feature matrix
• —SNP-gene edge index
• —SNP-disorder edge index
• —variant and disorder mappings

These are provided by the associated public graph artifact.

Outputs

For a scored (variant, disorder) pair, the model produces:

• —a link score indicating whether the SNP-disorder edge should exist under the masked-edge task
• —an effect-size prediction on the normalized training scale

The primary output of this release is the masked-edge link score.

How to use

Minimal checkpoint loading:

import torch
from huggingface_hub import hf_hub_download

ckpt_path = hf_hub_download(
    "lighteternal/psychgnn-masked-edge-imputation-model",
    "model.pt",
    repo_type="model",
)
checkpoint = torch.load(ckpt_path, map_location="cpu", weights_only=False)

print(checkpoint["hyperparams"])
print(checkpoint["feature_dims"])
print(checkpoint["report"]["masked_split"]["test_auroc"])

To run inference, instantiate a heterogeneous GraphSAGE-style model matching the architecture above, load checkpoint["state_dict"], and score variant-disorder pairs against lighteternal/psychgnn-psychiatric-graph.

Files in this repository

• —model.pt
• —evaluation_report.json

Limitations

• —This model is evaluated on a masked-edge recovery task, not on patient outcomes.
• —The headline benchmark covers only disorders with enough maskable positive edges under the current graph construction.
• —The graph uses a strict GWS edge definition; this reduces benchmark coverage for some disorders.
• —The checkpoint should not be interpreted as evidence of robust generalization to completely unseen disorders.
• —The effect-size head is trained on harmonized summary-statistics edges and does not constitute a causal estimate.

Appropriate use

Reasonable uses:

• —cross-disorder variant prioritization
• —exploratory pleiotropy analysis
• —follow-up prioritization for psychiatric genetics studies
• —downstream research tooling built on the published graph

Inappropriate uses:

• —clinical decision-making
• —diagnosis or screening
• —patient-level risk prediction
• —treatment selection