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
Venue: ICLR 2026 Poster (2026)
Keywords: Knowledge graph
molecule protein interaction
optimal transport
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)

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.

Published as a conference paper at ICLR 2026
KGOT: UNIFIED KNOWLEDGE GRAPH AND OPTIMAL TRANSPORT PSEUDO-LABELING FOR MOLECULE- PROTEIN INTERACTION PREDICTION
Changyou Chen² David Zeevi Zhiqiang Xu² l DDpartmeet of Computer Science and Engineering, University at Buffalo, USA
Jiayu Qin Zhengquan Luo² Guy Tadmor'
2 Mohamed bin Zayed University of Artificial Intelligence, UAE
"Weizmann Institute of Science, Israel
{j Jjiayugin, changyou] (erryouuubufuaooo @buffaloo edu zhengquan.luo, zhiqians xu) euuzembuaai ac ae (guy. tadmor, david. zeevi @@ee zmann ac. il
ABSTRACT
Predicting molecule-protein interactions (MPIs) is a fundamental task in compu- tational biology, with crucial applications in drug discovery and molecular func- tion annotation. However, existing MPI models face two major challenges. First, the scarcity of labeled molecule- polecule-protein pairs significantly limits model perfor- mance, as available datasets capture only a small fraction of biologically relevant interactions. Second, most methods rely solely molecular and protein fea- tures, ignoring broader biological context-such as genes, metabolic pathways, and functional annotations-thar could provide essential complementary infor- mation. To address these limitations, our framework first aggregates diverse bi- ological datasets, including molecular, protein, genes and pathway-level interac- tions, and then develops an optimal transport- based approach to generate high- quality pseudo-labels for unlabeled molecule-protein pairs, leveraging the under- lying 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 diction. We evaluate our framework on multiple MPI datasets including virtual screening tasks and protein retrieval tasks, demonstrating substantial improve- ments over state-of-the-aar methods in prediction accuracy 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 tra- ditionally constrained by single- bi-modal learning, paving the way for future advances in computational biology and drug discovery.
of heterogeneous data to enhance MPI pre-
INTRODUCTION
Molecular and protein representation learning is an increasingly important topic in computational biology and drug discovery (Jumper al. 202T: Zhou et al. (202)), fueled by the availability of
large-scale unlabeled datasets (Consortium, 2024 Guo et 2005| AlQuraish] 2019 Nakata et al., 2020). These resources have enabled the development powerful molecular and protein encoders, which serve foundational models for various downstream tasks. For example, self-supervised learning approaches (Rives et al.) 2019 Zhou et al., 2023) leverage massive sequence and structure databases to capture intricate biochemical properties without requiring explicit supervision. This line research not only advances standalone molecularr/rotein modeling but also lays the founda- tion for pushing the boundaries of computational biology and accelerating medicine and life science discoveries.
What is new KGOT? While prior work explores knowledge-enhane enhanned drug- target prediction, existing approaches typically (i) treat knowledge graphs auxiliary feature source and train observed edges, or (ii) apply pseudo-labeling heuristics that do not explicitly enforce global con- sistency. KGOT instead introduces unified Optimal Transportation Knowledge Graph(KG)

---

Published as a conference paper at ICLR 2026
A
B
Preddoollpiic wileey
Encader
Malk: Predecud
Valre IOT1
Excyees
Cateydddddio
Agerto
Packe
(Grund Tualty
to
CROA
Vanoar
Encade
Source
Thende CaMd advit
Medel (itrat))
Milactee
Pasulo Laad
Danax
Camay
RLSG
Tulwers
Tokers
Vikailll
Figure Overview of KGOT. (A) Knowledge integration (B) Supervised score learning (C)
Pseudo-labe label generation (D) KG augmentation & link prediction
framework in which (1) an inverse optimal-transpon (IOT) objective supervises mutual molecule- protein retrieval (Sec. 2.3), (2) the resulting OT transport plan is written back into the knowledge graph as new issudd_interaa
on relation, and (3) KG embeddings and retrieval models are trained jointly the augmented graph (Sec. 2.4). This closes the loop between score learn- ing, pscudo-label generation, and knowledge-graph training, enabling label-efficient MPI prediction with multimodal biological context.
Despite these advances, retrieving molecule- proteule-protein interactions remains a formidable challenge. core issue is the scarcity of large-scale labeled datasets for MPIs, due to the experimental com- plexity and cost of validating interactions. Fach new molecule-protein interaction must typically be confirmed via laborious assays, often with regulatory oversight drug discovery, which severely limits data growth. High-throughput screening assays are expensive and slow, and even computa- tional docking simulations (Di Nola et al. 1994) are constrained by accuracy and scale. Existing MPI datasets (Chandak et al. 2023 Gao et al. (223)) tend to be small, biased toward specific protein Gao
families, or inconsistent across annotations, making it difficult to train deep models that generalize across diverse interactions better than traditional techniques (such docking and energy scoring (Wang et al., 2020)). If more labeled molecule-proter pairs were available spanning diverse chem- istry and targets, deep learning models could learn richer interaction patterns and capture complex biophysical properties that traditional methods struggle with. Recent studies (Xia et al. (024) un- derscore the need for novel frameworks that combine modern representation learning with strategies to cope with limited labels, in order to yield more robust and generalizable interaction predictions.
Another major challenge is the narrow reliance on bi-modal data, i.e., using only molecular and protein features while ignoring other relevant biological information. In real-world biology and drug discovery, molecular interactions are influenced by a broad spectrum of factors. For instance, genetic variations can alter protein function and a molecule's binding efficacy; biochemical path- ways and networks can modulate the downstream effects of drug and indicate which proteins are likely involved. These additional modalities provide essential context for understanding interactions beyond what comes with molecular structure or protein sequence alone.
Large-scale biological knowledge graphs (KGs) offer promising way to integrate these heteroge- neous modalities. Resources like PRIMEKG (Chandak et al., 2023) aggregate diverse biological entities and relations, but they contain relatively few direct molecule- prottenn interactions and not tailored for MPI prediction. Our work aims to bridge this gap by systematically integrating mul- timodal biological data into the MPI prediction process. By using biological KG structured
2

---

Published as a conference paper at ICLR 2026
prior, contextualize molecule protein pairs with genetic, biochemical, and phenotypic insights, which can improve prediction robustness and generalization. In our framework, we extract relevant information from diverse sources into unified graph and then refine the resulting representations through a pseudo-labeling mechanism based on optimal transport, aligning the multimodal knowl- edge with the MPI prediction task.
To address the data scarcity challenge, we constructed a large-scale multimodal biological knowl- edge graph by integrating six high-quality public datasets. The resulting KG contains over three million relations, encompassing molecules, proteins, genes, pathways, and other biomedical enti- ties. 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-prose labels. It also enables the model to capture indirect interaction paths, potentially improving generalization to new interactions.
Building on this knowledge graph, we propose a pseudo-label generation framework based on op- timal 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, we align predicted interaction scores with the underlying biological structure to produce high-confidence pseudo-labels for unlabeled molecule-proted pairs. Concretely, our method first trains an interaction scoring function top of molecule and protein encoders. We then infer scores pairs to construct an initial dense score matrix. Next, we apply
for all unlabeled molecule-proter an OT-baaed algorithm to assign pseudo-label interaction probabilities by minimizing transport cost between molecule and protein distributions. The resulting pscudo-labeled interactions, combined with ground-truth labels, are used to augment training of the final model. This approach effectively transforms the MPI task into semi-supervised learning problem, where the model benefits from an expanded training set of confident interaction examples.
In summary, our contributions are threefold:
We construct large-scale multimodal knowledge graph for MPI prediction by integrat- ing 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-base pseudo-labeling strategy that treats label assign- ment as a point-set matching problem. This yields pseudo-labels aligned with the under- lying 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 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.
By bridging optimal transport with knowledge-drre representation learning, our approach provides a scalable and efficient solution for MPI prediction. Extensive experiments validate its effectiveness, particularly in leveraging multimodal biological signals to enhance interaction inference.
2
METHODOLOGGY
2.1
OVERVIEW
We focus on the task of molecule-protein interaction prediction, more specifically, mutual retrieval of molecule-protein, which involves two complementary objectives: (I) Given molecule retrieve the protein capable of best catalyzing it; (2) Given protein y. retrieve the molecules that can best bind to it.
To address this task, we collect a biological knowledge graph dataset that integrates molecular, protein, and knowledge graph (KG))based methodologies. We then propose novel framework that leverages optimal transportation (OT) to generate pseudo-label of molecule-protee interactions. As shown in Figure I approach is designed follows:
(1) Collection of knowledge graph dataset. We collected over 1.35 million interactions between molecules and proteins from UniProt and CHEBI datasets, and more interactions related to genes,
3

