Predicting adverse events caused by drug-drug interactions (DDIs) by combining Graph Neural Networks with Large Language Model embeddings. The model uses two distinct types of GNNs:
- GATv2 on molecular structure graphs (what does the drug look like chemically?)
- GraphSAGE on a drug-protein interaction (DPI) bipartite graph (what biological pathways does the drug affect, and which drugs share those pathways?)
Plus BioBERT (a biomedical LLM) embeddings for textual pharmacological knowledge. Given two drugs, the model predicts which of 9,599 adverse event types their combination is likely to cause.
- Problem Overview
- Data
- Architecture
- Results
- Training Details
- Repository Structure
- Usage
- Key Design Decisions
When two drugs are taken together, they can produce adverse events that neither causes alone. A patient on Warfarin (blood thinner) who starts taking Aspirin (anti-inflammatory) faces increased bleeding risk because both drugs interfere with clotting through different mechanisms. Predicting these interactions before they happen is critical for patient safety.
The challenge: with ~800 drugs in our dataset, there are over 300,000 possible pairs. Clinical trials cannot test them all. We train a model to predict interactions from drug properties.
This is a multi-label binary classification problem. For each drug pair, the model outputs a probability for each of 9,599 adverse event types. A single pair can trigger many events simultaneously.
The key challenges:
-
Extreme class imbalance. On average, each drug pair triggers ~209 events out of 9,599 (~2.2% label density). The rarest events appear in fewer than 10 out of 200,000 pairs, giving positive-to-negative ratios exceeding 15,000:1.
-
Three types of drug information matter. Molecular structure tells us what the drug looks like chemically. Text descriptions tell us what the drug does pharmacologically. Protein targets tell us which biological pathways the drug affects. All three are needed.
-
Scale. 200,347 drug pairs, 9,599 event classes, ~17.8M model parameters.
Two primary sources:
DrugBank provides per-drug information:
- Molecular structures (SMILES strings converted to graphs via RDKit)
- Text descriptions of mechanism of action, pharmacodynamics, indications
- Protein interaction profiles: target proteins, metabolising enzymes, membrane transporters (UniProt IDs)
TWOSIDES provides the interaction labels:
- Polypharmacy side effects mined from FDA adverse event reporting systems
- Each record: (drug1, drug2, adverse_event)
- Grouped by pair: ~200,347 unique pairs with multi-hot label vectors over 9,599 events
Raw data is preprocessed once and stored on disk:
| Feature | Format | Source | Dimension |
|---|---|---|---|
| Molecular graphs | PyG Data objects | RDKit from SMILES | variable nodes/edges per molecule |
| Text embeddings | HDF5 float arrays | BioBERT (dmis-lab/biobert-v1.1) | 768 floats per drug |
| Protein vocabulary | JSON mapping | UniProt IDs from DrugBank | 1,841 unique proteins |
| Event mapping | JSON mapping | TWOSIDES event types | 9,599 classes |
Each atom in a molecular graph carries 5 features: atomic number, degree, formal charge, aromaticity, and hydrogen count. Each bond carries 4 features: bond type, conjugation, ring membership, and stereochemistry.
Built at load time from DrugBank protein annotations. This is a bipartite graph:
| Component | Count |
|---|---|
| Drug nodes | 1,729 |
| Protein nodes | 1,841 |
| Edges (bidirectional) | 11,329 |
| Edge type: target | drug directly binds the protein |
| Edge type: enzyme | protein metabolises the drug |
| Edge type: transporter | protein shuttles the drug across membranes |
Data split: 60% train / 20% validation / 20% test at the pair level (seed=42), yielding ~120k / 40k / 40k pairs.
The model processes a pair of drugs through 7 steps. Both drugs go through identical encoders (shared weights), so the model learns general drug representations rather than memorising specific positions.
Here is the full pipeline:
INPUT: Drug A (Aspirin) + Drug B (Warfarin)
GLOBAL (run once per forward pass):
+---------------------------------------------+
| STEP 3: DPI GraphSAGE |
| Drug-Protein Interaction bipartite graph |
| 1,729 drug nodes + 1,841 protein nodes |
| 3 edge types x 3 layers of message passing |
| -> 512-d embedding per drug node |
+---------------------------------------------+
| |
| look up drug A | look up drug B
v v
+--------------------------+ +--------------------------+
| STEP 1: GATv2 Encoder | | STEP 1: GATv2 Encoder | (shared weights)
| Molecular graph | | Molecular graph |
| -> 384-d -> project 512 | | -> 384-d -> project 512 |
+--------------------------+ +--------------------------+
| STEP 2: BioBERT Proj | | STEP 2: BioBERT Proj | (shared weights)
| 768-d -> project to 512 | | 768-d -> project to 512 |
+--------------------------+ +--------------------------+
| |
+--------------------------+ +--------------------------+
| STEP 4: Intra-Drug | | STEP 4: Intra-Drug |
| Fusion | | Fusion |
| [mol 512 + text 512 | | [mol 512 + text 512 |
| + dpi 512] = 1536-d | | + dpi 512] = 1536-d |
| -> MLP -> 512-d | | -> MLP -> 512-d |
+--------------------------+ +--------------------------+
| |
+--------------------------+ +--------------------------+
| STEP 5: Interaction | | STEP 5: Interaction |
| Projection | | Projection |
| Linear+LayerNorm | | Linear+LayerNorm |
| 512-d -> 512-d | | 512-d -> 512-d |
+--------------------------+ +--------------------------+
| |
+----------- STEP 6 ----------+
| |
v v
+------------------------------------------------+
| STEP 6: Symmetric Pair Features |
| |
| d1 + d2 (what they share) = 512-d |
| |d1 - d2| (how they differ) = 512-d |
| d1 * d2 (interaction signal) = 512-d |
| |
| concat = 1536-d |
+------------------------------------------------+
|
v
+------------------------------------------------+
| STEP 7: Residual MLP Classifier |
| |
| Input: 1536-d |
| -> Linear(1536, 1024) + BN + GELU + Drop |
| -> Linear(1024, 1024) + BN + GELU + Drop |
| + residual skip from first layer |
| -> Linear(1024, 9599) -> sigmoid |
| |
| Output: 9,599 probabilities |
+------------------------------------------------+
Let's trace through the model with a concrete example: predicting what happens when a patient takes Aspirin and Warfarin together.
Aspirin (acetylsalicylic acid):
- Structure: a small molecule with a benzene ring, an ester group, and a carboxyl group
- Text: "Aspirin irreversibly inhibits cyclooxygenase (COX-1 and COX-2)..."
- Targets: PTGS1 (COX-1), PTGS2 (COX-2)
- Enzymes: CYP2C9 (metabolised by)
- Transporters: ABCC4
Warfarin (coumadin):
- Structure: a larger molecule with two aromatic rings and a lactone
- Text: "Warfarin inhibits vitamin K epoxide reductase..."
- Targets: VKORC1
- Enzymes: CYP2C9, CYP3A4, CYP1A2
- Transporters: ABCB1, ABCC2
What it does: Converts each drug's molecular structure into a fixed-size vector that captures chemical properties like functional groups, ring systems, and bond topology.
How it works:
-
Build the graph. Each atom becomes a node; each bond becomes an edge. Aspirin has 13 heavy atoms and ~13 bonds. Warfarin has 21 heavy atoms and ~23 bonds.
-
Embed atoms. Each atom's 5 categorical features (atomic number, degree, formal charge, aromaticity, hydrogen count) are looked up in separate learned embedding tables and summed (OGB-style). This produces a 384-dimensional vector per atom.
Carbon atom (atomic_num=6, degree=3, charge=0, aromatic=1, numH=0) -> embed each feature -> sum -> 384-d atom vector -
Message passing (4 rounds of GATv2). In each round, every atom looks at its neighbours using attention. GATv2 applies attention after concatenating source and target features, making it strictly more expressive than standard GAT. Each round expands the receptive field by one hop, so after 4 rounds each atom knows about atoms up to 4 bonds away. Residual connections at every layer prevent information loss.
Round 1: each atom sees 1-hop neighbours Round 2: each atom sees up to 2-hop neighbourhood Round 3: up to 3-hop Round 4: up to 4-hop (captures functional group context) -
Jumping Knowledge (JK). Instead of only using the final round's output, the model concatenates outputs from all 4 rounds (384 x 4 = 1536-d per atom) and projects them back to 384-d. This lets the model choose, for each atom, whether local (round 1) or global (round 4) context is more useful.
-
Graph pooling. All atom vectors are aggregated into a single molecular vector using both sum pooling (preserves molecule size information) and mean pooling (normalises for size). These are added together.
-
Projection. Linear(384, 512) + LayerNorm brings the molecular vector into the shared 512-d fusion space.
Output: one 512-d vector per drug capturing its chemical structure.
Aspirin graph (13 atoms) -> GATv2 x4 -> JK -> pool -> project -> [512-d vector]
Warfarin graph (21 atoms) -> GATv2 x4 -> JK -> pool -> project -> [512-d vector]
What it does: Converts each drug's pharmacological text description into a vector that captures what the drug does, how it works, and what it's used for.
How it works:
-
Pre-compute embeddings (offline). Each drug's DrugBank text description is passed through BioBERT (a BERT model pre-trained on PubMed and PMC biomedical literature). The last hidden layer is mean-pooled across tokens to produce a single 768-d embedding per drug. This happens once and is stored to disk.
-
Project (learned). A small neural network (Linear -> LayerNorm -> GELU -> Dropout) maps the frozen 768-d BioBERT embedding down to 512-d.
BioBERT has already learned biomedical semantics from millions of papers. Words like "inhibits", "agonist", "hepatotoxic", and "CYP2C9" have meaningful representations. The projection layer learns how to translate this knowledge into the format the rest of the model expects.
Output: one 512-d vector per drug capturing pharmacological semantics.
"Aspirin irreversibly inhibits COX-1 and COX-2..."
-> BioBERT -> 768-d (frozen) -> project -> [512-d vector]
"Warfarin inhibits vitamin K epoxide reductase..."
-> BioBERT -> 768-d (frozen) -> project -> [512-d vector]
What it does: This is the key innovation. Instead of encoding each drug's proteins independently (the old DeepSets approach), we build a global bipartite graph connecting ALL drugs to ALL their proteins and run Relational GraphSAGE on it. After multi-hop message passing, each drug's embedding captures not just its own protein interactions, but also which other drugs share those proteins.
The DPI graph looks like this:
Drug Nodes (~1,729) Protein Nodes (~1,841)
Aspirin ──target──────> PTGS1 (COX-1)
Aspirin ──target──────> PTGS2 (COX-2)
Aspirin ──enzyme──────> CYP2C9 <─enzyme── Warfarin
Aspirin ──transporter──> ABCC4
Warfarin ──target──────> VKORC1
Warfarin ──enzyme──────> CYP3A4
Warfarin ──transporter──> ABCB1
Ibuprofen ──target─────> PTGS1 <─target── Aspirin
Ibuprofen ──target─────> PTGS2 <─target── Aspirin
Ibuprofen ──enzyme─────> CYP2C9 <─enzyme── Aspirin, Warfarin
How it works:
-
Initialise nodes. Each drug node and each protein node gets a learnable embedding vector (256-d). These are randomly initialised and learned during training.
-
Relational message passing (3 rounds). Each round has three separate GraphSAGE convolutions, one per edge type (target, enzyme, transporter). This is important because a protein appearing as a metabolic enzyme has different pharmacological significance than the same protein appearing as a direct target.
For each layer, every node:
- Receives its self-contribution through a dedicated linear layer (applied once)
- Receives aggregated messages from target-type neighbours (SAGEConv with mean aggregation)
- Receives aggregated messages from enzyme-type neighbours (separate SAGEConv)
- Receives aggregated messages from transporter-type neighbours (separate SAGEConv)
- All contributions are summed, normalised (LayerNorm), activated (ReLU), and added back via residual connection
-
What each hop captures:
Hop 1: Drug sees its protein neighbours directly Aspirin learns: "I target PTGS1/PTGS2, I'm metabolised by CYP2C9" (This replaces what the old DeepSets encoder did) Hop 2: Drug sees OTHER DRUGS through shared proteins Aspirin -> CYP2C9 -> Warfarin Aspirin -> PTGS1 -> Ibuprofen (This replaces what the old SharedProteinEncoder did, but captures ALL drugs sharing those proteins, not just the current pair being classified) Hop 3: Drug sees proteins of those other drugs Aspirin -> CYP2C9 -> Warfarin -> VKORC1 (Aspirin now "knows about" VKORC1 even though it doesn't directly interact with it -- transitive biological context) -
Output projection. Linear(256, 512) + LayerNorm maps drug node embeddings into the shared 512-d fusion space.
Critical implementation detail: The DPI graph is the same for all drug pairs (it's a global property of the drug universe), so GraphSAGE runs once per forward pass to produce embeddings for all 1,729 drugs. Each batch then just looks up the relevant drug indices. With only ~3,500 nodes, this is extremely fast (~0.5ms).
Output: one 512-d vector per drug, enriched with protein interaction context and inter-drug relationships.
DPI graph (1,729 drugs + 1,841 proteins + 11,329 edges)
-> Relational GraphSAGE x3 layers -> project
-> look up Aspirin node -> [512-d vector]
-> look up Warfarin node -> [512-d vector]
What it does: Combines the three modality vectors for each drug into a single representation. After this step, each drug is described by one vector that encodes its structure, its text description, and its protein interaction context.
How it works:
The three 512-d vectors from Steps 1-3 are concatenated into a 1536-d vector, then compressed:
Linear(1536, 512) -> LayerNorm -> ReLU -> Dropout(0.1)
The MLP learns to weight and combine information across modalities. For a drug where the molecular structure is the most informative signal, the learned weights will emphasise that portion.
Output: one 512-d fused vector per drug.
Aspirin: [mol_512 | text_512 | dpi_512] = 1536-d -> MLP -> [512-d]
Warfarin: [mol_512 | text_512 | dpi_512] = 1536-d -> MLP -> [512-d]
What it does: Refines each drug's vector through a learned projection before computing pair features. This gives the model a dedicated layer to prepare representations specifically for the comparison that follows.
How it works: Linear(512, 512) + LayerNorm applied to each drug independently with shared weights.
Output: one 512-d refined vector per drug (d1, d2).
What it does: Combines the two drug vectors into a single representation of the drug pair. This representation is order-invariant: DDI(Aspirin, Warfarin) produces the exact same features as DDI(Warfarin, Aspirin).
How it works:
Three element-wise operations capture different aspects of the relationship:
| Operation | Formula | What it captures | Example intuition |
|---|---|---|---|
| Sum | d1 + d2 | Properties shared by both drugs | Both affect clotting pathways |
| Absolute difference | |d1 - d2| | Properties that differ between drugs | Different molecular sizes |
| Product | d1 * d2 | Multiplicative interactions | Features strong in both drugs amplify (e.g. CYP2C9 signal) |
All three are symmetric: a + b = b + a, |a - b| = |b - a|, a * b = b * a.
These 3 vectors (512-d each) are concatenated into a 1536-d pair vector.
d1 = Aspirin_512, d2 = Warfarin_512
pair_sum = d1 + d2 -> 512-d
pair_diff = |d1 - d2| -> 512-d
pair_prod = d1 * d2 -> 512-d
pair_features = [sum | diff | prod] = 1536-d
Note: the product d1 * d2 is particularly important here. Because both Aspirin and Warfarin have strong CYP2C9 signals in their DPI GraphSAGE embeddings (both are metabolised by CYP2C9), the corresponding dimensions in the product will be large. This is exactly the "shared protein" signal that was previously computed by hand.
What it does: Takes the 1536-d pair features and predicts probabilities for all 9,599 adverse events.
How it works:
The classifier is a 2-layer MLP with a residual skip connection:
Input: pair_features = 1536-d
Layer 1: Linear(1536, 1024) -> BatchNorm -> GELU -> Dropout(0.1)
|
+--- save as residual
|
Layer 2: Linear(1024, 1024) -> BatchNorm -> GELU -> Dropout(0.1)
|
+--- add residual back (skip connection)
|
Output: Linear(1024, 9599) -> Sigmoid
Result: 9,599 probabilities, each between 0 and 1
At inference time, probabilities above a learned threshold (optimised on validation set, typically ~0.52) are predicted as positive.
Output for Aspirin + Warfarin:
Event 0: "Haemorrhage" -> 0.87 (above threshold -> PREDICTED)
Event 1: "Nausea" -> 0.73 (above threshold -> PREDICTED)
Event 2: "Dizziness" -> 0.61 (above threshold -> PREDICTED)
...
Event 4521: "Hepatic failure" -> 0.02 (below threshold -> not predicted)
...
9,599 total probabilities
Each baseline uses the exact same classifier architecture (Step 7) and training setup. The only difference is which input modalities are used.
| Model | Input Modalities | Sample F1 | Micro F1 | Macro F1 |
|---|---|---|---|---|
| MLP Fingerprint | Morgan fingerprints (2048-bit) | 0.4225 | 0.4996 | 0.0969 |
| GNN Only | Molecular graphs (GATv2) | 0.4035 | 0.4832 | 0.0921 |
| Text Only | BioBERT embeddings | 0.4333 | 0.5298 | 0.1116 |
| Protein Only | Protein profiles (DeepSets) | 0.4219 | 0.5080 | 0.1132 |
| GNN-LLM-DDI (Ours) | All three combined | 0.4465 | 0.5419 | 0.1416 |
Key observations:
- Combining all modalities outperforms every unimodal baseline on every metric.
- Text embeddings alone are surprisingly strong, showing BioBERT encodes substantial pharmacological knowledge.
- Protein features contribute the most to Macro F1 (rare class performance), likely because shared protein targets are a strong biological prior for uncommon interactions.
- The full model's Macro F1 (0.1416) is 25% higher than the best unimodal Macro F1, indicating modalities are complementary.
With 97.8% of all labels being negative, standard binary cross-entropy encourages the model to predict everything as zero. Asymmetric Loss (Ridnik et al., ICCV 2021) solves this:
- gamma_neg = 4: Aggressively down-weights easy negatives. When the model is already 95% confident a label is negative, the loss contribution is reduced by ~16x.
- gamma_pos = 0: No down-weighting of positives. Every positive label contributes fully.
- Probability clipping (clip = 0.05): Negatives with predicted probability below 5% contribute zero loss.
| Parameter | Value |
|---|---|
| Optimiser | AdamW |
| Learning rate | 1e-4 (cosine decay to 1e-6) |
| Weight decay | 1e-5 |
| Batch size | 2048 pairs |
| Gradient clipping | max norm 1.0 |
| Mixed precision | FP16 (PyTorch AMP) |
| Early stopping | patience 25 epochs on validation Sample F1 |
| Max epochs | 150 |
After training, a grid search over thresholds (0.10 to 0.65, step 0.005) on the validation set finds the optimal decision boundary.
- Sample F1 (primary): F1 computed per drug pair, then averaged.
- Micro F1: Pools all predictions, then computes F1. Dominated by common events.
- Macro F1: F1 per class, then averaged. Gives equal weight to rare and common events.
- Hamming Loss: Fraction of individual (pair, event) predictions that are wrong.
ddi/
train.py Training entry point
model.py DDIModel (GATv2 + BioBERT + GraphSAGE fusion)
encoders.py MolecularEncoder, TextProjector, DPIGraphSAGE
dataset.py DDIDataset, DPI graph construction, data loaders
balanced_sampling.py AsymmetricLoss, weighted BCE
inference.py Predict interactions for new drug pairs
requirements.txt Python dependencies
scripts/
prepare_dataset.py Parse DrugBank XML, build protein vocab
precompute.py Generate BioBERT embeddings and molecular graphs
data/
drugbank_clean.csv Parsed drug features
twosides_mapped.csv Interaction labels mapped to DrugBank IDs
event_map.json Adverse event name -> class ID
precomputed/
molecular_graphs.pkl PyG graph objects per drug
text_embeddings.h5 BioBERT 768-d embeddings per drug
protein_vocab.json UniProt ID -> integer vocabulary
experiments/ Self-contained baseline experiments
common/ Shared encoders, data loading, training loop, metrics
mlp_fingerprint/ Morgan fingerprint MLP baseline
gnn_only/ GATv2 molecular graph baseline
text_only/ BioBERT text embedding baseline
protein_only/ Protein DeepSets baseline
multimodal/ Previous multimodal architecture (DeepSets proteins)
paper/ LaTeX source for the research paper
main.tex
references.bib
generate_figures.py
Python 3.10+, PyTorch 2.x, PyTorch Geometric, and dependencies in requirements.txt.
python train.py \
--batch_size 2048 \
--lr 1e-4 \
--weight_decay 1e-5 \
--epochs 150 \
--patience 25 \
--loss asl \
--device cudapython inference.py \
--drug1 "Aspirin" \
--drug2 "Warfarin" \
--top_k 20Why two different GNNs? The molecular structure graph and the drug-protein interaction graph encode fundamentally different information at different scales. Molecular graphs are small (10-50 atoms) with rich edge features (bond types), making GATv2 with its attention mechanism ideal. The DPI graph is large (~3,500 nodes) and heterogeneous (3 edge types), making Relational GraphSAGE with its efficient neighbourhood sampling a better fit.
Why GraphSAGE on the DPI graph instead of DeepSets? The previous architecture (v6) used DeepSets to encode each drug's proteins independently, then manually computed set intersections for shared proteins. This has two limitations: (1) each drug's protein embedding has no context about other drugs, and (2) the shared-protein signal is hand-crafted rather than learned. GraphSAGE solves both: after 2+ hops, each drug's embedding already "knows" about other drugs through shared proteins, and the message passing function is learned end-to-end.
Why relational (per-edge-type) convolutions? A protein appearing as a metabolic enzyme has different pharmacological significance than the same protein appearing as a direct target. CYP2C9 as an enzyme means "this drug is broken down by CYP2C9" (drug metabolism), while PTGS1 as a target means "this drug directly modulates PTGS1" (therapeutic effect). Separate SAGEConv per edge type lets the model learn these distinct roles.
Why GATv2 specifically for molecular graphs? Standard GAT computes attention weights using a static linear combination of source and target features. GATv2 fixes this by applying attention after concatenation, making it strictly more expressive. For molecular graphs, this matters because the importance of a neighbour depends on the specific source-target pair (a carbon bonded to oxygen behaves differently than a carbon bonded to nitrogen).
Why Asymmetric Loss over standard BCE? With 9,599 classes and ~2% label density, a naive BCE model achieves low loss by predicting all zeros. Standard focal loss helps but applies the same gamma to positives and negatives. Asymmetric Loss uses separate gammas (0 for positives, 4 for negatives), recognising that negative samples need much more aggressive down-weighting.
Why sum + |diff| + product for pair representation? These three operations capture complementary aspects of the drug pair: shared properties (sum), complementary properties (absolute difference), and multiplicative interactions (product). The product is particularly important with GraphSAGE: dimensions where both drugs have strong signals from shared proteins will naturally amplify, providing the pair-specific shared-protein signal without explicit set intersection.