XplainMD Part 2: Finding the Missing Links with Machine Learning

In **Part 1 of the XplainMD series*, we zoomed out to explore the architecture of biomedical knowledge — unpacking the rich topology of PrimeKG through centrality analyses, causal subgraphs, and community detection. By mapping how diseases, drugs, proteins, and phenotypes interconnect in a vast biomedical graph, the foundation has been laid for understanding not just what exists — but what might be missing*.

Now, in Part 2, the gears will shift from exploration to prediction.

Graphs open the door to a wide range of prediction tasks — from node classification (predicting properties of nodes) to link prediction (inferring missing or potential connections). Since our focus is on understanding hidden relationships in biomedical data, this series dives into link prediction.

So, what is link prediction? It’s the task of asking: “Given what we know about this graph’s structure, can we infer meaningful relationships that aren’t explicitly present?”

This is where representation learning steps in.

By applying Node2Vec, the rich topology of the biomedical graph is transformed into dense vector embeddings that capture both semantic proximity and structural context. These embeddings serve as the foundation for downstream tasks — in this case, predicting missing or unknown edges between entities like drugs, diseases, and phenotypes.

These embeddings further become the input to downstream machine learning models like Logistic Regression and XGBoost, enabling us to tackle the powerful task of link prediction — estimating whether a biologically plausible, yet currently unobserved, connection exists between entities such as a disease and a phenotype.

This is where XplainMD begins to evolve from graph understanding into graph reasoning.

Data Pre-processing : Structuring the Graph

Before training any Machine Learning or Deep Learning models, it is essential to ensure the input data is well-formatted, clean, and consistent. This preprocessing step converts the raw PrimeKG CSV into a form that can be used to build a graph structure for machine learning tasks and deep learning as well.

1. Data Loading

In this project, a filtered subset of PrimeKG was loaded into a pandas DataFrame to focus on the most clinically relevant biomedical relationships. Specifically, only the following relation types were extracted:

selected_relations = [
    "protein_protein",
    "disease_phenotype_positive",
    "bioprocess_protein",
    "disease_protein",
    "drug_effect",
    "pathway_protein",
    "disease_disease",
    "contraindication",
    "drug_protein",
    "indication"
]

What Does Each Row Represent?

Each row in the DataFrame corresponds to a single edge in the biomedical knowledge graph — that is, a meaningful connection between two biomedical entities.

The relevant columns include:

• **x_name, y_name**: The actual names of the two nodes connected by the relation (e.g., “Alzheimer’s disease”, “APP”).
• **x_type, y_type**: The entity types for each node — such as disease, protein, drug, phenotype, etc.
• **relation**: The type of biomedical relationship between the nodes — e.g., disease_protein or drug_effect.
• **x_source, y_source: These fields do not indicate directionality of the edge — instead, they refer to the original source database** (like NCBI or DrugBank) from which the node was extracted.

️ Direction ≠ Semantics

While the table structure follows a source → target format (x_name to y_name), this does not mean the graph is directed. According to the official PrimeKG paper, the graph is treated as undirected during analysis and modelling. This means that relationships are bidirectional, even though they are stored in a structured row format.

Why This Matters for Graph Construction

Understanding the true semantics of these edges is critical. When building the graph later in PyTorch Geometric (or any GNN library):

• Treat the edges as undirected for most graph algorithms and embeddings like Node2Vec.
• Ensure the edge type (relation) and node types (x/y_type) are preserved in a mapping — enabling construction of typed heterogeneous graphs.

2. Text Normalisation

Biomedical datasets often contain inconsistent casing, hidden unicode characters, or stray spaces. To ensure uniformity across node names, every node label is lowercased, stripped of whitespace, and normalised using unicodedata.

def clean_text(text):
    return unicodedata.normalize("NFKD", str(text)).strip().lower()

df["x_name"] = df["x_name"].apply(clean_text)
df["y_name"] = df["y_name"].apply(clean_text)

