I Trained Knowledge Graph Embeddings to Find a Cure for My Disease. The Model Found One That Causes It.

I have ulcerative colitis. It's an autoimmune disease where your immune system attacks your colon. I manage it with medication, but the medication options aren't great; immunosuppressants with side effects you wouldn't wish on anyone. Some days are fine, while others aren't so great.

I'm not a biologist. I'm a software engineer. But there's this whole field of computational drug repurposing that promises something remarkable and frankly delusional: take a massive database of known biological relationships, train a machine learning model on the patterns, and it can predict which existing drugs might treat new diseases. The drugs are already approved for other conditions, so you skip years of safety testing. Papers get published. Startups get funded. Impressive-looking ranked lists of "novel drug candidates" get produced.

I decided to build one for UC. I realise that the decision of going through this whole exercise was ambitious but it wasn't to personally discover a treatment, but to find out whether the approach actually works. Whether "45% of known treatments recovered" is a meaningful statement about drug discovery or a meaningless statement about graph retrieval. What I found was more interesting than any drug prediction could have been.

---

The Setup

A knowledge graph stores facts as triples:

(Infliximab,    binds,            TNF-alpha)
(TNF-alpha,     associated_with,  Ulcerative Colitis)
(Azathioprine,  treats,           Ulcerative Colitis)

Millions of these (drugs, genes, proteins, pathways, diseases) linked into one enormous network.

I used DRKG (Drug Repurposing Knowledge Graph), compiled by Amazon and several universities from six curated biomedical databases:

Source	What it captures
DrugBank	Drug-target interactions
Hetionet	Disease-gene relationships
STRING	Protein-protein interactions
+ 3 others	Pathways, side effects, pharmacology

97,000 entities · 5.9M edges · 107 relationship types. Biomedical knowledge as it existed circa 2020.

Every entity becomes a vector in high-dimensional space. Train the model so known-true facts score high and random corruptions score low:

(Infliximab, binds, TNF-alpha)    → TRUE    → push score UP
(Infliximab, binds, Hemoglobin)   → FALSE   → push score DOWN

Repeat millions of times. The vectors shift until geometry encodes biology. Drugs treating similar diseases cluster together. The "treats" relationship becomes a learned geometric transformation. Score every drug against UC and you get a ranked list of candidates.

The experiment that makes this rigorous: remove all 108 known "drug treats ulcerative colitis" edges. Train on the remaining 5.9M edges. Ask: can the model recover those drugs using only indirect paths?

If infliximab (whose "treats UC" edge was removed) still ranks high because the model learned infliximab → TNF-alpha → inflammatory pathways → UC, then it has learned meaningful biological structure.

Three models (TransE, ComplEx, RotatE) trained on a P100 GPU, each with different geometric scoring functions. The metric: Hits@K, what fraction of held-out drugs land in the top K predictions?

---

The First Results Looked Great

Training embeddings on 5.9 million edges is slow and expensive, so I started with a smaller subgraph: 250,000 triples sampled via breadth-first search from UC and the holdout drugs. 128-dimensional embeddings, 30 epochs, basic negative sampling. Standard recipe. On that subgraph, RotatE recovered 45% of known UC treatments in the top 100 out of roughly 5,850 compounds. By any published standard, this is a strong result.

Drug	What it is	Rank (out of 5,850)
Cyclosporine	Immunosuppressant	#1
Infliximab	Anti-TNF biologic	#9
Prednisone	Corticosteroid	#54
Azathioprine	Immunosuppressant	#105
Methotrexate	Anti-inflammatory / chemo agent	#133

Real UC drugs, ranked near the top, without the model ever being told they treat UC. It felt like the system was working.

Then I trained on the full 5.9-million-edge graph, 23 times more data, with 256-dimensional embeddings. The results didn't improve. They collapsed.

Drug	Subgraph rank	Full graph rank
Cyclosporine	#1	#1,683
Infliximab	#9	#5,449
Prednisone	#54	#3,970
Azathioprine	#105	#4,863
Methotrexate	#133	#7,786

Hits@100 dropped from 45.4% to 2.8%. 105 of 108 held-out drugs degraded. "More data = better" is one of the most reliable heuristics in machine learning. When it fails this catastrophically, something structural is wrong. Either the subgraph numbers were inflated, or the full graph training was broken, or both. I spent some time finding out which.

---

The Honest Numbers

Both. Two problems, compounding each other.

