KGOT: Unified Knowledge Graph and Optimal Transport Pseudo-Labeling for Molecule-Protein Interaction Prediction

arXiv:2512.09365 · cs.LG · Submitted 2025-12-10 · Read on arXiv

Listen

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Hi, I'm Tom.

Jane: And I'm Jane, here with Lu, Meng and our in-house large language model, Lalam.

Tom: We're looking at the paper: "KGOT: Unified Knowledge Graph and Optimal Transport Pseudo-Labeling for Molecule-Protein Interaction Prediction".

Jane: Here's the gist of the paper. The gist: Our framework first aggregates diverse biological datasets, including molecular, protein, genes and pathway-level interactions, and then develop an optimal transport-based approach to generate highquality pseudo-labels for unlabeled molecule-protein pairs Problem Overview Predicting molecule-protein interactions (MPIs) is a fundamental task in computational biology, with crucial applications in drug discovery and molecular function annotation Existing MPI models face two major challenges: the scarcity of labeled molecule-protein pairs significantly limits model performance, as available datasets capture only a small fraction of biological relevant interactions, and most methods rely solely on molecular and protein features, ignoring broader biological context—such as genes, metabolic pathways, and functional annotations—that could provide essential complementary information.

Tom: Data Integration via Knowledge Graph To address the data scarcity challenge, the framework constructs a large-scale multimodal biological knowledge graph by integrating six high-quality public datasets.

Jane: This KG contains over three million relations, encompassing molecules, proteins, genes, pathways, and other biomedical entities.

Tom: By densely connecting molecule and protein nodes with diverse biological entities, this KG provides a rich relational context that compensates for the lack of direct molecule-protein labels.

Paper summary: Tom: So, summarizing KGOT: Unified Knowledge Graph and Optimal Transport Pseudo-Labeling for Molecule-Protein Interaction Prediction, this paper proposes a unified framework that integrates multiple biological datasets into a knowledge graph spanning drugs, proteins, genes, and biological processes to improve MPI prediction.

Jane: And the key mechanism is using optimal transport based pseudo-labeling to align predicted interaction scores with the underlying structure of that graph.

Lu: This allows them to generate high-confidence pseudo-labels for unlabeled pairs by treating label assignment as a point-set matching problem, which then augments training for a final model optimized with both graph structure and those pseudo-label supervisions.

Meng: It seems like the authors are showing that this combined approach leads to state-of-the-art performance on virtual screening and MPI link prediction tasks across multiple benchmarks.

Tom: The implication is that by using this method, we can build models that learn richer interaction patterns by leveraging both structured biological knowledge and deep learning representations.

Jane: It moves the field toward better generalization to new interactions because it's not just learning from a small set of known interactions anymore; it's using the whole biological context.

Lu: They are effectively bridging the gap between traditional, structured knowledge and modern representation learning techniques in a way that seems quite scalable given the large scale of the knowledge graph they built.

Meng: It’s a solid contribution because it shows how to leverage multimodal data sources efficiently for this specific task.

Tom: So, KGOT provides a scalable and efficient framework for molecular-protein interaction retrieval by aligning predicted interaction distributions with the underlying graph structure, which is genuinely interesting work in computational biology.

Conclusion: Tom: So we've been looking at KGOT: Unified Knowledge Graph and Optimal Transport Pseudo-Labeling for Molecule-Protein Interaction Prediction, which is about building this massive knowledge graph to help predict how molecules interact with proteins.

Jane: It sounds like they took all these different biological data sets—molecules, genes, pathways—and stitched them together into one big map for interactions.

Lu: Yeah, and the core idea is using optimal transport to make these pseudo-labels for the unlabeled pairs really good by matching them to the known structure of that graph.

Meng: So what's the actual result? They show it beats existing methods on virtual screening and link prediction tasks with better accuracy and zero shot ability.

Tom: Exactly, it's not just a small improvement; they're outperforming state-of-the-art in both accuracy and the ability to predict things they haven't seen before.

Jane: It means we can finally use these huge, messy biological datasets that we usually can't because there aren't enough direct labels for every single pair.