3.Type Mapping

Node types in PrimeKG can vary in format — e.g., “gene/protein”, “chemical/drug”, or redundant variants like “bioprocess”. These are mapped to canonical categories to simplify modelling and ensure consistency.

node_type_mapping = {
    "gene/protein": "protein",
    "chemical/drug": "drug",
    "drug": "drug",
    "disease": "disease",
    ...
}

4. Extracting Node Names, Types, and Normalized Relations

Before applying any graph machine learning technique, we need to structure the biomedical data in a way that respects its semantic complexity. In PrimeKG, each row represents a biologically meaningful link — such as a drug treating a disease, a gene associated with a phenotype, or a protein interacting with another protein.

But models like R-GCN or Node2Vec don’t just want a list of edges — they need a clear map of what each node is, what role it plays, and how it’s connected.

Step 1: Assign Global Node IDs

The first step is to collect all unique node names across both columns (x_name and y_name) and assign each one a global integer ID. This gives us a consistent reference for each entity throughout the graph.

all_nodes = pd.concat([df["x_name"], df["y_name"]]).dropna().unique()
node_maps = {name: i for i, name in enumerate(sorted(all_nodes))}
print(f"[INFO] Total unique nodes: {len(node_maps):,}")

Analogy: Think of this like assigning a library index number to every book — whether it’s in the “Medicine” section or “Biochemistry,” a unique ID is aasigned to keep everything organised.

Step 2: Normalise the Relation Map

In a heterogeneous biomedical graph, relationships connect different types of nodes:

• A disease–protein interaction is different from a drug–effect link
• Some relationships are directional, others symmetrical

