Graph embeddings (TransE/ComplEx) predict novel drug-disease links for repurposing candidates
A drug repurposing knowledge graph (KG) is a heterogeneous network stitching together disease ontologies, gene/target annotations, and drug-target-indication records into a single queryable structure. Building it well — resolving identifiers, deduplicating relations, and encoding provenance — is the unglamorous foundation on which every downstream embedding and prediction ultimately depends.
The KG is built as a directed, labeled multigraph. Every node carries a type — Disease, Gene/Protein, Drug/Compound, or Phenotype — and every edge carries a relation drawn from a fixed schema: drug–TARGETS–gene, gene–ASSOCIATED_WITH–disease, drug–INDICATED_FOR–disease, disease–HAS_PHENOTYPE–phenotype, gene–INTERACTS_WITH–gene, and roughly forty other predicates.
This schema matters because embedding models learn one vector per relation type in addition to one vector per entity — the relation vector encodes the geometric "translation" that the model must apply to move from a head entity to a plausible tail entity. A cleaner, more consistent schema produces a more learnable geometry.
Entities are resolved to stable identifiers before ingestion: diseases to MONDO/ICD-10, genes to HGNC/Entrez, drugs to ChEMBL/DrugBank IDs. Without this normalization step, the same disease appearing under three different string labels would fragment into three separate nodes and silently destroy the very connectivity the model needs to learn from.
No single database captures the full biomedical picture, so the graph is assembled from complementary sources, each contributing a different edge type:
• DGIdb (Drug-Gene Interaction Database) — curated drug-gene interaction claims aggregated from 30+ underlying sources, used for drug–TARGETS–gene edges • ChEMBL — bioactivity data (IC50, Ki, EC50) linking compounds to protein targets, filtered by potency threshold • Open Targets Platform — evidence-scored gene-disease associations combining genetic, somatic, and pathway data into a single association score • DrugBank — approved and investigational drug indications, used both as positive training triples and as the held-out ground truth against which novel predictions are later compared
Each source is versioned and each edge retains a provenance tag and a confidence weight, because sources disagree — DGIdb and ChEMBL report contradictory targets for roughly 6% of shared drug entries — and the training pipeline needs to know which edges to trust more heavily during negative sampling.
The central statistical fact motivating this entire pipeline: of the ~480,000 biologically plausible (drug, disease) combinations implied by the 10×48 thousand entity space, fewer than 1% are recorded as known indications anywhere in the literature. The graph is extremely sparse by construction — most true relations are simply unobserved rather than false.
This is why knowledge graph embedding models are trained under the open-world assumption (OWA) rather than the closed-world assumption used in classic relational databases: an absent (drug, treats, disease) triple is treated as unknown, not as negatively confirmed. Negative examples for training must instead be synthesized by corrupting true triples (swapping the head or tail entity for a random one), not read directly from the graph.
Fewer than 1 in 10,000 plausible drug-disease pairs is a recorded indication. The entire discipline of link prediction exists because the other 99.99% of cells in that matrix are not "no" — they are "not yet observed."
Knowledge graph embedding compresses every entity and relation into a dense vector in R^d, chosen so that the geometry of the vector space mirrors the relational structure of the graph. TransE, ComplEx, and RotatE are the three workhorse scoring functions; all are trained with frameworks like PyKEEN via stochastic gradient descent over millions of observed and corrupted triples.
TransE, introduced by Bordes et al. (2013), embeds every entity and relation as a vector in the same space and enforces a strikingly simple geometric constraint: for a true triple (h, r, t), the head vector plus the relation vector should land close to the tail vector — h + r ≈ t. The scoring function is the negative distance: f(h,r,t) = −‖h + r − t‖₂ or ‖·‖₁.
During training, gradient descent nudges the embedding of every entity and relation so that observed triples satisfy this translation as closely as possible, while corrupted (false) triples are pushed away. After convergence, the relation vector for "treats" becomes, geometrically, a consistent displacement — walking from any drug's vector along the "treats" direction lands near the vectors of diseases it is known to treat, and near the vectors of diseases with similar underlying biology that it does not yet have a recorded indication for. That geometric neighborhood is exactly what candidate scoring in Stage 3 exploits.
Training minimizes a margin-based ranking loss: for each true triple, one or more corrupted triples are generated by replacing the head or tail with a random entity, and the loss pushes the true triple's score above the corrupted triple's score by at least a margin γ (typically 1.0):
L = Σ max(0, γ + f(corrupted) − f(true))
Uniform random corruption wastes most gradient updates on "easy" negatives that are already scored far apart. RotatE popularized self-adversarial negative sampling, which weights harder (higher-scoring) negatives more heavily in the loss, meaningfully accelerating convergence. PyKEEN — the standard open-source library for this class of model — implements TransE, ComplEx, RotatE, DistMult, and 30+ other architectures behind a common training loop, making head-to-head comparison on the same graph straightforward.
Self-adversarial negative sampling (RotatE, 2019) can cut the epochs needed to reach a target Hits@10 by roughly 40% compared to uniform negative sampling, because gradient signal is concentrated on the negatives the model is actually confusing with true triples.
TransE is fast and interpretable but structurally cannot represent symmetric relations (if a treats b, TransE geometry forces b to not equally translate back to a) or one-to-many relations well. ComplEx and RotatE were designed specifically to fix this by moving into complex-valued embedding spaces, at the cost of roughly double the parameters and reduced interpretability. The table below summarizes the tradeoffs across the four models most commonly benchmarked on biomedical KGs.
| Product | Indication | Trial Design | Key Result |
|---|---|---|---|
| TransE | f(h,r,t) = −‖h + r − t‖ | Simple, fast, highly interpretable geometric translation | Cannot model symmetric or 1-to-N relations |
| ComplEx | Re(⟨h, r, t̄⟩) complex bilinear | Naturally captures asymmetric and antisymmetric relations | Larger parameter count, harder to interpret |
| RotatE | Relations as rotations in ℂ plane | Models symmetry, inversion, and composition patterns | Needs self-adversarial negative sampling to shine |
| DistMult | Bilinear diagonal ⟨h, r, t⟩ | Extremely fast, very few parameters, strong baseline | Structurally cannot represent asymmetric relations |
Once the embedding space is trained, generating predictions requires no new model architecture — only a scan. Every candidate (drug, treats, disease) triple absent from the training graph is scored by the same translation function used during training, and the resulting ranked list is what turns a static graph into a hypothesis-generation engine.
For every drug in the graph, the model scores its embedding translated by the "treats" relation vector against the embeddings of every disease, producing a full ranked list per drug. Naively, this would rank already-known true indications near the top too — which tells us nothing new. The standard filtered evaluation protocol removes all other known-true triples from the candidate list before computing rank, so that a model is never penalized (or credited) for correctly re-surfacing a fact it was already trained on.
The same filtering logic carries over to candidate generation: known (drug, treats, disease) edges present in the training graph are excluded from the repurposing candidate list entirely — the pipeline only ever surfaces pairs the graph has never asserted, by construction.
Two metrics dominate KG completion benchmarks and are reported for every embedding model before it is trusted for repurposing work:
• Mean Reciprocal Rank (MRR) — the average of 1/rank across all held-out test triples; an MRR of 0.31 means the true tail entity lands, on average, around rank 3 among competitors, though the harmonic averaging is dominated by the easy cases • Hits@k — the fraction of held-out triples for which the true entity appears within the top k predictions; Hits@10 of 47% means that for roughly half of withheld true drug-disease pairs, the correct disease is recovered within the model's own top 10 guesses
Both metrics are computed on a held-out test split of known triples never shown during training, giving a proxy for how trustworthy the model's scores on genuinely novel, unlabeled pairs are likely to be.
A random baseline scoring 48,000 candidate diseases per drug would achieve a Hits@10 near 0.02%. A trained TransE/ComplEx ensemble reaching 47% represents roughly a 2,000-fold enrichment over chance — the gap that makes automated hypothesis generation worth pursuing at all.
Scoring 2.4 million candidate pairs sounds expensive, but the translation scoring function reduces to a single batched matrix subtraction and norm computation, which GPUs execute extremely efficiently. All drug embeddings (10 × d), the relation embedding (1 × d), and all disease embeddings (10 × d, broadcast against every drug) are held in memory simultaneously; the entire candidate matrix is scored in one forward pass rather than one triple at a time.
At production scale (tens of thousands of drugs and diseases, d=256), a full re-score after each model update still completes in seconds on a single GPU, which is what makes iterative refinement — retrain, re-score, inspect, adjust the graph, retrain again — a practical weekly cycle rather than a multi-day one.
A raw score is not a recommendation. Candidate ranking layers a top-K cutoff on top of the scored list, then applies mechanistic plausibility filters — does a credible gene-level path connect the drug to the disease? — so that the small number of candidates that reach a human reviewer are both statistically confident and biologically explainable.
Top-K is a deliberately blunt lever: a small K (5–10) yields a short list dominated by the model's highest-confidence predictions, at the cost of missing true positives that scored just below the cutoff; a larger K (25–30) recovers more true positives but dilutes the list with lower-confidence noise that consumes reviewer time.
In practice, K is chosen empirically against a held-out validation set of known repurposing successes (drugs later approved for a second indication after initial approval for a first) — the cutoff that maximizes precision without collapsing recall below a usable floor becomes the operating point, and it is re-validated whenever the underlying graph or embedding model changes.
A geometrically close embedding is a statistical hint, not a mechanism. Before a candidate is presented as a repurposing hypothesis, the pipeline checks whether an explainable path exists through the graph connecting the drug to the disease — typically drug → TARGETS → gene → ASSOCIATED_WITH → disease, possibly through one intermediate gene-gene interaction hop.
This is a form of Path Ranking Algorithm (PRA), layered on top of the embedding score rather than replacing it: embeddings surface candidates that structured path search alone would miss (because embeddings generalize across paths, not just enumerate them), while PRA supplies the human-readable "why" — the actual gene or pathway a reviewer can cross-check against pharmacology literature. Candidates lacking any such path are down-weighted, not discarded outright, since some genuinely novel mechanisms have no prior path in the graph by definition.
78% of top-ranked candidates in this pipeline resolve to at least one explainable gene-level path — the remaining 22% are the most interesting and the riskiest: predictions the embedding space finds compelling but that current pathway annotations cannot yet explain.
Ranking pipelines are validated conceptually against a small set of famous repurposing precedents, each of which followed exactly this drug → gene → disease logic after the fact: sildenafil (developed for angina, its PDE5A-inhibition mechanism made it effective for pulmonary hypertension and erectile dysfunction), thalidomide (a sedative later found to inhibit TNF-alpha and angiogenesis, now used in multiple myeloma and leprosy), and metformin (a first-line diabetes drug whose AMPK-activation mechanism is under active investigation for oncology and longevity indications).
Each of these would, in retrospect, have scored highly in a graph embedding trained on the drug's known target and that target's other disease associations — which is precisely the pattern the candidate ranking stage is built to detect prospectively rather than retrospectively.
The final stage grounds the model in reality: every top-ranked candidate is checked against ClinicalTrials.gov, PubMed literature, and post-market prescribing data to see whether independent clinical evidence already exists. A match is reassuring face validity; the absence of a match does not disprove a candidate — it may simply be genuinely novel.
Each candidate (drug, disease) pair is queried against ClinicalTrials.gov by normalized drug name and disease/condition MeSH term, capturing trials in any phase or status — completed, active, terminated, or withdrawn — since even a terminated trial indicates the hypothesis was considered clinically plausible enough to test. Matches are additionally cross-checked against PubMed abstracts for case reports or observational studies that predate formal trials.
This produces a simple three-way classification for every candidate: (a) independently corroborated — a matching trial or publication exists, strong face validity; (b) partially corroborated — evidence exists for a closely related disease subtype or the same target class rather than the exact pairing; (c) uncorroborated — no independent evidence found, the candidate is a genuinely novel hypothesis awaiting testing.
It is tempting to treat "no matching trial" as a false positive, but that inference is backwards for a discovery pipeline. Precision@K measured against trial registries is a lower bound on the model's true precision, not an unbiased estimate — it can only credit the model for candidates the clinical community has already independently thought to test. A candidate the model ranks highly with no matching trial anywhere is exactly the output the entire system exists to produce: an untested, mechanistically supported hypothesis.
Conversely, a matched trial that failed (terminated for lack of efficacy) is informative in a different way — it suggests the embedding geometry captured a real target relationship that turned out not to translate into clinical benefit, useful negative signal for refining the graph's confidence weighting on that relation type.
The most cited validation of this exact methodology occurred in February 2020, when BenevolentAI used a biomedical knowledge graph to identify baricitinib — a JAK1/JAK2 inhibitor approved for rheumatoid arthritis — as a candidate COVID-19 treatment, reasoning through a target-based path: baricitinib inhibits AAK1, a regulator of viral endocytosis, while its JAK-STAT inhibition also dampens the cytokine storm driving severe respiratory disease.
The prediction was published as a hypothesis before large-scale clinical evidence existed. It was subsequently validated in the NIH-sponsored ACTT-2 trial, and baricitinib received FDA Emergency Use Authorization for hospitalized COVID-19 patients in November 2020 — a rare, well-documented case where a graph-based repurposing prediction preceded, rather than followed, its clinical confirmation.
Baricitinib went from a knowledge-graph-generated hypothesis to FDA Emergency Use Authorization in under nine months — the timeline repurposing pipelines are built to compress, since a molecule already through Phase I-III safety trials for its original indication can skip years of preclinical toxicology.
Validation results feed back into the graph itself rather than terminating the pipeline. Confirmed candidates become new positive training triples for the next training cycle; failed or terminated-trial candidates inform down-weighting of the relation paths that produced them; and the independent validation set of 512 historically known repurposing events is periodically refreshed as new approvals occur, keeping the precision benchmark from going stale.
This closes an active-learning loop: graph → embeddings → predictions → validation → updated graph → re-trained embeddings. Each iteration is intended to shift the uncorroborated-but-plausible candidates that survive multiple rounds toward prospective wet-lab or clinical testing, which is ultimately the only validation that fully resolves a repurposing hypothesis.