Project
Drug–Drug Interaction Extraction with Knowledge Graph Embeddings
A reproducibility study of a CIKM 2022 model that fuses PubMedBERT text features with RotatE knowledge-graph embeddings for 5-class drug–drug interaction classification.
Problem
The DDIExtraction-2013 benchmark asks a model to classify drug–drug interactions described in biomedical text into 5 classes. A CIKM 2022 model combines a language model (PubMedBERT) with knowledge-graph embeddings of the drugs. As a two-person reproducibility study for Big Data for Healthcare (CSE 6250) at Georgia Tech, we set out to rebuild it and check whether its gains hold up.
Approach
Each sentence is encoded with PubMedBERT. In parallel, a drug knowledge graph built from the DrugBank XML database and the DDI corpus is embedded with RotatE. The text features and the knowledge-graph embeddings are fused and passed to a classifier trained with a custom KGE–multi-focal loss, and a loss-function ablation measures what the knowledge-graph embeddings add.
Key results
- Knowledge-graph embeddings raised micro-F1 from 52.0% to 80.7% in the loss-function ablation.
- Drug knowledge graph of 5,719 entities, 264 relation types and 479K triplets.
- DrugBank version drift (5.6× fewer triplets in v5.1) identified as the main source of the reproducibility gap.
My role
- Built the drug knowledge graph from the DrugBank XML database and the DDI corpus (5,719 entities, 264 relation types, 479K triplets).
- Trained the RotatE embeddings and integrated them into the classifier.
- Implemented the custom KGE–multi-focal loss and ran the loss-function ablation.
- Diagnosed DrugBank version drift as the main source of the reproducibility gap, and documented it.
Two-person team. My teammate handled data preprocessing, the text features and data augmentation.
Tools
PyTorch, PubMedBERT