To build a flexible yet structured graph, the relation types were normalised by sorting their source and target node types alphabetically. This ensures consistency and avoids duplication (e.g., drug→disease is treated the same as disease→drug if the model doesn't care about directionality).

relation_map = {}

for rel in df["relation"].unique():
    subset = df[df["relation"] == rel]
    if subset.empty:
        continue

    type_pairs = set(zip(subset["x_type"], subset["y_type"]))

    for x_type, y_type in type_pairs:
        if x_type in node_type_mapping.values() and y_type in node_type_mapping.values():
            normalized_pair = tuple(sorted([x_type, y_type]))
            relation_map[rel] = normalized_pair

print(f"[INFO] Total unique normalized relations: {len(relation_map):,}")

Analogy: This is like grouping roads on a map based on which areas they connect, regardless of direction — a road from “Hospital to Lab” is still the same route as “Lab to Hospital.”

Step 3: Build the Node Metadata Table

To keep track of each node’s type and global ID, a unified node_df was created that holds every unique node, its type (e.g., "gene", "disease"), and the global ID that was previously assigned.

node_df = pd.concat([
    df[["x_name", "x_type"]].rename(columns={"x_name": "node_name", "x_type": "node_type"}),
    df[["y_name", "y_type"]].rename(columns={"y_name": "node_name", "y_type": "node_type"})
]).dropna().drop_duplicates().reset_index(drop=True)

node_df["global_id"] = node_df["node_name"].map(node_maps)

Analogy: This is like creating a clean catalog where every book (node) has its genre (type) and index number (ID) — critical for graph construction.

Constructing the Graph with Global Node Mapping

Step 1: Global Node Indexing

Earlier in the pipeline, each unique biomedical entity (e.g., gene, disease, phenotype) was assigned a global integer ID using:

node_maps = {name: i for i, name in enumerate(sorted(all_nodes))}

This ensures that every node — regardless of its type — is mapped to a unique identifier, creating a flat, consistent index space that simplifies downstream processing.

Step 2: Adding Edges to the Graph

With node IDs in hand, looping through each relationship type (from the normalised relation_map) and add the corresponding edges:

G = nx.Graph()

for rel in relation_map:
    rel_df = df[df['relation'] == rel]

    src_indices = rel_df['x_name'].map(node_maps).fillna(-1).astype(int)
    dst_indices = rel_df['y_name'].map(node_maps).fillna(-1).astype(int)

    valid_edges = [(s, d) for s, d in zip(src_indices, dst_indices) if s != -1 and d != -1]
    G.add_edges_from(valid_edges)

print("\n[INFO] Graph constructed successfully with global node map.")

Here’s what this does:

• For each relation, it selects the relevant rows from the dataset.
• It converts the source and target node names into global IDs using the node_maps dictionary.
• Any missing or invalid mappings are filtered out (using -1 as a sentinel).
• All valid edges are added to the graph.

Learning Node Representations with Node2Vec

Using Node2Vec, a model is trained to convert nodes into dense, continuous embeddings — capturing semantic and topological relationships between the entities. It is an unsupervised learning algorithm that learns low-dimensional embeddings for nodes by simulating random walks on the graph.

At its core, Node2Vec learns by simulating random walks across the graph — just like how Word2Vec learns word embeddings from natural language. It treats each node like a “word” and each walk like a “sentence.” By walking through the graph in flexible, biased ways (some walks stay local, others explore far), it captures both structural roles (e.g., hubs, bridges) and semantic proximity (e.g., diseases linked by shared phenotypes or pathways).

Image Generated using ChatGPT-4o

The result? A high-dimensional representation where nodes with similar roles or connections are embedded close together — even if they aren’t directly connected.

This makes it possible for ML models to detect missing links, suggest biological analogies, and uncover latent similarities — all from the geometry of the graph.

Converting NetworkX Graph to PyTorch Geometric Format

Once the undirected graph G is constructed using NetworkX, the next step is to embed its nodes into a continuous vector space using the Node2Vec algorithm. These embeddings are designed to capture the structural roles and semantic context of nodes based on their local and global neighbourhoods.

To do this effectively in PyTorch Geometric (PyG), the graph needs to be transformed into a format that PyG understands. This is done using:

pyg_graph = from_networkx(G)

This line converts the G object (a standard NetworkX graph) into a torch_geometric.data.Data object. The resulting pyg_graph includes PyG-friendly attributes like edge_index, a 2D tensor that defines the graph's connectivity in terms of source and target node indices.

This format allows PyG models like Node2Vec, GCN, or RGCN to efficiently process the graph, optimise over its structure, and learn expressive embeddings.

The edge_index serves as the backbone of all graph-based computations in PyG, enabling operations like random walks, message passing, and convolution to be implemented seamlessly.

Sending to Device (CPU/GPU)

pyg_graph.edge_index = pyg_graph.edge_index.to(device)

To leverage GPU acceleration (if available), the graph’s edge list is moved to the appropriate device.

Initializing Node2Vec

Node2Vec doesn’t just look at who’s connected to whom — it walks the graph like a tourist, exploring local and global neighbourhoods to uncover hidden structural patterns.

It combines two clever ideas:

• Random Walks: For each node, Node2Vec simulates multiple random walks — like sending out a curious explorer to roam the neighbourhood. These walks create sequences of nodes, kind of like sentences in a language.
• Skip-Gram Model: Inspired by Word2Vec, the skip-gram model learns to predict a node’s neighbours (context) from these sequences. It treats nodes like words and walk sequences like sentences, capturing how often and in what order nodes appear together.

node2vec = Node2Vec(
    pyg_graph.edge_index,
    embedding_dim=128,
    walk_length=10,
    context_size=5,
    walks_per_node=20,
    num_negative_samples=1
).to(device)

The configuration used in this project is carefully tuned to balance exploration and efficiency during the Node2Vec training process:

• embedding_dim=128: Each biomedical entity—be it a disease, gene, or drug—is represented by a 128-dimensional vector, capturing its structural and semantic context in the graph.
• walk_length=10: Each simulated random walk explores 10 steps from a starting node, allowing it to traverse across nearby biological relationships (e.g., a disease → protein → drug → pathway).
• context_size=5: For every node, only its 5 closest neighbours in a walk are treated as context. This is akin to saying: “Which other genes are typically discussed near BRCA1 in biomedical pathways?”
• walks_per_node=20: The model simulates 20 random walks per node, giving it enough exposure to both local and global graph structure. For instance, breast cancer might co-occur with immune genes, metabolic pathways, or co-morbid phenotypes in different walks.
• num_negative_samples=1: For every positive pair (e.g., Breast Cancer ↔ TP53, which co-occur in a walk), one negative pair is generated by randomly sampling unrelated nodes (e.g., Breast Cancer ↔ Toe curvature). This teaches the model to pull meaningful pairs closer while pushing irrelevant pairs apart.

This setup enables Node2Vec to learn embeddings that reflect biomedical semantics, even though the training is entirely unsupervised. The internal buffers rowptr and col, essential for sampling operations, are also moved to the correct device (GPU or CPU) to ensure efficient execution.

Edge Sampling for Training

train_edges, val_edges = train_test_split(...)

Instead of training on the entire graph at once, the training loop samples batches of edges and trains on mini-batches. A 90–10 split is used for training and validation.

Training Loop with Early Stopping

The model is trained using Adam optimizer. For each epoch:

1. A batch of nodes is sampled based on the edge list.
2. Positive random walks and negative random walks are generated.
3. The model computes the loss, backpropagation, and updates the weights.
4. Validation loss is computed and monitored.

if val_loss.item() < best_loss:
    ...
    torch.save(...)

If the validation loss improves, the model is saved. Otherwise, a counter is incremented. Training stops early if no improvement is seen for 200 consecutive epochs.

Visualizing Node2Vec Embeddings with t-SNE

The Node2Vec model was trained on an undirected biomedical graph to learn vector representations for each node based on its local and global connectivity. After training, these embeddings were projected into two dimensions using t-SNE, a non-linear dimensionality reduction technique that preserves local structure and neighbourhood relationships.

Before Node2Vec training

Embedding after Node2Vec training

Visualising the Graph Embeddings with t-SNE

The scatter plot above presents a t-SNE projection of node embeddings from the PrimeKG graph — with each point representing a node (e.g., disease, phenotype, drug, protein, pathway, or biological process). The axes (Component 1 and Component 2) are abstract latent dimensions created by t-SNE and don’t correspond to specific biomedical properties. Instead, they are used to help us visualise structural similarity in a 2D space.

Nodes that appear closer together in this plot were likely embedded with similar structural contexts — meaning they share common neighbors, appear in similar paths, or participate in similar types of relationships within the graph.

This version of the plot is subsampled to improve visibility while maintaining the distributional structure of the full graph.

Key Observations

• No single dominant cluster is present, but we observe dense regions of overlap where nodes of different types co-locate — reflecting the highly interconnected nature of biomedical entities in PrimeKG.
• Drugs and proteins appear more uniformly scattered, consistent with their broad connectivity across multiple biomedical contexts (e.g., a drug linking to diseases, pathways, and targets).
• Phenotypes and diseases still form partial clusters, often overlapping — which makes sense biologically, as phenotypes are often clinical manifestations of diseases, and their embeddings are shaped by similar neighbourhood structures.
• Pathways and biological processes are sparsely spread out, possibly due to lower edge density or fewer random walk interactions — indicating they may function more as semantic anchors in the graph than highly connected hubs.

What This Means

Despite the noise introduced by subsampling and the non-deterministic nature of t-SNE, there are clear semantic signals emerging from the structure:

• Nodes of similar types often drift toward local neighborhoods, showing that the graph structure preserves contextual semantics.
• The fact that different biomedical entities aren’t isolated but rather entangled in shared regions is reflective of real-world biology, where interdependencies are the norm.

These patterns validate that the graph construction and embedding pipeline is working — capturing not just node proximity, but meaningful biomedical associations that reflect the underlying complexity of healthcare knowledge.

Cosine Similarity Between Top Disease Embeddings

After training Node2Vec embeddings on the PrimeKG graph, it becomes possible to quantify how “similar” any two nodes are in the latent space using cosine similarity. The heatmap below visualises pairwise similarities between a curated set of disease nodes, helping assess whether the learned embeddings reflect intuitive medical relationships.

The heatmap above shows the cosine similarity between the learned embeddings of three related disease nodes: hypertension, insulin resistance, and metabolic syndrome. Each cell represents the cosine similarity score between two disease embeddings. As expected:

• A score of 1.00 (along the diagonal) reflects perfect self-similarity.
• Values closer to 0.00 suggest low or orthogonal similarity.
• Higher off-diagonal values indicate that the model sees those diseases as structurally or contextually similar in the graph.

Interpreting the Patterns

• The embedding similarity between hypertension and metabolic syndrome is the highest among the pairs (0.18), which may reflect their shared connection to cardiovascular and metabolic pathways.
• Insulin resistance has modest similarity with metabolic syndrome (0.14) and hypertension (0.07), indicating weaker but non-random alignment — possibly due to sparse shared phenotypes or indirect links via common co-morbidities.
• Despite their biomedical relevance to each other, the similarity values are still low in absolute terms, which highlights the structural sparsity and specificity of disease nodes in PrimeKG.

Why This Is Useful

This kind of similarity analysis provides a semantic lens into the embedding space — giving us clues about how the model interprets disease relationships based on graph structure. It can be valuable for:

• Clustering diseases by mechanism or shared phenotypes
• Identifying potential co-morbidities based on shared neighbourhoods
• Prioritising links for drug repurposing or phenotype prediction
• Filtering noise by removing structurally irrelevant candidates in downstream tasks

Why Are Similarity Scores Still So Low?

At first glance, one might expect diseases like insulin resistance and metabolic syndrome to have much higher similarity. But here’s why the scores remain low — and why that’s not necessarily a flaw:

1. Node2Vec is structure-aware, not domain-aware

Node2Vec learns from walk patterns, not domain semantics. Two diseases might be biologically related but embedded in separate neighbourhoods if they don’t share enough graph connectivity.

2. Cosine similarity focuses on direction, not magnitude

Cosine similarity captures directional alignment, but ignores vector magnitude. So even two influential nodes with meaningful overlap might show low similarity if they vary in connectivity or feature strength.

Link Prediction with Logistic Regression

Once the Node2Vec model has learned low-dimensional vector representations (embeddings) for each node in the graph, the next step is to predict whether two nodes should be connected — even if they currently aren’t. This task is called link prediction.

Think of it like asking:

“Based on their embedding vectors, is there a high chance that disease A and phenotype B are biologically connected?”

Why Embeddings Matter Here

Each node (like asthma or IL6 protein) is now represented by a 128-dimensional vector that encodes its structural and contextual role in the graph. These embeddings serve as features for traditional machine learning models.

Step 1: Extract Positive Edges for a Specific Relation

This begins by collecting real edges for a biomedical relation of interest — for example, "disease_phenotype_positive". These are the known connections that serve as positive training examples.

relation_edges = np.sort(
    df[df["relation"] == relation_name][["x_name", "y_name"]].values.astype("U"),
    axis=1
)
relation_edges = np.unique(relation_edges, axis=0)

• The dataset is filtered for the desired relation type.
• The node pairs are sorted alphabetically to treat edges as undirected.
• The duplicates are removed with np.unique to ensure each positive edge is counted only once.

Step 2: Collect Valid Nodes by Type

The list of valid nodes is extracted for the given source and target types (e.g., disease, phenotype) from the cleaned node metadata:

src_nodes = list(node_df[node_df["node_type"] == src_type]["node_name"])
tgt_nodes = list(node_df[node_df["node_type"] == tgt_type]["node_name"])

This ensures that sampling is being done from the correct subsets when generating negatives.

Step 3: Map Positive Edges to Global Node IDs

All the valid node pairs in the positive edge set into their corresponding global integer IDs (as required for embedding lookup and modelling):

pos_edges = np.array([
    [node_maps[x], node_maps[y]]
    for x, y in relation_edges
    if x in node_maps and y in node_maps
])

Step 4: Generate Negative Samples

Since link prediction is a binary classification task, we also need negative examples — node pairs that are not connected in the graph. For this reason the same number of fake edges are generated by randomly sampling node pairs that don’t exist in the original relation set:

num_samples = len(pos_edges)
neg_edges = np.array([
    [node_maps[random.choice(src_nodes)], node_maps[random.choice(tgt_nodes)]]
    for _ in range(num_samples)
])

Note: These are synthetic and may occasionally include real but unlabelled edges — which introduces noise, but is common in graph-based negative sampling.

Step 5: Train–Test Split

The positive and negative edges are split separately into 80% training and 20% testing:

pos_train, pos_test = train_test_split(pos_edges, test_size=0.2, random_state=42)
neg_train, neg_test = train_test_split(neg_edges, test_size=0.2, random_state=42)

Step 6: Compute Edge Features

Each edge (whether positive or negative) is represented by a dot product of its two node embeddings:

def edge_features(edges):
    return (embeddings[edges[:, 0]] * embeddings[edges[:, 1]]).sum(dim=1).view(-1, 1)

• The dot product measures vector alignment — a simple proxy for similarity.
• Higher values suggest a stronger connection between the two nodes.

Then the feature matrices and labels are constructed:

X_train = torch.cat([edge_features(pos_train), edge_features(neg_train)], dim=0).cpu().numpy()
y_train = np.array([1] * len(pos_train) + [0] * len(neg_train))

X_test = torch.cat([edge_features(pos_test), edge_features(neg_test)], dim=0).cpu().numpy()
y_test = np.array([1] * len(pos_test) + [0] * len(neg_test))

Step 7: Train Logistic Regression

model = LogisticRegression(class_weight="balanced", max_iter=1000)
model.fit(X_train, y_train)

• class_weight="balanced" helps account for class imbalance.
• max_iter=1000 ensures convergence for larger datasets.

Step 8: Score a Specific Node Pair

Lastly, the model can be queried for a specific disease–phenotype pair — like:

"permanent neonatal diabetes mellitus" ↔ "retinopathy"

The dot product is computed, passed through the classifier, and a probability is returned:

pair_feat = (embeddings[u] * embeddings[v]).sum().item()
pair_score = model.predict_proba(np.array([[pair_feat]]))[:, 1][0]

Logistic Regression for Disease–Phenotype Link Prediction

This evaluation tests whether simple logistic regression on Node2Vec embeddings can predict meaningful biomedical associations. Specifically, it targets the **disease_phenotype_positive** relation—i.e., known links between diseases and observable phenotypes.

The classifier was trained using dot-product-based features between node embeddings for each disease–phenotype pair. The results are broken down into:

Evaluation Metrics

Metric Value Description Accuracy 0.5009. The model is only marginally better than random guessing (which would be ~0.50 in a balanced binary setup). Precision is 0.5010 which is slightly more than half the predicted links are correct. Recall is 0.4677 which means the model missed out a fair number of actual links. F1 Score is 0.4709 which is a harmonic mean of precision and recall — reflects overall balance. ROC-AUC 0.5100 Shows poor separation between true and false links — close to chance level (0.5).

These metrics highlight a limitation: although embeddings are informative, logistic regression alone cannot capture the complexity of biomedical graph structures.

Specific Pair Score

The logistic regression model was also queried for a specific edge:

“permanent neonatal diabetes mellitus” ↔ “retinopathy”

• Predicted probability: 0.5027

This score is barely above 0.5, suggesting the model has low confidence in this edge’s existence.

Summary

This experiment demonstrates that logistic regression over simple dot-product embeddings is insufficient for nuanced biomedical link prediction. While this setup works as a baseline, it motivates the use of more powerful models like XGBoost, GNNs, or transformer-based approaches for improved prediction quality.

Will XGBoost be any better?

XGBoost for Link Prediction on Biomedical Graphs

To improve upon the earlier baseline, an XGBoost classifier was trained using concatenated Node2Vec embeddings for disease–phenotype pairs. This setup leverages tree-based learning to better capture nonlinear relationships between node representations in the biomedical graph.

The evaluation focused again on the **disease_phenotype_positive** relation.

Performance Metrics

• Accuracy: 0.80 Over 80% of the predictions were correct.
• Precision: 0.80 High precision means most predicted links are actually relevant.
• Recall: 0.79 The model successfully retrieved a large portion of the true links.
• F1 Score: 0.80 Reflects a strong balance between precision and recall.
• ROC-AUC: 0.88 Indicates a decent discrimination between positive and negative link predictions.

These results reflect a clear performance boost over logistic regression, highlighting XGBoost’s ability to capture richer, non-linear patterns from the node embeddings.

Specific Pair Score

The model was queried for the link:

“permanent neonatal diabetes mellitus” ↔ “retinopathy”

• Predicted probability: 0.4467

Interestingly, while overall performance is strong, the probability for this specific pair is lower than expected, possibly due to data sparsity or lack of direct co-occurrence in the walk-based embedding generation process.

XGBoost proves to be a powerful link predictor in the biomedical domain when trained on structural node embeddings. It outperforms logistic regression by a large margin and serves as a strong baseline for future comparison with more complex models like Graph Neural Networks or attention-based link predictors.

Conclusion

This blog explored the use of Node2Vec embeddings on the PrimeKG biomedical graph, followed by traditional machine learning models (Logistic Regression and XGBoost) for link prediction between disease and phenotype nodes. While XGBoost outperformed Logistic Regression with significantly better precision and AUC scores, both models struggled to capture the complex semantics of biomedical relationships. The cosine similarity heatmap further revealed that even with high-dimensional embeddings, the latent space remained weakly informative when it came to reflecting true biological proximity.

This outcome highlights an important limitation: traditional ML models operating on static embeddings are not sufficient for relational reasoning in multi-relational graphs like PrimeKG. They treat the problem as a classification task over vector pairs, overlooking the rich contextual interactions and multi-hop dependencies within the graph.

The full code is available on Github

Coming Up Next:

XplainMD Part 3: Relational GCN + GNNExplainer: Learning & Explaining Links

In this blog we explored how shallow models like Node2Vec + XGBoost can uncover patterns in biomedical graphs. Now, it’s time to level up.

In the next part of this series, we dive into Relational Graph Convolutional Networks (R-GCN) — a graph-native neural architecture built to learn directly from multi-relational knowledge graphs like PrimeKG.

Unlike traditional pipelines, R-GCN dynamically updates node representations based on both edge types and neighbourhood structure, capturing the true semantics of biomedical relationships.

But we won’t stop at prediction.

Explainability will take center stage as GNNExplainer will be introduced, a tool that reveals the “why” behind each link prediction — uncovering the subgraph structures and features that drive the model’s decisions.

This next post will show how R-GCN + GNNExplainer work together to produce trustworthy, interpretable insights — a must-have in domains like drug discovery, clinical reasoning, and precision medicine.

Stay tuned — as this one’s where machine learning meets meaning.

References:

1. Grover, A. and Leskovec, J., 2016, August. node2vec: Scalable feature learning for networks. In Proceedings of the 22nd ACM SIGKDD international conference on Knowledge discovery and data mining (pp. 855–864): https://arxiv.org/abs/1607.00653