First, the subgraph was sampled starting from the drugs I was trying to predict, which meant the graph was built around the answer. I redesigned it from scratch: start from the disease itself and take its biological neighborhood, every edge touching a direct neighbor of UC. No random cap, no circular seeding. The graph composition changed fundamentally.

Property	Original (flawed)	Redesigned
Sampling seed	108 holdout drugs + UC	UC disease nodes only
Drug-drug interaction edges	37.1%	0.7%
Gene-gene edges	27.0%	57.7%
Training triples	250,000 (randomly capped)	409,826 (natural size)

Second, the evaluation protocol itself was broken.

The entity map bug

The entity-to-embedding lookup table was built from training and holdout triples combined. So drugs whose only edge was "treats UC" still got an embedding vector even after that edge was removed. Zero training signal, randomly initialized, never updated. Three such drugs ranked in our top 7 by pure luck.

Fix: Rebuilt the entity map from training triples only. Evaluated only the 91 drugs reachable from UC through indirect paths.

Third, the training was undertrained. I upgraded the negative sampling strategy, embedding dimensions, and epoch count: the specifics are in the table below. Then retrained on both graphs.

	Original runs	Redesigned runs
Embedding dim	128 (subgraph) / 256 (full)	128 (subgraph) / 400 (full)
Epochs	30	50
Negatives per positive	1 (uniform)	256 (NSSA)
Evaluation	All 108 holdout drugs	91 reachable only

Results:

Configuration	Training triples	Hits@100
Original subgraph	250,000	45.4%
Full DRKG (original training)	5,874,125	2.8%
Redesigned subgraph (rebuilt)	409,826	12.1%
Full DRKG (rebuilt)	5,874,125	17.6%

The first two rows used the same protocol: biased subgraph sampling, basic negative sampling, and the entity map bug. The subgraph looked great, the full graph collapsed. The bottom two rows fixed all three problems. The 45% was an artifact of biased sampling and a graph dominated by drug-drug interaction edges that clustered compounds by pharmacological class rather than therapeutic relevance. The full graph's jump from 2.8% to 17.6% came from better training, harder negatives, and honest evaluation combined.

When I examined which drugs the model learned, the pattern was stark:

Drug connectivity	# drugs	Hits@100	Interpretation
High (51+ training edges)	8	62.5%	Model learned these
Medium (11-50 edges)	38	13.2%	Weak signal
Low (1-10 edges)	45	2.2%	Essentially random

Dexamethasone, a corticosteroid used to suppress inflammation, has 120 edges connecting it to dozens of UC-relevant genes. It ranked #11. The model had rich context to work with: a well-studied drug connected through multiple pathways to UC's gene network. Theobromine, the stimulant in chocolate with a single edge in the graph, ranked #5,611. The model cannot learn a useful embedding from one data point.

If the model is really just ranking by connectivity, that should be measurable. So I checked how well a drug's number of training edges predicts its rank:

Not by learned biological insight. By connectivity.

This might seem like a model failure. Maybe better hyperparameters would learn deeper patterns. But connectivity isn't a confound to be removed. It IS the signal.

---

The Core Finding: Epistemological Circularity

The model ranks by connectivity because the graph already encodes the answer.

Consider the held-out drugs, the ones whose "treats UC" edges I removed.

The graph already separates known UC drugs from everything else through indirect structure alone.

Why? Because drugs that treat UC have been extensively studied in the context of UC. Every study produces edges: drug targets this gene, this gene is associated with UC, this pathway is relevant to inflammation. The more a drug has been researched for UC, the more indirect connections it accumulates in the graph. High connectivity isn't an independent feature; it's a direct consequence of being a known UC treatment.

Here's the causal chain, using infliximab (an anti-TNF biologic I actually take for UC) as a concrete example:

This isn't data leakage in the ML sense. I correctly removed the target edges. The holdout is technically clean. It's a deeper problem: epistemological circularity. Epistemology is how we know what we know, and the circularity here is that the evidence we're using to "predict" the answer was only generated because someone already knew the answer. The features used for prediction were placed in the graph by the same human research process that established the answer. The Drug → Gene and Gene → Disease edges are consequences of the Drug → Disease relationship I'm trying to "predict."

Every structure-aware method picks up this signal because it's baked into the graph topology. To confirm, I ran Personalized PageRank, a textbook graph algorithm that reads graph proximity with zero learning. Two minutes on my Mac.

Method	Training time	Hits@100
RotatE (400d, NSSA, P100 GPU)	~8 hours	17.6%
Personalized PageRank	~2 minutes	34.1%

