KGOT: Unified Knowledge Graph and Optimal Transport Pseudo-Labeling for Molecule-Protein Interaction Prediction
Authors: Jiayu Qin, Zhengquan Luo, Guy Tadmor, Changyou Chen, David Zeevi, Zhiqiang Xu
Organizations: Department of Computer Science and Engineering, University at Buffalo, USA · Mohamed bin Zayed University of Artificial Intelligence, UAE · Weizmann Institute of Science, Israel
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.
Figures & tables
Figure 1: Overview of KGOT. (A) Knowledge integration (B) Supervised score learning (C) Pseudo-label generation (D) KG augmentation & link prediction
Table 1: Results on the DUD-E virtual screening benchmark (zero-shot setting). Higher is better for all metrics. Our OT + KG framework outperforms both traditional docking methods and modern learning-based approaches across all evaluation metrics.
Table 2: Results on the LIT-PCBA benchmark (zero-shot setting). Our method achieves the best performance across all metrics, illustrating its robustness on this more challenging dataset.
Method
Hits@1
Hits@3
Hits@5
PairRE
10.0%
17.0%
21.4%
PairRE + KGOT
10.9%
20.5%
26.4%
RotatE
48.5%
61.6%
66.6%
RotatE + KGOT
52.0%
63.9%
68.0%
MuRE
15.9%
24.7%
31.0%
MuRE + KGOT
13.7%
25.9%
31.2%
Table 3: Knowledge graph link prediction performance (Hits@K) with and without KGOT generated pseudo-label augmentation.
Appendix figures & tables7 assets
Supplementary material from the paper’s appendix.
Appendix
Method
Unified multimodal KG
KG denoising
Multi-task learning
Global OT pseudo-labeling
KG-MTL Ma et al. [2022]
×
×
✓
×
BioKDN Ma et al. [2024]
×
✓
×
×
KGOT (ours)
✓
×
×
✓
Appendix
Table 4: Qualitative comparison between representative KG-based molecular interaction prediction methods and KGOT. KGOT targets weakly-/unlabeled MPI retrieval via global OT pseudo-labeling and KG write-back, rather than multi-task supervised learning or KG denoising.
Head Entities
Tail Entities
Type
Quantity
Type
Quantity
UNIPROT
5,956,325
GO
3,191,321
CHEBI
336,374
CHEBI
1,678,407
KEGG
92,184
PFAM
792,235
GO
89,235
KEGG_KO
407,307
EC
8,459
EC
304,428
Appendix
Table 5: Entities Distribution, where KMODUL represents KEGG_MODUL, KPATHWAY represents KEGG_PATHWAY, KCOMPOUND represents KEGG_COMPOUND, and KMODULE represnts KEGG_MODULE.
Table 8: Performance of the scoring model with OT-based loss vs. contrastive InfoNCE loss. OT loss yields a modest but consistent improvement in link prediction metrics.
Strategy
Hits@5 (%)
No pseudo-label augmentation
68.5
Random pseudo-labels (matched count)
66.0
Top- k per protein ( k=5 )
72.0
OT without similarity ( λ=0 )
73.8
OT with high entropy ( ϵ=0.1 )
73.0
OT + similarity ( λ=0.1,ϵ=0.01 )
74.9
Appendix
Table 9: Ablation on pseudo-label generation (illustrative Hits@5 on link prediction).
Department of Computer Science and Engineering, The Hong Kong University of Science and Technology, Clear Water Bay, Hong Kong. · Lee Shau Kee School of Business and Administration, The Hong Kong Metropolitan University, Ho Man Tin, Kowloon, Hong Kong.