---

Published as a conference paper at ICLR 2026
genomes, protein families and enzyme commission numbers. Please check Appendix B for details of our dataset.
(2) Training small labeled dataset: We use the subset including only the molecule-prose pairs to train a scoring model that predicts the interaction strength for given pairs.
(3) Pseudo-label generation molecule-proteii pairs using all existing molecules and proteins in the dataset, these are mostly without pairwise labels and then we employ digimal-transpott based method to assign high-quality pseudo-labels.
large unlabeled dataset: Using the trained model, score
(4) Knowledge graph augmentation: We extract the predicted molecule-protein pseudo- as supervision to the training on the full knowledge graph dataset for molecule-protein link predictions.
This framework enables the effective utilization of labeled and unlabeled data while leveraging relation information from KGs to enhance molecule-protein interaction predictions.
2.1.1 DESIGN CHOICES AND IMPLEMENTATION DETAILS.
OT cost and marginals. We define the OT cost Cjj 1 Syy. where Siy is the predicted interaction score between molecule I; and protein Yj (see Section 2.3). Unless additional priors are available, we use uniform marginals r; = M7 and e, t. which encourages balanced coverage over molecules and proteins; we discuss alternative choices and robustness in Appendix B
Entropy regularization. We solve the entropically regularized OT problem with Sinkhorn itera- tions. In our experiments we variant (e = 0.1) in the ablation study given in Appendix E.2
0.01 (see Appendix C) and additionally report a higher-entropy
How the KG encoded. KGOT is agnostic the underlying KG embedding model. We report sults with a range of standard KGE architectures (PairRE, RotatE, MuRE, TorusE, ComplEx-FFi see Section 3.2 all trained on the same augmented KG with the additional pseudo interaction
relation (see Section
2.4
Leakage control. For KG link prediction, we withhold a disjoint set of molccule-protein edges for evaluation and exclude them from both pscudo-label generation and KG training. For virtual screening, we additionally apply molecule- and protein-side filtering based on scaffold similarity and sequence identity (see Section 3.1 Appendix D).
2.2 PREPARATIONS FOR PSEUDO-LABEL GENERATION
Pretrained backbone represent molecule and protein structural information effectively, we adopt pretrained encoders based on the Uni-Mol (Zhou et al.) 2023) framework the backbone for feature extraction. Uni-Mol molecule and protein pretraining model specially designed to process 3D conformations, and it has achieved strong performance on a variety of downstream tasks including molecular property prediction and binding pose prediction.
Two encoders produce embeddings f(x) and g(y) for molecule and protein respectively. These embeddings are normalized using the Euclidean norm to ensure consistency and compatibility for later tasks.
Molecular similarity calculation To measure the similarity between molecules, we utilize the embeddings extracted by the pretrained molecular encoder. Given two molecules x, and Tjs their respective embeddings f(x;) and f2;, first computed using the molecular encoder.
The similarity between molecules x; and Iy is then quantified using the cosine similarity, which is defined
Sim(z,, = ((a)). [[(=))))
where f( is the dot product of the two embeddings, and [f l| and 1 their respective Euclidean norms. This similarity measure serves as a foundation for incorporating molecular relationships into optimal transport-based pseudo-label generation.
4

---

Published as a conference paper at ICLR 2026
2.3 PSEUDO LABEL GENERATION WITH OPTIMAL TRANSPORTATION
To estimate interaction probabilities for molecule-protein pairs, we first train a scoring function S(x,, y) the labeled subset. We extract features with pretrained encoders obtain embeddings f(x)) and g(y), and then compute scores for all molecule-protein pairs to form a score matrix S E RMxN where Sty S(x;, My). Inspired by [hh (2023), we model the global matching structure by learning an optimal transport plan T 12
VIa:
min
s.t. TIN = T 1MMM =c
Here, the cost matrix C R4X and Ci.j = 1- Stg shows the for pairing. Source and sink distributions are defined as: The source distribution E RM satisfies r; , and the sink distribution e E RR- satisfies Gj - N. 1M, IN are all-one vectors with length M and N, respectively.
Training strategy of score function on labeled dataset To train the score function S(x. y) effec- tively labeled dataset, take an inverse optimal transportation perspective. The interaction score S(x, y) between a molecule r and a protein is modeled as a learned function of their em- beddings:
S(x v) = W(x y),
where W a trainable model, and @ represents concatenation.
During training, we construct ground truth cost matrices Cgt based on current batch of positive and negative molccule-proted pairs. For given positive molccule-proten pair (æ,, yr), the ground truth cost ji) is defined such that:
Cgrli. (0: j = (positive pair) = [1, if j * (negative pair)
(3)
For a sampled batch of N molecule- proteei--pttnnnn pairs, the theoretical optimal transport matrix Tg is computed using Cg the cost matrix, following the Sinkhorn-Knopp algorithm enforce marginal constraints:
N ΣΣ (i,,),
T y = arg min
s.t. TIN = TI T
(4)
T20
where and are uniform source and sink distributions respectively.
Then we calculate the predicted transport matrix Tpeed based on the cost matrix derived from the S(x, Vj). The loss function is defined the KL divergence
predicted scores Cpeea (i,
between Teed and Tgi
(t.) (i.c log Tz((i,j)
Lycore KL pred (Ta)
(5)
1=1
This formulation ensures that learned score function S(x,, y) aligns the predicted transport matrix with the theoretical optimal transport matrix derived from ground truth labels. As described in Shi et al. (2023), contrastive learning with the InfoNCE loss can be viewed as a special form this method. By minimizing Locore get the optimal model W that maximizes the scores of true molccule-protein pairs while minimizing those of false pairs in the batch.
Pseudo label generation on large unlabeled dataset Let us consider the scenario of the whole knowledge graph where we have M molecules and N proteins. The pairing degree between molecules and proteins is represented as a matrix 5 E M where each element St reflects the score between molecule and protein j. calculated form the model we get the former step. Our goal is to generate pseudo-label matrix T E RRMx that can be used for further training. We also treat this problem as an optimal transport problem.
We define the optimal transport problem follows:
5

---

Published as a conference paper at ICLR 2026
(1) The cost matrix C' B is defined as C2, Si.. where smaller cost corresponds to a higher pairing degree.
(2) The source distribution R represents the initial distribution over the molecules, with each entry given by ri M
(3) The target distribution e E R represents the distribution over the proteins, with each entry given by Cj N-
The optimal transport problem to find a transportation matrix T E Mx that minimizes the total cost while satisfying the marginal constraints:
M
N
min
CiyCr. subject to 7IN =, T 1M =
(6)
=1j=1
where 1M and 1y are uniform distributions over lengths M and N, respectively.
In our case, we have additional information about molecular similarity represented as a matrix Sim e RMXM,, where Sim,, quantífies the similarity between molecule and molecule k. To leverage this information, we introduce additional constraint: the similarity of pseudo-labels between molecules and k; denoted by Simţ, Σÿ= TeyTke should be close possible Sim;,
The modified objective function becomes:
M N
M
M
Cig
+
(Sima.x
Sim
where À > 0 is weighting factor balancing the cost and similarity terms, we take À 0.1
As shown in the algorithm, we employ the Sinkhorn- Sinhhor-Knoop algorithm to solve the optimal transport problem efficiently and extend handle the similarity constraints.
The Sinkhorn- Knopp algorithm solves the regularized optimal transport problem by introducing an entropic regularization term:
Ci. +E
Th log Tig
(8)
1=1
The solution can be computed iteratively: (1) Initialize u 1a, v 1n. and K exp(-C/e); (2) Iterate until convergence: u f- K V & (3) Compute 7 as: T diag(u) K diag(v). We take = 0.01 for the experiments. Ablations can found in Appendix E
To incorporate similarity constraints, we modify T using gradient-basee based optimization. The gradient of the similarity term with respect to T given by:
TT., 22
(Sim Sim?.)) Tk.j
The algorithm alternates between updating T using the Sinkhorn steps and refining T with the similarity constraint gradient:
T V
(10)
where 77 > 0 is the learning rate, and T is projected back to the feasible set if needed.
Pseudo-label extraction from the OT plan. After solving the regularized OT problem, we obtain a transport plan 7 RxM that assigns soft matching mass Ti, each molecule-protein pair (I., We treat entries with large transport mass high-confidence pseudo-positive interactions:
P3 = ((x,, 1i) Ti 2 8),
(11)
use 8 = 0.5 in all main experiments unless stated otherwise. In practice, 8 controls precision- coverage trade-off: a larger yields fewer but cleaner iseudo-labels. while a smaller increases coverage but may introduce noise.

---

Published as a conference paper at ICLR 2026
Algorithm 1: Training strategy on large unlabeled dataset
Input:
Pairing score matrix S E R molecular similarity matrix Sim R source distribution IR target distribution E IR entropic regularization parameter t, similarity weight X, learning rate 1), maximum iterations max_iter. Output: Optimal transport matrix E R M2M
Initialization:
MxN Set K exp(-C/e),
Define cost matrix C' E R4
as Ci. = St.j. M. IN. and T 0. for t = to max max_iter do
Step 1: Sinkhorn-Knopp Iteration.
while not converged do
Update <- Kg'
Update V
K
Compute = diag(u) diag(uu)
Step 2: Similarity Constraint Adjustment.
Compute Sim = -Σ Tk. for all k.
Compute gradient V r., = 22 k=1 (Sim;,s Sim Tk.j.
Update T + Vr.
Step 3: Projection onto Feasible Set.
Ensure T 0, and normalize T such that TlN = and T 1M
return T.
Table I: 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 camming-based approaches across all evaluation metrics.
Model
AUROC (%)
BEDROC (%)
EF@0.5%
EF@ 1%
EF 2%
Glide SP Halgren et al. 2004
76.70
40.70
19.39
16.18
7.23
Vina Trod and Olson 2010
71.60
9.13
7.32
4.44
DDurant McCammon
3.12
RFscore Batteterr and Miteren 2010
65.21
12.41
4,90
4.52
2.98
Pafnucy Stephiewska- Dztuoinska et 2018
63.11
16.50
4.24
3.86
3.76
OnionNer Zhenn ala 2019
59.71
2,20
Planet Zhain et Bl. 2023
71.60
10.23
8.83
5,40
DrugC TIP
Dia cl ala 20446
80.93
50.52
38.07
31.89
10.66
KGOT
83.45
0.42
51.20 0.35
39.10
0.50
33.00
0.47
11,20
0.30
2.4 UNIFIED FRAMEWORK FOR LINK PREDICTION
To enhance the protein-molecut link prediction task, we propose unified framework that inte- grates the pseudo-label matrix T, generated by the optimal transport-based training strategy, with the structured knowledge from the knowledge graph (KG).
The relations in KG dataset include all observed interactions in the KG as well as new relation which encodes the pseudo-label scores T.
type, ppee.pseudo.inteeracee
Multi-objective learning framework Our model optimized with multi-objective loss that jointly leverages the knowledge graph structure and pseudo-label supervision. Specifically, we com- bine a graph embedding loss over KG triples with a pseudo-label alignment term that encourages predicted interaction scores to match the generated pseudo-label matrix. This joint formulation allows the model to balance structural knowledge with data-driven signals. Full mathematical defi- nitions and implementation details provided in Appendix
EXPERIMENTS AND RESULTS
Our evaluation spans two settings that highlight different aspects of the proposed framework: (1) virtual screening benchmarks, which test the model's ability to retrieve active molecules for given

---

Published as a conference paper at ICLR 2026
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.
Model
AUROC (%)
BEDROC (%)
HF@0.5%
EF@
EF@5%
Surflex (Jain 2003
51.47
2.50
Glide-SP" Saale-SPaallgreen
Halgren et al. 2004
53.15
4.00
3.17
3.41
2.01
Planct Zhhang et al. 2023
57.31
4.64
3.87
2.43
Gnina (McNutt al. 2021
60.93
5.40
4.63
DeepDTA (Oztürk al.
2018
56.27
2.53
1.47
BigBind Brocidiacono el
2023
60.80
3.82
DrugCLIP et 2026
57.17
6.23
8.56
5.51
2.27
KGOT
62.45 + 0.38
6.52 +: 0.22
9.12 +: 0.40
5.90 +: 0.28
2.50 0.15
protein targets in a zero-shot manner, and (2) knowledge graph link prediction, which tests the model's 's ability to identify held-out molecule protein links using the integrated KG.
pelocule-protein
3.1
EVALUATION ON VIRTUAL SCREENING TASKS
Virtual screening benchmarks evaluate how well models can rank candidate molecules for partic- ular protein target, based on predicted binding likelihood. We consider two widely-used datasets: DUD-E and LIT-PCBA. We follow the evaluation protocol of recent works like DrugCLIP (Gao et al. 2023 to ensure fair comparison.
DUD-E benchmark. The DUD-E dataset (Mysinger et 2012) contains 102 protein targets, each with set of known active molecules (binders) and a large set of decoys (inactive molecules). In total, DUD-E includes 22, 286 active molecule- target pairs. We frame each target's screening as ranking task: the model must assign higher scores active molecules than to decoys. We evaluate performance using three metrics commonly used in virtual screening: (1) AUROC (area under the ROC curve), a thresholddindeeendeen measure of ranking quality; (2) BEDROC with a 20, which emphasizes early recognition of actives (important in screening scenarios); (3) Enrichment Factor (EF) at certain top fractions (c.g., EF@0.5% 1%, 2%). which measures how many actives are found among the top-ranked subset compared to random expectation.
Leakage control. To preclude train-tess leakage the molecule side, remove from the training pool any ligand whose Tanimoto similarity to any DUD-E test active is above a threshold of 2 0.60; additionally report a stricter Hurcko-scaffold-a Bemis-Murcko scaffold with any test active are excluded (see Appendix D). On the protein side, we exclude from training any protein whose sequence identity (computed by MMseqs2) any DUD- E test target exceeds 60%; we also provide family-out control by removing all training proteins mapped (via HMMER Pfam-A) the same families DUD-E test targets. We use a single model (no ensembling or post-hoc re-ranking) for all runs.
variant in which all training ligands sharing the
Our model produces an interaction score for each molecule targgtt pair, which we use to rank molecules for each target. Table [ summarizes the results on DUD-E in the zero-shot setting. We compare to classical docking methods such as Glide-SP (Halgren et al. 2004) and AutoDock Vina (Trott and Olson, 2010), and learned baselines including similarity- based models such as NN-score DDuraat and McCammon (01I) and RFscore (Ballester and Mitchell 2010)., structure-based deep Saanniewska-Dziubinne
et al.) (0188) and OnionNer (Zheng et al.) 2019). and recent cross-modal representation models such as Planer (Zhang et al. 2023) and DrugCLiP (Jia et al. 2026). Our method achieves the highest scores on all metrics. Notably, it outperforms the previous best model by a significant margin in early recognition metrics (e.g.. BEDROC and EF), indicating that the OT-guided pseudo-labeling helps identify actives at the top of the ranking more effectively. In terms of AUROC, we obtain about 83.5%, which is an improvement of ~2.5% over DrugCLIP (80.9%) and substantially higher than traditional docking (Glide SP: 76.7%). These results demonstrate the benefit of augmenting the limited training interactions with pseudo-labeled examples and KG-derived information. Here we also report mean standard deviation over multiple runs with 3 different random seeds
models such as Pafnucy
LIT-PCBA benchmark. LIT-PCBA is another benchmark from the Therapeutic Data Commons, designed to be more challenging and address some biases in DUD-E. It comprises 15 protein targets
8

---

Published as a conference paper at ICLR 2026
with 7,844 experimentally confirmed actives and 407,381 inactives. The class imbalance is more extreme here, making early retrieval metrics critical. We evaluate our model on LIT-PCBA under the same zero-shot setting.
Table shows the results. We compare to several docking and deep learning baselines reported for this benchmark. including Surflex (Jain 2003). Glide-SP. Gnina (McNutt et al., 2021). DeepDTA Öztürk et al., 2018). BigBind (Brocidiacono et al. 2023). and DrugCLIP. Our framework again achieves the top performance on all metrics. In particular,
see improvements in AUROC (62.45% vs. 60.93% for the best baseline, Gnina) and in early enrichment (EF@0. 0.5% of 9.12 vs. 8.56 for DrugCLIP). Although the absolute values are lower than DUD-E (due LIT-PCBA's difficulty). the consistent gains indicate that our pseudo-label + KG approach generalizes well.
3.2 MOL.CCU.E- LINK PREDICTION TASK
We next evaluate our unified framework on the task of predicting pooeecule-protiie links on knowl- edge graph. For this, created held-out set of known molecule-protein interactions from collected KG data. Specifically, we withhold 60,000 molecule-proteie pairs (edges) from the KG serve as test examples, these pairs KG training. Each test pair is true interaction. The model must predict these links purely from the remaining training data the KG and the pseudo-label enriched representations.
not included in the training pseudo-label generation
We cast this as link prediction problem: given a molecule node, rank candidate protein nodes by the predicted likelihood of interaction (head entity prediction in KG terms). We evaluate using standard information retrieval metrics Hit@K, K: Hits@I I, Hits@3 and @its@5, which measure the proportion of test queries for which the correct partner appears the top top 3, top predic- tions, respectively. These metrics reflect the model's 's ability to place the true interaction at near the top of the candidate list.
We compare our full model (which uses pseudo-labels and KG, as described) against ablated ver- sions to quantify the effect pseudo-labeling. graph embedding models with and without our pseudo-label augmentation: PairRE, RotatE, MuRE, TorusE, and ComplEx-FF are five different KG embedding architectures. For each, we train one model using only the real KG edges (baseline) and another including the aseudo.interactiin edges and the alignment loss. We observe that incorporating pseudo-labels leads to consistent improvements across all model types. The absolute gains vary by model, but the trend is clear: the additional pseudo-labeled interactions help the KG embedding models better discriminate true links. Among the embedding methods, found TorusE performed strongly, and with pseudo-labels achieved the highest H (74.9%). ComplEx-FFF benefited remarkably from pseudo-labels (Hits@) rising from 30.8% to 43.6%). These results highlight that our pseudo-labeling approach is model-agnostic and can enhance range of link prediction algorithms by providing extra supervi- sion.
Table we report results for several knowledge
Table 3: Knowledge graph link prediction performance (Hits K) with and without KGOT gener- ated pseudo-label augmentation.
Method
Hits@1 1
Hits@3
Hits@5
PairRE
10.0%
17.0%
21.4%
PairRE + KGOT
10.9%
20.5%
26.4%
RotatE RotatE KGOT
48.5%
61.6%
66.6%
52.0%
63.9%
68.0%
MuRE
15.9%
24.7%
31.0%
MuRE KGOT
13.7%
25.9%
31.2%
TorusE
49,4%
64.2%
70.0%
TorusE + KGOT
53.4%
65.2%
74.9% %
ComplEX-FF
30.8%
40.2%
44.4%
ComplEx-FF + KGOT
43.6 %
54.3%
58.6%

---

Published as a conference paper at ICLR 2026
The consistent improvements demonstrate that our pseudo-label generation successfully injects ful information into the KG. By effectively "filling in" likely interactions that were not explicitly in the KG, the model can learn from of held-out interactions, validating the core idea of our approach: leveraging unlabeled pairs through OT-based pseudo-labeling boosts prediction performance.
complete interaction network. This leads to better retrieval
3.3 ABLATION STUDY.
We conduct ablation experiments to assess the contributions of our design choices, including (i) the of OT loss versus standard InfoNCE contrastive loss, (ii) different pseudo-labeling strategies, and (iii) the integration of multiple biological knowledge sources. Across all settings, our proposed OT-based formulations and multimodal knowledge graph integration consistently yield higher link prediction accuracy. For example, OT loss improves over InfoNCE, +T+immilarit pseudo-labeling achieves the best @its@5, 5, and adding GO, protein family, and pathway relations each provide cremental gains in Hits 1. Due to space limit, full details and complete results are reported in Appendix
CONCLUSION
In this work, we provide a unified biomedical knowledge graph model to tackle the challenge of molecular-protein interaction retrieval by integrating multi-modal biological data to improve molecule-protein interaction prediction and proposing a novel pseudo-label generation framework based on optimal transport (OT) to mitigate the scarcity of large-scale labeled datasets. Our approach integrates multiple biological datasets into unified knowledge graph, spanning drugs, proteins, genes, and biological processes, aligning predicted interaction distributions with the underlying graph structure, significantly improving retrieval performance across multiple benchmark datasets.
Extensive experiments validate the superiority of approach, demonstrating consistent perfor- mance gains across various molecular-protein interaction prediction tasks.
Overall, our work presents scalable and efficient framework for molecular-p palecular-protein interaction trieval, bridging the gap between structured biological knowledge and deep learning-based repre- sentation learning. We hope that our work can provide a paradigm for other data-lacking tasks in the field of computational biology, and ultimately point the way to using machine learning methods build unified foundation models in the biological field.
RECOMMENDED ATTACHMENTS
5.1 ETHICS STATEMENT
This work uses publicly available biomedical resources to construct a multimodal knowledge graph and evaluate molccule-protein interaction prediction established benchmarks (DUD-E, LIT- PCBA). No human subjects, PHI/PII, or clinical interventions are involved. We follow the licenses of all data providers.
5.2 REPRODUCIBILIIT STATEMENT
We provide end-to-end details to enable reproduction: data sources, entity/edge schemas, and counts for the integrated KG (Appendix B): leakage-contre splits and zero-shot protocols for DUD-E and LIT-PCBA (Section 3.1 Appendix [D] the OT-based pseudo-labe generation objective and algorithm (Section 2.3] Alg. 1), along with hyper-parameters and hardware specs (Appendix C An anonymous repository with code and data is referenced in the appendix.
5.3
USAGE OF LLM
LLMs were not used to generate data, labels, or experimental results. The core methodology re- lies pretrained molecular/protein encoders and an optimal-transport pseudo-labeling framework integrated with a biomedical KG. LLMs were used only as general-purposs assist tools for light copy-editing and minor code polishing; all technical ideas, algorithms, experiments, and analyses were conceived and implemented by the authors, who take full responsibility for accuracy and in- tegrity.
10

---

Published as a conference paper at ICLR 2026
REFERENCES
Mohammed AlQuraishi. Proteinnet: a standardized data set for machine learning of protein struc- ture. BMC bioinformatics, 20:1-10, 2019.
Michael Ashburner, Catherine A Ball, Judith Blake, David Botstein, Heather Butler, J Michael Cherry, Allan P Davis, Kara Dolinski, Selina Dwight, Janan T Eppig, al. Gene ontology: tool for the unification of biology. Nature genetics, 25(1):25- 2000.
Amos Bairoch. The enzyme database in 2000. Nucleic acids research, 28(1)3304-305, 2000.
Pedro J Ballester and John BO Mitchell. A machine learning approach to predicting protein ligaad binding affinity with applications to molecular docking. Bioinformatics, 26(9): 169-11967770 1175, 2010.
Michael Brocidiacono, Paul Francoeur, Rishal Aggarwal, Konstantin I Popov, David Ryan Koes, and Alexander Tropsha. Bigbind: leamiing from nonstructural data for structure- based virtual screening. Journal of Chemical Information and Modeling, 6A(7):2488-2495, 2023.
Payal Chandak, Kexin Huang, and Marinka Zitnik. Building knowledge graph enable precision medicine. Nature Scientific Data, 2023, doi: https:///oiiiooo http:///000..o0//00000001/110000
URL
nature https: http:///wwwnaauuee.m cww taraiiilleee 841 197 023
The UniProt Consortium.
Uniprot: the universal protein knowledgebase in 2025. Nucleic Acids Research, 530D1):609--D617, 11 2024. ISSN 0305-1048. doi: 10. (093/nar//kaee010 URL
Kirill Degtyarenko, Paula de Matos, Marcus Ennis, Janna Hastings, Martin Zbinden, Alan Mc- Naught, Rafael Alcántara, Michael Darsow, Mickaël Guedj, and Michael Ashburner. Chebi: database and ontology for chemical entities of biological interest. Nucleic Acids Research, 36
(1:pl_1):D334 D350, 10 2007. ISSN 0305-1048. doi: 10.1093/nar/gkm771. URL https: //dei. 10. 1093 nar /g
Alfredo Di Nola, Danilo Roccatano, and Herman JC Berendsen. Molecular dynamics simulation of the docking of substrates to proteins. Proteins: Structure, Function, and Bioinformatics, 19(3):
174-182, 1994.
Jacob D. Durrant and Andrew McCammon. Nnscore 2.0: A neural-network receptor receptor-ligand scoring function. Journal of Chemical Information and Modeling, 51 ():287)-2993 2011. doi: 10.1021/ci2003889 URL https: org/10.../...oo...ooooooo.......
htt.:t/s//......rg..0000///////0 1021/c12003889
PMID: 22017367.
Qizhe Fan, Xiaoqin Shen, Shihui Ying, and Shaoyi Du. Otclda: Optimal transport and contrastive learning for domain adaptive semantic segmentation. IEEE Transactions on Intelligent Trans- portation Systems, 2024.
Yin Fang, Qiang Zhang, Ningyu Zhang, Zhuo Chen, Xiang Zhuang, Xin Shao, Xiaohui Fan, and Huajun Chen. Knowledge graph-enhanced molecular contrastive learning with functional prompt. Nature Machine Intelligence, ((5::54-5533, 2023.
Robert D Finn, Alex Bateman, Jody Clements, Penelope Coggill, Ruth Eberhardt, Sean R Eddy, Andreas Heger, Kirstie Hetherington, Liisa Holm, Jaina Mistry, et al. Pfam: the protein families database. Nucleic acids research, 42(D1):D222 22212222222,2222222222222222,,,2 2014.
Bowen Gao, Bo Qiang. Haichuan Tan, Yinjun Jia Minsi Ren, Minsi Lu, Jingjing Liu, Wei-Ying Ma, and Yanyan Lan. Drugelip: Contrastive protein-molecule rep- resentation learning for virtual screening. In A. Oh, T. Naumann, A. Glober- son, K. Saenko. M. Hardt, and S. Levine, editors, Advances in Neural Informa- tion Processing Systems, volume 36. pages 44595 44614. Curran Associates, Inc.,
www.inng...uurii.cccmm wwtpsps.ccpppppppppe. files//aper/2022 2023. URL Balaad8B9a31133 1ee.55999feeeee f www.eeeeeeee774ddee oon File/Bbd3 288ad88
Anna Gaulton, Louisa J Bellis. A Patricia Bento, Jon Chambers, Mark Davies, Anne Hersey, Yvonne Light, Shaum McGlinchey, David Michalovich. Bissan Al-Lazikani, et al. Chembl: a large-scale bioactivity database for drug discovery. Nucleic acids research, 40(D1):D1 100-D1107, 2012.
11

---

Published as a conference paper at ICLR 2026
Federico Gossi, Pushpak Pati, Panagiotis Chouvardas, Adriano Luca Martinelli, Marianaa Kruithof- de Julio, and Maria Anna Rapsomaniki. Matching single cells across modalities with contrastive learning and optimal transport. Briefings in bioinformatics, 24(3):bbad130) 2023.
Siyuan Guo, Lexuan Wang, Chang Jin, Jinxian Wang, Han Peng, Huayang Shi, Wengen Li, Jihong Guan, and Shuigeng Zhou. M9-20m: A large-scale multi-modal molecule dataset for ai-driven drug design and discovery. Journal of Bioinformatics and Computational Biology, 2025, doi:
10.1142750211722025500000 URL https: www wor www.worrddcieetifii...c. wtt.tti....coo...oommmmmmo.oo.tt doi/10. 1142/802197202250005
Thomas A Halgren, Robert B Murphy, Richard A Friesner, Hege S Beard, Leah I. Frye, W Thomas Pollard, and Jay Banks. Glide: a new approach for rapid, accurate docking and scoring. 2. en- richment factors in database screening. Journal of medicinal chemistry, 47(7):1750-1755 2004.
Ajay Jain. Surflex: fully automatic flexible molecular docking using molecular similarity- search engine. Journal of medicinal chemistry, 46(4):499-51 511.99951111111 2003. 2003.
Yinjun Jia, Bowen Gao, Jiaxin Tan, Jiqing Zheng, Xin Hong, Wenyu Zhu, Haichuan Tan, Yuan Xiao, Liping Tan, Hongyi Cai, Yanwen Huang, Zhiheng Deng, Xiangwei Wu, Yue Jin, Yafei Yuan, Jiekang Tian, Wei He, Weiying Ma, Yaqin Zhang, Wei Zhang, Lei Liu, Chuangye Yan, and Yanyan Lan. Deep contrastive learning enables genome-wide virtual screening. Science, 2026. doi: 10.1 hetp////ccccccccs5555500 URL/tttss: http:////wwwwssienee.oooo wtt.://oc...rg/////0.1..1.1..... aalence.ads995
John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Židek, Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. nature, 596(7873):583-589 2021.
Minoru Kanehisa and Susumu Goto. Kegg: kyoto encyclopedia of genes and genomes. Nucleic
acids research, 28(1):27-30, 2000.
Craig Knox, Mike Wilson, Christen M Klinger, Mark Franklin, Eponine Oler, Alex Wilson, Al- lison Pon, Jordan Cox, Na Eun Chin, Seth A Strawbridge, al. Drugbank 6.0: the drugbank knowledgebase for 2024. Nucleic acids research, 52(D1):D126S-DI77) 2024.
Seonghyeon Lee, Dongha Lee, Scongbo Jang, and Hwanjo Yu. Toward interpretable semantic tex- tual similarity via optimal transport- contrastive sentence learning. Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Pa- pers), pages
Dublin, Ireland, may 2022. Association for Computational Linguis- doi: 10. (0.116653/1122022.cd// acl 012.4127
tics.
URL https: / /aclaalaa anthelooy..rrg2222.
Xuan Lin, Zhe Quan, Zhi-Jie Wang, Tengfei Ma, and Xiangxiang Zeng. Kgnn: Knowledge graph neural network for drug-drug interaction prediction. In IJCAI, volume 380, pages 2739-2745, 2020.
Tengfei Ma, Xuan Lin, Bosheng Song, Philip S. Yu, and Xiangxiang Zeng. Kg-mtl: Knowledge graph enhanced multi-task learning for molecular interaction. IEEE Transactions Knowledge and Data Engineering, 35:7068-7082 2022. doi: 10. 1099//DD..002..33 144.
Tengfei Ma, Yujie Chen, Wen Tao, Dashun Zheng, Xuan Lin, Patrick Cheong-Iao Pang. Yiping Liu, Yijun Wang, Longyue Wang. Bosheng Song. Xiangxiang Zeng. and Philip S. Yu. Learning to de- noise biomedical knowledge graph for robust molecular interaction prediction. IEEE Transactions on Knowledge and Data Engineering, 36:8682- 8694, 2024. doi: 10. 1..10099KDDEE.22444477
Andrew T McNutt, Paul Francoeur, Rishal Aggarwal, Tomohide Masuda, Rocco Meli, Matthew Ragoza, Jocelyn Sunseri, and David Ryan Koes. Gnina 1.0: molecular docking with deep learn- ing. Journal of cheminformatics, 13(1):43, 2021.
Christian von Mering, Martijn Huynen, Daniel Jaeggi. Steffen Schmidt, Peer Bork, and Berend Snel. String: a database of predicted functional associations between proteins. Nucleic acids research, (111:258-261. 2003.
12

---

Published as a conference paper at ICLR 2026
Michael M. Mysinger, Michael Carchia, John. J. Irwin, and Brian K. Shoichet. Directory of use- ful decoys, enhanced (dud-e): e): Better ligands and decoys for better benchmarking. Journal of Medicinal Chemistry, 55(14):6582-6333 6594, 2012. doi: 10.1021/jm300877. PMID: 227 22716043.
Maho Nakata, Tomomi Shimazaki, Masatomo Hashimoto, and Toshiyuki Maeda. Pubchemqe pm6:
Data sets of 221 million molecules with optimized molecular geometries and electronic properties. Journal of Chemical Information and Modeling. 60(12):5891-589) 2020.
Hakime Özüürk, Arzucan Özgür, and Elif Ozkirimli. Deepdta: deep drug-target binding affinity prediction. Bioinformatics, 34(17):1221ii82) 2018.
Jiayu Qin, Jian Chen, Rohan Sharma, Jingchen Sun, and Changyou Chen. A probability contrastive learning framework for 3d molecular representation learning. In A. Globerson, L. Mackey, D. Bel- grave, A. Fan, U. Paquet, J. Tomezak, and C. Zhang, editors, Advances in Neural Information Pro- cessing Systems, volume 37, pages 58058-58076. Curran Associates, Inc., 2024. doi: 10.52202/ httpsperrr_eeeeppppeeeeeee
079017-1851. LLLL:aates//pprcedddgggeneoooo..
Alexander Rives, Joshua Meier, Tom Sercu, Siddharth Goyal, Zeming Lin, Jason Liu, Demi Guo, Myle Ott, C. Lawrence Zitnick, Jerry Ma. and Rob Fergus. Biological structure and function emerge from scaling unsupervised learning to 250 million protein sequences. PNAS, 2019. doi: 10.1101/622803, URLhttps: /tttpps::/www.... www.iiv...vv.oommmmoom... org http://www..aaaaroooooo..tttt./. wtt.:/ttetn/n...11/222222222
Liangliang Shi, Gu Zhang, Haoyu Zhen, Jintao Fan, and Junchi Yan. Understanding and generalizing contrastive learning from the inverse optimal transport perspective. In International conference on machine learning, pages 31408-31421. PMLR, 2023.
Richard Sinkhorn and Paul Knopp. Concerning nonnegative matrices and doubly stochastic matri- ces. Pacific Journal of Mathematics, 21(2):343-348, 1967.
Tengwei Song, Long Yin, Zhen Han, and Zhiqiang Xu. Improving enzyme prediction with chemical reaction equations by hypergraph-enhanced knowledge graph embeddings. In Proceedings of the 32st ACM SIGKDD Conference on Knowledge Discovery and Data Mining, KDD 2026, Jeju Island, Republic of Korea, August 9-13, 2026. ACM, 2026. doi: 10. 10.114553770854.3888777
Marta M. Ceepneeska--ziiubing
Piotr Zielenkiewiez, and Pawel Siedlecki. Development and evaluation of deep learning model for protein- binding affinity prediction. Bioinformatics, 34(21):3666 3674, 2018.
doi: 10. 1001000/bbommmmmmcccc/tty..34444 URL /ttttt:p///aaeeettt/n/
Oleg Trott and Arthur J Olson. Autodock vina: improving the speed and accuracy of docking with a new scoring function, efficient optimization, and multithreading. Journal of computational chemistry, (12)):5554442 2010.
The UniProt Consortium. Uniprot: the universal protein knowledgebase. Nucleic acids research, 46 (5):2699-2699, 2018.
Joana Vilela, Muhammad Asif, Ana Rita Marques, Joao Xavier Santos, Célia Rasga, Astrid Vicente, and Hugo Martiniano. Biomedical knowledge graph embeddings for personalized medicine: Pre- dicting discase-gene associations. Expert Systems, 40(5) 18111, 2023.
Debby D Wang, Mengxu Zhu, and Hong Yan. Computationally predicting binding affinity in protein-ligand complexes: free energy-based simulations and machine learning-based scor- ing functions. Briefings in Bioinformatics, 22(3):bbaa107, 06 2020. ISSN 1477-4054. doi:
10.1093/bib/bbaa107 UtLhttps::/ org/1...rgg.... wtt.://///....//.i..////./ 1t993//bbbbbba1001 bbbal007
1093/
Zifeng Wang. Zichen Wang, Balasubramaniam Srinivasan, Vassilis N. Ioannidis, Huzefa Rang- wala, and Rishita Anubhai. Biobridge: Bridging biomedical foundation models via knowl- edge graphs. In International Conference on Learning Representations (ICIR), 2024. URL. Mtttps:/cppezzett
net wwr.moormmmmmm JJdMMiWWeMMiH
Ying Xia, Xiaoyong Pan, and Hong-Bin Bin Shen. comprehensive survey on protein ligand binding site prediction. Current Opinion in Structural Biology, 86: 16:102,,3, 2024.
13

---

Published as a conference paper at ICLR 2026
Xiangying Zhang, Haotian Gao, Haojie Wang, Zhihang Chen, Zhe Zhang, Xinchong Chen, Yan Li, Yifei Qi, and Renxiao Wang. Planet: multi- multi-objectiiv graph neural network model for protein- ligand binding affinity prediction. Journal of Chemical Information and Modeling. 64(7):2205- 2220, 2023.
Liangzhen Zheng, Jingrong Fan, and Yuguang Mu. Onionnet: a multiple-layer intermolecular- contact-based convolutional neural network for protein ligand binding affinity prediction. ACS omega, 1(14): 2019.
Gengmo Zhou, Zhifeng Gao, Qiankun Ding, Hang Zheng, Hongteng Xu, Zhewei Wei, Linfeng Zhang, and Guolin Ke. Uni-mol: universal 3d molecular representation learning frame- work. In International Conference on Learning Representations (ICLR), 2023. URL https: //openreview. 60
14

---

Published as a conference paper at ICLR 2026
A RELATED LITERATURE
A.1 OPTIMAL TRANSPORT IN REPRESENTATION LEARNING
Optimal Transport (OT) is a mathematical framework originally developed to solve resource al- location problems by minimizing the cost of transporting mass from one distribution another. Recently, it has been increasingly applied in machine learning and representation learning tasks due its ability to compare distributions in a geometrically meaningful way. Unlike traditional dis- lance metrics such as Euclidean or cosine distances, OT considers the geometry of the distributions, enabling it to capture grainge relationships between data points.
In the context of multimodal representation learning, OT has been used to align embeddings from different modalities by computing an optimal coupling between them. This approach ensures that structurally similar elements across modalities are closely matched, thus enhancing the quality of learned representations. Applications of OT have been particularly impactful in cross-modal tasks, such as image-text retrieval, molecular-prose prroteiin interaction modeling, and domain adaptation. No- table works include the Sinkhorn-Knopp algorithn and Knopp (1967). which makes OT computationally efficient for lcage-scale datasets by introducing entropy regularization to the trans- port problem. Shi et al.] (2023) aim to understand and generalize contrastive learning as a form of inverse optimal transport. The paper also shows that InfoNCE contrastive loss is a specific case of the proposed IOT loss. Other notable works in applying Optimal Transportation in feature learning
include Fan al. (2024), Lee al. (2022) al. (2024) and Gossi al. (2023)
A.2 KNOWLEDGE GRAPHS IN MULTI-MODALITY MOLECULAR AND PROTEIN TASKS
Knowledge graphs (KGs) have emerged as powerful tools for integrating and modeling complex, heterogeneous data in multi-modality tasks, including molecular and protein-related studies. A KG represents entities (e.g.. molecules, proteins, and biological pathways) as nodes and their relation- ships (e.g., binding, inhibition, interaction) as edges, enabling the incorporation of prior biological knowledge into machine learning models.
In molecular and protein studies, KGs such as DrugBankKnox al. (2024). ChEMBI Gaulton al. (2122,, and STRINGMering et al.) (2003) have been widely used to capture pileculla--roteinn in- teractions and other biological relationships. By encoding domain-specific knowledge into graph structures, KGs enhance downstream tasks like molecular property predictionFang et al. (2023), drug-drug interaction predictior[Lin et al. (2020), and disease-gene association studics Vilela et al. 2023).
Recent advancements have also explored cross-modality tasks involving molecules and proteins by leveraging KGs common representation space. For example, BioBridge Wang et al. (2024) aligns molecule and protein embeddings through a knowledge graph and evaluates cross- cross-domain retrieval tasks. [Song al. 2026) is another application of knowledge graphs in biological model- ing. Similarly, graph neural networks (GNNs) have been applied to propagate information across graph nodes, enabling the extraction of contextual embeddings that incorporate relational knowl- edge. Despite these successes, integrating KGs with multi-modality data remains challenging due the heterogeneous nature of molecular and protein features, as well as the sparsity of some entity relationships.
Positioning with KG-based MPI frameworks. Several recent KG-based approaches are closely related our setting but differ their supervision and how they exploit KG structure. (Ma al.) 2022) propose KG-MTL, which jointly leverages biomedical KG (for drug entities) and molecu- lar graphs (for compound structures) under a multi-tast learning objective for DTT/CPI prediction. (Ma et al. 2024) propose BioKDN, a robustification module that denoises local KG neighborhoods and smooths relation semantics to mitigate noisy edges when predicting molecular interactions. In contrast, KGOT is designed for unlabeled weakly-labeled MPI retrieval: we (i) construct uni- fied, multimodal KG spanning compounds, proteins, pathways, GO terms, and protein families, (ii) generate global pseudo-labels via entropy-regulariieed optimal transport that enforces detaset-level matching constraints, and (iii write back siggh--ofidencc pseudo interactions into the KG to enable downstream KG completion.
15

---

Published as a conference paper at ICLR 2026
Method
Unified multimodal KG
KG denoising Multi-task learning
Global OT
KG-MTL Ma et al. 2022
BioKDN Ma al. KGOT (ours)
2024
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 KG denoising.
B DETAILS OF DATASET CONSTRUCTION
To develop comprehensive knowledge graph to study molecule and protein interactions, sidered 6 primary resources of biological and clinical information. The data resources provide widespread coverage of biomedical entities, including proteins, genes, drugs, molecules, biologi- cal processes, cellular components and protein families. These were high-quality datasets, either expertly curated annotations such as KEGG, widely-used standardized ontologies such the Gene Ontology, or direct readouts of existing large scale unimodality dataset, such as CHEBI and UniProt.
(I) 1.35 million potential catalytic activity relationships between over 30,000 molecules from CHEB Degtyyrngya et al. (2007) and 500,000 proteins from UniPro[UniProt Consortium (2018).
(2) Genetic and genomic information from KEGGKanchisa and Goto (2000). capturing how molec- ular interactions are influenced by gene regulation and metabolic pathways.
(3) Functional annotations from Gene Ontology Ashburner (2000), enriching molecular repre- sentations with biological process and cellular component insights.
(4) Protein family classifications from PFaMFinn al. (2014), improving interaction predictions by incorporating shared functional domains.
(5) Enzyme Commission numbers from ENZYMEBairoch (000), allowing for enzyme-substrate relationship modeling within our graph.
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
KMODUL
1,275
KPATHWAY
89,989
KCOMPOUND 16,695
KMODULE 3,470
Table 5: Entities Distribution, where KMODUL represents KEGG_MODUL KPATHWAY rep- resents KEGG PATHWAY, KCOMPOUND represents KEGG COMPOUND, and KMODULE represnts KEGG MODULE.
Our knowledge graph dataset a multimodal knowledge graph with types of nodes, 29 types of directed edges, 6,483,852 relationships between entities.
C IMPLEMENTATION DETAILS
Backbone architecture. We use Uni-Mol Zhou al. (2023) for both molecular and protein coders. Hidden dimension is set to 512 for both. The scoring MLP has two layers size [512, 256, 1]. with ReLU activations.
16

---

Published as a conference paper at ICLR 2026
Labeled dataset training. Training the labeled set uses a batch size of 128, learning rate le-4, and Adam optimizer with weight decay 0.01. The score function S(x, y) trained for 50 epochs with early stopping.
Pseudo-label generation. For full-batch OT, we construct score matrix for 10k molecules 5k proteins. Sinkhorn 0.01, similarity weight À 0.1, learning rate 7j 1.0, and 50 iterations. Top-k baseline selects the top 5 pseudo-labels per protein.
Knowledge graph training. KG embeddings are trained with embedding size 256 and margin 6.0. We use 1:1:1 sampling for real triples and batch size of 1024. Pseudo-interactions are weighted via the loss term ppeudo-
Hardware. All experiments are run on 4 NVIDIA A6000 48GB GPUs with 256GB RAM. Total training time for each variant is under 12 hours.
D TEAKAGE-CONTRO) EXPERIMENTS
Protocol. To further mitigate train test leakage beyond the default setup in the main text, we evaluate our model under a strict protocol that combines molecule- and protein-side filtering: (i) on the molecule side, we adopt Murek--xcafoold-ou setting where training ligands sharing the
scaffold with any test active removed; (ii) the protein side, we adopt family-out setting by removing from the training pool all proteins that map (via HMMER to Pfam- A) to any family present among test targets. All other training, inference, and evaluation details (single model, zero-shot evaluation, fixed seeds) follow the main-text configuration.
Results on DUD-E (strict leakage control). Table reports performance under the combined Nlurck--caafold-ort (ligands) Pfam family-out (proteins) protocol. As expected, absolute num- bers are lower than in Table [ yet early recognition remains strong under strict filtering.
Table 6: DUD-E under strict leakage control.
Model
AUROC (%)
BEDROC (%)
EF@0.5%
EF@ 1%
EF@29
KGOT (strict)
81.78
51.04
38.91
32.47
10.35
Table summarizes performance LIT--PCA under the same strict pro- tocol. Given the greater class imbalance and challenge level of LIT-PCBA, we continue to emphasize early enrichment metries.
Results LIT-PCBA.
Table 7: LIT-PCBA under strict leakage control
Model
AUROC (%)
BEDROC (%)
EF@0.5%
EF@ 1%
EF@ 55%55%
KGOT (strict)
61.22
5.96
8.92
5.34
2,27
E ABLATION EXPERIMENTS
E. OT LOSS VS. CONTRASTIVE Loss (INFONCE) IN SUPERVISED TRAINING
We first compare the proposed optimal transport loss to standard contrastive loss (InfoNCE) for supervising the scoring model. This experiment evaluates whether our OT-based loss offers an advantage over more conventional pairwise contrastive approach.
Experimental Setup:
We train the scoring model (MuRE-base encoder) on known entity pairs using either (i) OT loss, (ii) InfoNCE loss. For InfoNCF, each positive pair is contrasted against
17

---

Published as a conference paper at ICLR 2026
multiple randomly sampled negative pairs (with temperature tuned to 0.1). All other training settings (learning rate, epochs, etc.) are kept identical. We evaluate the models on the link prediction task.
Following Table shows the performance comparison. The model trained with the OT loss achieves slightly higher accuracy than with InfoNCE. For instance, OT loss yields an MRR 0.256 vs. 0.243 with InfoNCE, and @its@ 10 of 45.1 vs. 42.7.
Training Loss
MRR
Hits 10
OT Loss (ours)
0.256
45.1%
InfoNCE Loss
0.243
42.7%
Table 8: Performance of the scoring model with OT-based loss vs. contrastive InfoNCE loss. OT loss yields modest but consistent improvement in link prediction metrics.
The OT-trained model outperforms the InfoNCE variant, indicating that the OT formulation provides beneficial supervisory signal. We attribute this gain the OT loss's ability consider the full distribution of true associations for each entity, rather than only one positive against negatives at time. In knowledge graph link prediction, an entity can have multiple correct targets; the OT loss naturally accommodates multiple positive matches by optimizing a transport plan over them.
E.2 PSEUDO-LABELING STRATEGIES
We compare pseudo labels, random selection, top-k per protein, OT without similarity, high- entropy OT, and our full simiiilrar
OT yields balanced pseudo labels with broader protein/molecule coverage than top-k, while the similarity prior filters implausible assignments. Excessive entropy (e large) makes labels diffuse and less informative.
E.3 COMPONENTS ABLATION
To further analyze the contributions of different components, we perform an ablation on the knowl- edge graph variant of model. Table [0]sshww an ablation using the TorusE embedding method on the link prediction task, where we incrementally add sources information. Starting with only molecule-protein edges (using the small labeled set, akin baseline without any additional knowl- edge), we then add gene ontology (GO) relations, protein family relations, and metabolic process (pathway) relations from the knowledge graph. Each addition yields an improvement in Hits@ 1. Specifically, incorporating GO relations (which provide gene-function context) boosts Hits from 36.3% to 43.7%; adding protein family info further raises it 47.4%; and including metabolic pathways brings it to 53.4%. This ablation highlights that each modality of biological knowledge contributes to better performance. It underscores the importance of multimodal data integration: the model with full knowledge (last row) performs significantly better than using only the labeled molecule-protel pairs.
F MULTI-OBIECTIVE LEARNING FRAMEWORK
The framework is trained using a multi-objective loss function that integrates information from both the pseudo label matrix and the KG structure. The loss components are as follows:
Graph embedding loss A knowledge graph embedding model is trained to learn representations of entities and relations in the KG. The graph embedding loss defined as:
LKG
Σ
logo (f(,
Σ
logo (ogo(-f(h'r,tt))
(12)
(h.r:))DRa
(M (.
where f(h, r. t) is the scoring function for a triple (h, r. t), a is the sigmoid function, and Dxc and DiG are the sets of positive and negative samples, respectively.
18

---

Published as a conference paper at ICLR 2026
Table Ablation on pseudo-label generation (illustrative Hits@ on link prediction).
Strategy
Hits @ (%)
No pseudo-lahel augmentation
68.5
Random pseudo-labelt (matched count)
66.0
Top-k per protein (A: 5)
72.0
OT without similarity (A 0)
73.8
OT with high entropy (E = 0.1)
73.0
OT similarity (A = 0.11, 0.01)
74.9
Table 10: Ablation study on effect of integrating different knowledge graph relation types (using TorusE).
Model Configuration
Hits @ (Link Prediction)
TorusE (molecules & proteins only)
36.3%
Gene Ontology relations
43.7%
Protein family relations (GO + PF)
47.4%
Metabolic process relations (GO + PF + KEGG + EC)
53.4%
The pscudo labels T represent predicted scores for protein- moleeule interactions. We incorporate these into the model by defining loss term that aligns the KG- predicted scores with T:
Pseudo label alignment loss
M N ΣΣ (fees peeundo Eÿ) Tij)²
Lpouedo
(13)
=1
where e, and E, are the embeddings of molecule and protein j. respectively, and rpaeudo is the learned embedding for the sseudo-interacee
relation.
The total loss function for training the model is weighted sum of the above components:
Lital = Lxg + a.-eeuulo
(14)
where a is the hyperparameter that balance the contributions of the pseudo label and similarity terms, we take = 0.1 in the experiments.
19


Reviews:
[{"review_id": "1zbpHIpANG", "reviewer": "Reviewer_Css5", "summary": "The authors propose a framework to predict molecule–protein interactions. It generates high-quality pseudo-labels to leverage diverse biological modalities.", "questions": "Could the authors compare the performance and clarify the advantages over relevant KG-based methods such as KG-MTL and BioKDN?", "limitations": "", "rating": 2, "confidence": 5, "soundness": 1, "presentation": 1, "contribution": 1, "strengths": "The paper proposes a unified framework that aims to integrate biological entities such as pathways and genes.", "weaknesses": "(1) The proposed framework mainly combines existing methods without introducing any new modeling components or theoretical insights.\n\n(2) The performance is not compared with highly relevant works such as KG-MTL and BioKDN.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "hp7xT98vYw", "reviewer": "Reviewer_YDJx", "summary": "This paper proposed a unified framework KGOT that integrates knowledge graph and optimal transport to generate high-quality pseudo-labels for unlabeled molecule-protein pairs in molecule-protein interactions (MPIs) prediction tasks.", "questions": "- What is the specific cost function used in the Optimal Transport plan? Is it based on the model’s predicted scores, molecular/protein embeddings, or a combination?\n- How is the marginal distribution for the OT problem defined? Are uniform distributions assumed for molecules and proteins, or is there a prior (e.g., based on node degree in the KG) that biases the label assignment?\n- How is the entropy regularization parameter chosen and tuned? This parameter critically balances between fitting the data and the smoothness of the transport plan, directly impacting pseudo-label quality.\n- What is the architecture of the initial scoring model (Step 2)? Is it a simple neural network, or does it already incorporate some KG information?\n- How is the knowledge graph itself encoded in Step 4? Are you using TransE, ComplEx, a Graph Neural Network (GNN), or another Knowledge Graph Embedding (KGE) method? The choice here significantly affects the model’s ability to capture complex relational paths.\n- The mutual retrieval objective suggests a dual-encoder architecture. How are the molecule and protein encoders designed and aligned? Is the contrastive loss used, and if so, what is the strategy for mining hard negatives?\n- How do you prevent data leakage between the small labeled dataset (Step 2) and the large unlabeled dataset (Step 3)? If proteins/molecules from the labeled set appear in the unlabeled KG, it could artificially inflate performance.\n- What is the criteria for a high-quality pseudo-label? Is there a threshold on the OT-assigned probability, and how is this threshold determined?\n- Is the proposed KGOT process iterative? In other words, is the KG-augmented model (Step 4) used to re-score pairs and generate new, improved pseudo-labels in a self-training loop?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "S1: This paper is to leverage large-scale multimodal knowledge graphs and propose an optimal transport-based pseudo-labeling strategy for the MPIs prediction.\n\nS2: In experiment part, the proposed KGOT outperforms existing MPIs prediction methods in terms of AUROC, early recognition metrics, and generalization to unseen interactions.", "weaknesses": "W1: A primary weakness of KGOT is its reliance on the critical yet potentially unstable step of pseudo-label generation via optimal transport. The quality of the entire approach hinges on the assumption that the optimal transport mechanism can accurately infer the underlying distribution of known interactions to assign reliable labels to unknown molecule-protein pairs.\n\nW2: Aggregating molecular, protein, gene, and pathway-level information requires sophisticated fusion techniques to handle the disparate scales, formats, and sparsity levels of each modality. However, the integration of diverse biological data introduces significant complexity and potential challenges in data harmonization and model interpretation.\n\nW3: Without validation on truly de novo targets or wet-lab confirmation, the practical utility and robustness of the framework for unseen interactions, iterative drug discovery pipeline remain unproven.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "K36uYiAAh8", "reviewer": "Reviewer_ZaTk", "summary": "This study presents a new method for predicting protein-molecule interactions. The method relies on extracting the knowledge from a multimodal knowledge graph. The construction and composition of the multimodal knowledge graph is not novel. The two main novel contributions are: using optimal transport-based pseudo-labeling strategy to leverage a large unlabelled dataset and augment the original knowledge graph. The method shows robust improvement over the SOTA.", "questions": "1. Why did you not use a pre-made KG like PrimeKG, for example?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The optimal transport-pseudo labelling strategy seems a good contribution to the overall field and opens the door to creating more sophisticated augmented knowledge graphs.\n2. The results are shown with some measure of the performance dispersion, it is unclear which or where it is derived from, but it allows for some determination of the statistical significance of the results.\n3. The ablation study is comprehensive and convincingly demonstrates that the different components of the method improve prediction accuracy.", "weaknesses": "1. It is unclear what the error represents in Tables 1 and 2.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "yv5H10oN52", "reviewer": "Reviewer_Tsmw", "summary": "In the paper, the authors studied methods for predicting ligand-target binding affinity as link prediction problems. In order to resolve the issue of lacking large-scale training data. The authors proposed methods for augmentation of labelled data with useful information from multi-modal knowledge graphs. \n\nThe proposed method following 4 different steps:\n\n+ use a small labelled datasets to train a predictive model that predicts the probability of binding between any pair of molecule and protein.\n\n+ use the given model to create a pseudo-label data on a larger unlabelled graph. The pseudo labels are enforced to give overall consistency prediction on the entire unlabelled graph by Optimal Transport and the regularization of the similarity between nodes on the latent spaces.\n\n+ the pseudo-labelled graph together augmented with the edges of the large multimodal knowledge graph is then used to train a final link prediction model\n\nIn order to evaluate the proposed approach the author compared there KGOT method with baseline methods that do not use external information on the DUDe and  LIT-PCBA with leakage removal from the training graphs. Their proposed approach yields better results then the baseline methods.\n\nIn another experiments for link prediction, KGOT was used in addition to other graph embedding methods. the experimental results show that the data augmentation procedure in KGOT via OT helps simple graph embedding methods improve linking prediction accuracy.", "questions": "Could you please compare your approach with the given related works listed in the weakness section?\n\n Could you please do more experiments on the datasets listed in the weakness section?", "limitations": "", "rating": 4, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "An interesting idea was proposed regarding using OT to control the consistency of pseudo-labels overall.\n\nThe paper was well written and easy to follow.", "weaknesses": "The idea of using multimodal knowledge graphs to augment training data for ligand-target affinity prediction is not a new idea. The given idea has been proposed and studied in the following works:\n\n+ N. Zhang, Z. Bi, X. Liang, S. Cheng, H. Hong, S. Deng, Q. Zhang, J. Lian, and H. Chen. Ontoprotein: Protein\npretraining with gene ontology embedding. In International Conference on Learning Representations, 2022.\n\n+ H.-Y. Zhou, Y. Fu, Z. Zhang, B. Cheng, and Y. Yu. Protein representation learning via knowledge enhanced\nprimary structure reasoning. In The Eleventh International Conference on Learning Representations, 2023.\n\n+ Lam H. T., Sbodio M. L., Martínez Galindo M., et al. “Otter-Knowledge: benchmarks of multimodal knowledge graph representation learning from different sources for drug discovery.” (2023).\n\n+ Ye Q., Hsieh C-Y., Yang Z., et al. “A unified drug–target interaction prediction framework based on knowledge graph and recommendation system.” Nature Communications 12:1 (2021).\n\nI think the authors have not discussed those related work carefully and compare to those approaches.\n\n\nThe results on DUDe and LIT-PCBA are interesting but I think the authors should compare to the above approaches on the following standard benchmarks:\n\n+ TDC DTI https://tdcommons.ai/benchmark/dti_dg_group/overview/\n\n+ DAVIS\n\n+ KIBA", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)