Link-prediction on Biomedical Knowledge Graphs

<p>Release of code and experimental data from the paper <em>Towards Linking Graph Topology to Model Performance for Biomedical Knowledge Graph Completion&nbsp;</em>(<em>Machine Learning for Life and Material Sciences</em> workshop @ ICML2024) and <a href="https://arxiv.org/abs/2409.04103" rel="nofollow">The Role of Graph Topology in the Performance of Biomedical Knowledge Graph Completion Models</a>.</p> <div> <div>Knowledge Graph Completion has been increasingly adopted as a useful method for several tasks in biomedical research, like drug repurposing or drug-target identification.&nbsp;To that end, a variety of datasets and Knowledge Graph Embedding models has been proposed over the years. However, little is known about the properties that render a dataset useful for a given task and, even though theoretical properties of Knowledge Graph Embedding models are well understood, their practical utility in this field remains controversial. We conduct a comprehensive investigation into the topological properties of publicly available biomedical Knowledge Graphs and establish links to the accuracy observed in real-world applications. By releasing all model predictions we invite the community to build upon our work and continue improving the understanding of these crucial applications.</div> <div>&nbsp;</div> <div>Experiments were conducted on six datasets: five from the biomedical domain (<a href="../records/268568">Hetionet</a>, <a href="https://dataverse.harvard.edu/dataset.xhtml?persistentId=doi:10.7910/DVN/IXA7BM">PrimeKG</a>, <a href="../records/4077338">PharmKG</a>, <a href="../records/5361324">OpenBioLink2020 HQ</a>, <a href="../records/7011027">PharMeBINet</a>) and one trivia KG (<a href="https://aclanthology.org/W15-4007.pdf">FB15k-237</a>). All datasets were randomly split into training, validation and test set (80% / 10% / 10%; in the case of PharMeBINet, 99.3% / 0.35% / 0.35% to mitigate the increased inference cost on the larger dataset).</div> <div>On each dataset, five different KGE models were compared:&nbsp;<a href="https://dl.acm.org/doi/10.5555/2999792.2999923">TransE</a>, <a href="https://arxiv.org/abs/1412.6575">DistMult</a>, <a href="https://arxiv.org/abs/1902.10197">RotatE</a>, <a href="https://arxiv.org/abs/2209.08271">TripleRE</a>, <a href="https://dl.acm.org/doi/10.5555/3504035.3504256">ConvE</a>. Hyperparameters were tuned on the validation split (see final train configurations in <code>train/scripts</code>). We release results for tail predictions on the test split. In particular, each test query&nbsp;<code>(h,r,?)</code> is scored against all entities in the KG and we compute the rank of the score of the correct completion <code>(h,r,t)</code> , after masking out scores of other <code>(h,r,t')</code> triples contained in the graph.</div> <div>Note: the ranks provided are computed as the average between the optimistic and pessimistic ranks of triple scores.</div> <div>&nbsp;</div> <div>Inside <code>experimental_data.zip</code>, the following files are provided.</div> <div> <ul> <li><code>datasets/{dataset}</code>: a folder for each dataset, containing <ul> <li><code>{dataset}_preprocessing.ipynb</code>: a Jupyter notebook for downloading and preprocessing the datasets. In particular, this generates the custom label-&gt;ID mapping for entities and relations, and the numerical tensor of&nbsp;<code>(h_ID,r_ID,t_ID)</code> triples for all edges in the graph, which can be used to compute graph topological metrics (e.g., using <a href="https://github.com/graphcore-research/kg-topology-toolbox">kg-topology-toolbox</a>)&nbsp; and compare them with the edge prediction accuracy.</li> <li><code>test_ranks.csv</code>: csv table with columns <code>["h", "r", "t"]</code> specifying the head, relation, tail IDs of the test triples, and columns <code>["DistMult", "TransE", "RotatE", "TripleRE", "ConvE"]</code> with the rank of the ground-truth tail in the ordered list of predictions made by the five KGE models;</li> <li><code>entity_dict.csv</code>: list of entity labels, ordered by entity ID (as generated in the preprocessing notebook);</li> <li><code>relation_dict.csv</code>: list of relation labels, ordered by relation ID (as generated in the preprocessing notebook).</li> </ul> </li> <li><code>train</code>: code to reproduce training (and validation) of the five KGE models, using the <a href="https://github.com/graphcore-research/bess-kge">BESS-KGE</a> distribution framework. <ul> <li><code>train/scripts</code>: executable scripts, with specifications of the final hyperparameters for all models and datasets.</li> </ul> </li> <li><code>notebooks</code>: Jupyter notebooks for data analysis and generation of all the figures in the paper.</li> </ul> <p>The separate <code>top_100_tail_predictions.zip</code> archive contains, for each of the test queries in the corresponding <code>test_ranks.csv</code> table, the IDs of the top-100 tail predictions made by each of the five KGE models, ordered by decreasing likelihood. The predictions are released in a <code>.npz</code>&nbsp;archive of numpy arrays (one array of shape <code>(n_test_triples, 100)</code> for each of the KGE models).&nbsp;</p> </div> </div>