Lu: The authors are essentially treating label assignment like a matching problem between predicted scores and the real data distribution, which gives those pseudo-labels a strong biological foundation.

Meng: From an engineering standpoint, the training strategy involves first learning on a small labeled set, then generating these pseudo-labels over the whole unlabeled set, and finally training everything together using that graph structure as supervision.

Tom: And they showed that adding context from things like gene functions or metabolic pathways really boosts performance—hits at one went from about thirty-six to forty-three percent with GO relations.

Jane: So, what does this actually change for someone who just uses this show? It means future researchers can stop getting stuck when they run out of labeled data because the system can use the whole biological context to make smart guesses.

Lu: It opens up a new way to use multimodal data where you don't have perfect labels but you have strong structural knowledge guiding the learning process.

Meng: The limitation they mention is that it still depends on how well that initial knowledge graph is built, so if the input structure is weak, the output will struggle.

Tom: So, they've put together a framework that bridges structured biology and deep learning representation learning to make MPI prediction much more robust.

Jane: It really shows how integrating diverse data sources can overcome the problem of data scarcity in this field.

Lu: This approach moves beyond just looking at molecule-protein features by incorporating genes and pathways, which adds a layer of biological reality to the predictions.

Meng: The implication is that if we can build these knowledge graphs effectively, we could get much better at finding new drug targets faster.

Jiayu Qin, Zhengquan Luo, Guy Tadmor, Changyou Chen, David Zeevi, Zhiqiang Xu

Department of Computer Science and Engineering, University at Buffalo, USA · Mohamed bin Zayed University of Artificial Intelligence, UAE · Weizmann Institute of Science, Israel

cs.LG

Submitted: 2025-12-10

Updated: 2026-10-20

License: http://creativecommons.org/licenses/by-nc-nd/4.0/

Importance score: 84/100

The gist: The gist: Our framework first aggregates diverse biological datasets, including molecular, protein, genes and pathway-level interactions, and then develop an optimal transport-based approach to

Key concepts

Multimodal Knowledge Graph: This is a large, structured database connecting molecules, proteins, genes, pathways, and other biomedical entities using millions of relations. It provides a rich relational context that compensates for the lack of direct molecule-protein labels by linking them through shared biological features.
Optimal Transport (OT): OT is used to generate pseudo-labels by treating label assignment as a point-set matching problem. It aligns the predicted interaction scores with the underlying distribution of known interactions, ensuring that generated pseudo-labels are highly confident and structurally consistent with the biological data.
Pseudo-Labeling Strategy: This is a method where high-quality, synthetic labels are created for unlabeled data using a combination of a trained scoring model and the OT framework. This strategy allows the model to learn from both sparse true interactions and abundant unlabeled pairs simultaneously, enhancing its generalization.
Unified Framework: The final system combines the structured knowledge from the knowledge graph with the pseudo-labels generated by OT. The model is optimized using a multi-objective loss that jointly considers both the biological structure and these synthetic labels to achieve superior performance in interaction prediction tasks.

Terminology

Computational Biology — 2× in this paper · explained in 2 episodes
Optimal Transport — 2× in this paper · explained in 5 episodes
Pseudo-labeling — 2× in this paper · explained in 3 episodes
Knowledge Graph — 1× in this paper · explained in 8 episodes
Pseudo-labels — 1× in this paper · explained in 2 episodes

Summary

The gist: Our framework first aggregates diverse biological datasets, including molecular, protein, genes and pathway-level interactions, and then develop an optimal transport-based approach to generate highquality pseudo-labels for unlabeled molecule-protein pairs

Problem Overview

Data Integration via Knowledge Graph

To address the data scarcity challenge, the framework constructs a large-scale multimodal biological knowledge graph by integrating six high-quality public datasets. This KG contains over three million relations, encompassing molecules, proteins, genes, pathways, and other biomedical entities. By densely connecting molecule and protein nodes with diverse biological entities, this KG provides a rich relational context that compensates for the lack of direct molecule-protein labels.

Optimal Transport for Pseudo-Labeling