A method with no training, no embeddings, no GPU doubles the best embedding model. It just reads the structure that's already there.

Holdout evaluation on curated knowledge graphs measures retrieval efficiency, not discovery potential. High Hits@K tells you your method can read what's already in the graph. It does not tell you it can find anything new.

---

This Isn't Just a Bio Problem

The same circularity appears wherever knowledge graphs meet prediction tasks, including in deployed production systems.

The most celebrated success story in KG drug repurposing illustrates it perfectly. In February 2020, days after the WHO declared COVID-19 a public health emergency, BenevolentAI published a prediction in The Lancet: baricitinib, a JAK inhibitor approved for rheumatoid arthritis, could treat COVID-19 by blocking viral entry through AAK1 inhibition. They identified it in 48 hours using their biomedical knowledge graph. The drug received FDA emergency authorization in November 2020 and full approval in 2022. It works. Patients are alive because of this prediction.

But every edge in their reasoning chain was already in published literature:

Edge	Published
Baricitinib binds AAK1	2016, Sorrell et al., Structure
AAK1 role in viral endocytosis	2017, Bekerman et al., J. Clinical Investigation
Coronaviruses use endocytic entry	Established virology textbook knowledge

Derek Lowe, the medicinal chemist who writes for Science/AAAS, called it "more a testimony to good literature searching and curation" than a triumph of artificial intelligence. Timothy Cernak at Michigan said most medicinal chemists could have identified the molecule using traditional database searches. The KG's genuine contribution was speed (systematic retrieval under extreme time pressure), not discovery of a connection hidden from human experts.

The pattern generalizes. If the same human research process that created your labels also shaped your features, your benchmark measures memory, not prediction. This applies to any task built on curated knowledge bases, any holdout evaluation where the test items are well-studied, any link prediction on a graph assembled by domain experts. It does not apply where labels are naturally generated (user clicks, sensor readings) or where features are genuinely independent of the labeling process.

---

What Would Real Discovery Look Like?

Three approaches would address the circularity:

Temporal holdout. Train on knowledge available before year X. Predict treatments approved after year X. If the model identifies drugs from post-X research using only pre-X graph structure, it's genuinely discovering something that wasn't already encoded. DRKG doesn't provide edge timestamps, but versioned knowledge graphs could be constructed from dated database releases and publication records. This is the cleanest experimental design because it mirrors what real drug discovery requires: predicting the future from the past.

Novel-target evaluation. Evaluate only on drugs with zero existing graph connection to the target disease. No shared genes, no pathway overlap, no indirect paths. If a model predicts such a drug and it turns out to work, that's prediction in the truest sense. Current benchmarks never test this because every structure-aware method scores disconnected drugs randomly. But that's precisely the frontier where computational methods could add value that humans can't.

External prospective validation. Score novel candidates and check against ClinicalTrials.gov for active trials the model wasn't told about.

What happened when I tried this

I ran my model's top 100 novel predictions through a clinical plausibility review. 59 of them were oncology drugs. A sample of what ranked in the top 50:

• Irinotecan (rank 23): a chemotherapy agent that appears in IBD clinical guidelines as a drug to watch for causing intestinal damage, not one for treating it.
• Hydroxyurea (rank 46): documented in case literature as a cause of intestinal injury that resolves when the drug is stopped.
• Indomethacin (rank 16): an NSAID. IBD patients are told to avoid this drug class because it tends to worsen the disease.

The model has no mechanism to distinguish "this drug is connected to UC inflammation research" from "this drug would help someone with UC." Those look identical in the graph.

---

When I ran my model's top novel predictions through a clinical plausibility review, 51 of the first 100 were flatly implausible. Not "probably not useful": actively contraindicated. Drugs that damage the gut. Drugs that worsen the disease. Mitomycin C at rank 14, a compound used in laboratory experiments to induce the condition I was trying to treat. Zero of 100 received a high plausibility rating.

The honest answer to "does this work" is: it works at something. Retrieving drugs already connected to UC in the graph: yes. Finding treatments that aren't already encoded in the data: no demonstrated evidence. Those are different tasks. For the first task, you don't need a trained embedding. PageRank does it better in two minutes. For the second task, you need a different benchmark entirely.

What that benchmark looks like is the more interesting question. I don't have the answer yet, but for the first time, I have the right question. That's probably the right outcome for a software engineer with a disease who decided to go poking around in biomedical ML.

sushrut1058/knowledge-graph-circularity