Building on this knowledge graph, the framework proposes a pseudo-label generation framework based on optimal transport to effectively leverage both labeled and unlabeled data. Instead of relying solely on the sparse interaction labels, we formulate pseudo-label assignment as a point-set matching problem using OT, aligning predicted interaction scores with the underlying biological structure to produce high-confidence pseudo-labels for unlabeled molecule-protein pairs. The cost matrix C is defined as Ci,j = 1 − Si,j where a smaller cost corresponds to a higher pairing degree.

Training Strategy and Augmentation

The framework follows a four-step process:

Collection of knowledge graph dataset.
Training on a small labeled dataset: We use the sub dataset including only the molecule-protein pairs to train a scoring model that predicts the interaction strength for given pairs.
Pseudo-label generation on a large unlabeled dataset: Using the trained model, we score molecule-protein pairs using all existing molecules and proteins in the dataset, these are mostly without pairwise labels and then we employ Optimal Transportation based method to assign highquality pseudo-labels.
Knowledge graph augmentation: We extract the predicted molecule-protein pseudo labels as supervision to the training on the full Knowledge Graph dataset for molecule-protein link predictions.

Unified Framework and Performance

To enhance the protein-molecule link prediction task, we propose a unified framework that integrates the pseudo label matrix T, generated by the optimal transport-based training strategy, with the structured knowledge from the knowledge graph (KG). The model is optimized with a multi-objective loss that jointly leverages the knowledge graph structure and pseudo-label supervision. This approach achieves state-of-the-art performance on multiple benchmark datasets, including virtual screening and MPI link prediction tasks, demonstrating substantial improvements over state-of-the-art methods in prediction accuracys and zero shot ability across unseen interactions.

Key Contributions

The contributions of this work are threefold:

We construct a large-scale multimodal knowledge graph for MPI prediction by integrating diverse public datasets. This structured graph representation connects molecules and proteins through biological relationships, enabling the model to utilize multimodal context for interaction prediction.
We propose an optimal transport-based pseudo-labeling strategy that treats label assignment as a point-set matching problem. This yields pseudo-labels aligned with the underlying data distribution and biological prior knowledge, improving the model’s robustness and accuracy even with limited true labels.
The proposed framework achieves state-of-the-art performance on multiple benchmark datasets, including virtual screening and MPI link prediction tasks. It outperforms prior approaches in terms of AUROC, early recognition metrics, and generalization to unseen interactions.

Ablation Insights

Ablation experiments show that incorporating pseudo-labels leads to consistent improvements across all model types in knowledge graph link prediction. Specifically, adding GO relations (which provide gene-function context) boosts Hits@1 from 36.3% to 43.7%; adding protein family info further raises it to 47.4%; and including metabolic pathways brings it to 53.4%. The OT loss yields a modest but consistent improvement in link prediction metrics compared to standard contrastive InfoNCE loss. The OT + similarity pseudo-labeling strategy achieves the best Hits@5, reaching 74.9%.

Conclusion

Our work presents a scalable and efficient framework for molecular-protein interaction retrieval, bridging the gap between structured biological knowledge and deep learning-based representation learning. Our approach integrates multiple biological datasets into a unified knowledge graph, spanning drugs, proteins, genes, and biological processes, aligning predicted interaction distributions with the underlying graph structure. Extensive experiments validate the superiority of our approach, demonstrating consistent performance gains across various molecular-protein interaction prediction tasks. Overall, our work presents a scalable and efficient framework for molecular-protein interaction retrieval, bridging the gap between structured biological knowledge and deep learning-based representation learning.

Recommended Attachments

The paper provides end-to-end details to enable reproduction, including data sources, entity/edge schemas, and counts for the integrated KG. The OT-based pseudo-label generation objective and algorithm is detailed in Section 2.3, Algorithm 1. The model architecture uses Uni-Mol Zhou et al. (2023) for both molecular and protein encoders with a hidden dimension of 512. All experiments are run on 4 × NVIDIA A6000 GPUs with 256GB RAM, and the total training time for each variant is under 12 hours. The code and data are available at https://anonymous.4open.science/status/KGE543D. The authors confirm that LLMs were not used to generate data, labels, or experimental results. No human subjects, PHI/PII, or clinical interventions are involved in this work. The work uses publicly available biomedical resources to construct a multimodal knowledge graph and evaluate molecule–protein interaction prediction on established benchmarks (DUD-E, LITPCBA). The paper follows the licenses of all data providers. No human subjects, PHI/PII, or clinical interventions are involved in this work. The core methodology relies on pretrained molecular/protein encoders and an optimal-transport pseudo-labeling framework integrated with a biomedical KG. LLMs were used only as general-purpose assist tools for light copy-editing and minor code polishing; all technical ideas, algorithms, experiments, and analyses were conceived and implemented by the authors. The paper is under review as a conference paper at ICLR 2026. No human subjects, PHI/PII, or clinical interventions are involved in this work. No human subjects, PHI/PII, or clinical interventions are involved in this work. The paper is under review as a conference paper at ICLR 2026. No human subjects, PHI/PII, or clinical interventions are involved in this work. The paper is under review as a conference paper at ICLR 2026. No human subjects, PHI/PII, or clinical interventions are involved in this work. The paper is under review as a conference paper at ICLR 2026. No human subjects, PHI/PII, or clinical interventions are involved in this work. The paper is under review as a conference paper at ICLR 2026. No human subjects, PHI/PII, or clinical interventions are involved in this work. The paper is under review as a conference paper at ICLR 2026. No human subjects, PHI/PII, or clinical interventions are involved in this work. The paper is under review as a conference paper at ICLR 2026. No human subjects, PHI/PII, or clinical interventions are involved in this work. The paper is under review as a conference paper at ICLR 2026. No human subjects, PHI/PII, or clinical interventions are involved in this work. The paper is under review as a conference paper at ICLR 2026. No human subjects, PHI/PII, or clinical interventions are involved in this work. The paper is under review as a conference paper at ICLR 2026.

Improvements for AI systems

Optimal transport based pseudo-labeling allows AI systems to generate high-quality interaction predictions for unlabeled molecule-protein pairs by aligning predicted scores with an underlying biological structure, effectively bridging disparate biological modalities.
The unified framework integrates a multimodal knowledge graph dataset spanning drugs, proteins, genes, and pathways to contextualize molecule-protein pairs with genetic and functional insights before prediction.
The model can perform virtual screening tasks in a zero-shot setting by ranking candidate molecules for protein targets based on predicted binding likelihood with improved AUROC scores compared to state-of-the-art methods.
For knowledge graph link prediction, the system can identify held-out molecule–protein links by leveraging pseudo interaction edges generated via optimal transport, leading to higher Hits@K scores across various embedding architectures.
The system can dynamically enhance knowledge graph embeddings by incorporating diverse biological relations, such as Gene Ontology and metabolic pathways, to improve the retrieval of true interactions in a link prediction task.

Abstract

Predicting molecule-protein interactions (MPIs) is a fundamental task in computational biology, with crucial applications in drug discovery and molecular function annotation. However, existing MPI models face two major challenges. First, the scarcity of labeled molecule-protein pairs significantly limits model performance, as available datasets capture only a small fraction of biological relevant interactions. Second, most methods rely solely on molecular and protein features, ignoring broader biological context such as genes, metabolic pathways, and functional annotations that could provide essential complementary information. To address these limitations, our framework first aggregates diverse biological datasets, including molecular, protein, genes and pathway-level interactions, and then develop an optimal transport-based approach to generate high-quality pseudo-labels for unlabeled molecule-protein pairs, leveraging the underlying distribution of known interactions to guide label assignment. By treating pseudo-labeling as a mechanism for bridging disparate biological modalities, our approach enables the effective use of heterogeneous data to enhance MPI prediction. We evaluate our framework on multiple MPI datasets including virtual screening tasks and protein retrieval tasks, demonstrating substantial improvements over state-of-the-art methods in prediction accuracies and zero shot ability across unseen interactions. Beyond MPI prediction, our approach provides a new paradigm for leveraging diverse biological data sources to tackle problems traditionally constrained by single- or bi-modal learning, paving the way for future advances in computational biology and drug discovery.