Multimodal out-of-distribution individual uncertainty quantification enhances binding affinity prediction for polypharmacology

Abstract

Polypharmacology, a single drug that targets multiple proteins, holds promise for addressing unmet medical needs. Achieving accurate, reliable and scalable predictions of protein–ligand binding affinity across multiple proteins is crucial to realizing the potential of polypharmacology. Machine learning offers a powerful tool for multitarget binding affinity prediction. However, three major challenges remain: generalizing predictions to out-of-distribution compounds that are structurally different from those in the training data; quantifying the uncertainty of predictions in out-of-distribution scenarios where the assumption underlying existing methods does not hold; and scaling to billions of compounds, which remains unattainable for current structure-based methods. Here, to overcome these challenges, we propose a model-agnostic anomaly detection-based individual uncertainty quantification method: embedding Mahalanobis Outlier Scoring and Anomaly Identification via Clustering (eMOSAIC). eMOSAIC features the divergence between the multimodal representations of known cases and unseen instances and quantifies individual prediction uncertainty on a compound-by-compound basis. We integrate eMOSAIC with a multimodal deep neural network for multitarget ligand binding affinity predictions, leveraging a structure-informed large protein language model. Comprehensive validation in out-of-distribution settings demonstrates that eMOSAIC significantly outperforms state-of-the-art sequence-based and structure-based methods as well as existing uncertainty quantification approaches. These findings underscore eMOSAIC’s potential to advance real-world polypharmacology and other applications that require robust predictions and scalable solutions.

Drug discovery is a highly complex process that takes up to 15 years and costs billions of dollars¹, but its failure rate is extremely high due to a lack of clinical efficacy or safety². Many complex diseases such as neurological and mental disorders are multigenic, multifactorial diseases. Thus, an effective therapy needs to modulate multiple genes^3,4. Polypharmacology that uses a single chemical to selectively bind to multiple drug targets has emerged as a new paradigm in drug discovery^5–10. Furthermore, a drug molecule rarely interacts exclusively with its intended target. Off-target binding is common and can lead to undesirable side effects and drug–drug interactions¹⁰. Thus, elucidating selective ligand binding profiles on a genome-wide scale is essential for developing effective and safe polypharmacology therapeutics.

Screening a library of billions of compounds against multiple drug targets to identify hit compounds and subsequently optimizing their binding affinity selectivity profiles for drug leads are critical steps in polypharmacology¹¹. A gigascale compound screening can increase the likelihood of identifying more potent or selective ligands^12–15, streamline lead optimization and expand the chemical diversity, chemical novelty and patentability of drug leads¹⁶. Although DNA-encoded libraries are an effective approach to generate and screen billions of compounds for a single target¹⁷, they have limited chemistries and a high rate of false positives. The computational approaches are expected to facilitate large-scale polypharmacology compound screening¹¹.

Protein–ligand docking is commonly used for compound screening¹⁸. However, protein–ligand docking suffers from a high rate of false positives due to poor modelling of protein dynamics, solvation effects and other challenges¹⁹. The reliability of protein–ligand docking substantially deteriorates when using predicted structures^20,21 and is not reliable for virtual screening²². Despite the success of AlphaFold2²³, it can reliably model only half of understudied human proteins whose small-molecule ligands are unknown²⁴. Thus, the application of protein–ligand docking in polypharmacology is hindered by the frequent lack of critical structural information, such as the experimentally determined holo structures and binding sites of many target proteins. Moreover, protein–ligand docking is computationally demanding. As a result, it is impractical to use this approach to screen billions of compounds across multiple targets in the context of polypharmacology.

Recent advances in artificial intelligence encourage increasing interest in applying deep learning to drug discovery^25,26. Sequence-based deep learning offers notable advantages. Sequence-based drug–target interaction prediction, which uses only three-dimensional (3D) structure information implicitly, enables fast drug–target interaction predictions when inputs consist solely of a molecular description of the chemical and the amino acid sequence of the target protein^27,28. By leveraging protein sequences instead of full 3D structures, we can reduce the computational burden and allow for the rapid evaluation of vast libraries of compounds, making it a more practical and scalable strategy for polypharmacology. However, the generalization power of deep learning methods for protein–ligand interaction (PLI) predictions remains poor. Furthermore, the chemical space of small organic molecules is astronomically vast. Although the number of possible small organic molecules is approximate 10⁶⁰ (ref. 29), only around 10⁶ compounds have annotated protein targets^30–32. The limited coverage of chemical genomics space makes it challenging to train generalizable deep learning models for binding affinity predictions in an out-of-distribution (OOD) scenario^33,34, in which unseen testing chemicals are substantially different from training data.

Because drug discovery is a high-stakes process, making decisions based on incorrect predictions can waste time and resources. Knowing the confidence level of an individual prediction will facilitate prioritizing lead compounds more effectively and efficiently. This requires an estimation of prediction errors on a compound-by-compound basis, not just a uniform confidence interval for all cases³⁵. The uncertainty of prediction comes from either the lack of labelled data or data noisiness. The Gaussian process (GP) is one of the popular approaches for uncertainty quantification. Several works^36,37 propose a combined GP and multilayer perceptron (MLP) approach for various biological tasks. However, the proposed GP+MLP algorithm is computationally intensive and requires modification of the architecture of predictive models. Many studies have leveraged ensembles of neural networks for uncertainty quantification^38,39. Similar to the ensemble of neural networks, Bayesian-regularized neural networks generate a distribution of network weights and variations on predictions⁴⁰. However, ensemble-based methods are time-consuming during both the training and inference stages, posing challenges for their application to large-scale polypharmacology. Along with these methods, conformal prediction (CP) aims to provide distribution-free uncertainty quantification under data exchangeability assumptions, a weaker version of independent and identically distribution^41,42. Many works have attempted to utilize CP for various aspects in drug discovery^43–45. A vanilla CP offers uniform confidence intervals for all cases but does not quantify the prediction error of a specific individual prediction. In practice, however, different cases warrant different levels of confidence: for example, OOD cases should naturally have larger intervals than in-distribution cases. Moreover, the exchangeability assumption underlying the conformer prediction does not hold in OOD cases.

To overcome the aforementioned challenges in large-scale polypharmacology compound screening, we propose a model-agnostic anomaly detection-based individual uncertainty quantification method, embedding Mahalanobis Outlier Scoring and Anomaly Identification via Clustering (eMOSAIC). eMOSAIC is specifically designed to meet the following requirements: (1) accuracy for proteins lacking structural, holo-conformational or binding site information, (2) reliability when applied to compounds with novel scaffolds (OOD scenarios) and (3) scalability to millions of compounds across thousands of proteins for large-scale screening. eMOSAIC explicitly models clusters of multiple local distributions within a multimodal embedding space and utilizes the divergence between each cluster distribution and the embedding of an unseen test case as features for predicting residuals of test cases. The residual is the difference between a predicted value and the actual observed value. This approach differs from stochastic OOD detection methods⁴⁶, which rely on the global distribution of the embedding space to classify cases as OOD or non-OOD. In contrast, eMOSAIC not only detects OOD cases but also quantifies the prediction error of an OOD case. We apply eMOSAIC to various base models, including but not limited to TrustAffinityNet developed by us⁴⁷.

Under rigorous OOD benchmark studies, eMOSAIC significantly improves performance for binding affinity predictions when applied to multiple base models. The combination of eMOSAIC with TrustAffinityNet is superior to both state-of-the-art deep learning baselines and structure-based methods in terms of both accuracy and scalability for understudied proteins. Furthermore, eMOSAIC outperforms existing uncertainty quantification methods, demonstrating its potential for other machine learning applications. Thus, eMOSAIC represents a important advance in deep learning applications to drug discovery.

Results

Overview of eMOSAIC

We developed eMOSAIC using the ChEMBL31 database³⁰. The distribution of training, validation and test data is shown in Fig. 1a. We evaluated the performance of eMOSAIC in a scaffold-based OOD setting, where chemicals in the testing set have different chemical scaffolds from those in the training/validation set. This setup mimics real-life scenarios, where we are likely to encounter new potential chemicals that have not been used in the training process of the deep learning model. In addition, the chemicals in the testing set are structurally dissimilar to those in the training/validation set, as shown in Fig. 1b. Our OOD split is almost identical to the Uniform Manifold Approximation and Projection (UMAP)-based split⁴⁸ in terms of chemical similarity distributions between testing and training data (Supplementary Fig. 1). For the purpose of comparison, we also evaluate eMOSAIC in the in-distribution setting of random split.

Fig. 1 |. Dataset used for the model development and overall architecture of the pipeline.

a, ChEMBL31 dataset with OOD scaffold split to ensure no overlap of scaffolds among the training, validation and test sets. b, Tanimoto similarity distribution of the training versus test set (top) and validation versus test set (bottom). c, Training process of eMOSAIC. d, Training process of the binding affinity prediction module, TrustAffinityNet. e, Binding affinity prediction with uncertainty quantification for new protein–ligand pair. PLI represents a pair of interacting protein and chemical.

Figure 1c–e provides an overview of eMOSAIC. eMOSAIC trains a secondary model for uncertainty estimation alongside the task-specific base models. However, unlike conventional methods, eMOSAIC generalizes well to OOD data. The training of eMOSAIC involves three main steps. First, the embeddings of training examples that are used to train a task-specific deep learning model are clustered. Second, given a case in the training set of eMOSAIC, the Mahalanobis distances (MDs) between its embedding and each cluster of trained embeddings from the base model are calculated, and its absolute residual—that is, the absolute value of the difference between predicted value and actual value—is obtained using the base model. Finally, eMOSAIC is trained using these MDs as features and the absolute residuals as labels. During the inference stage (Fig. 1e), given a new case, in addition to the task-specific label ŷ (for example, binding affinity) predicted by the task-specific base model, eMOSAIC will predict the absolute residual |ŷ-y|, where y is the actual binding affinity, using MD features derived from the trained based model. The prediction is classified as high-confidence (low uncertainty) if the predicted absolute residual is less than 0.5 or low-confidence otherwise.

eMOSAIC is model-agnostic and can be applied to any trained task-specific deep learning models with embeddings. In this paper, we apply it to several protein–ligand binding affinity prediction models, including TrustAffinityNet developed by us⁴⁷ and other state-of-the-art models such as BIND⁴⁹, BACPI⁵⁰, DeepDTA⁵¹ and DeepPurpose⁵². As shown in Fig. 1d, TrustAffinityNet takes a protein sequence and a ligand SMILES as inputs. Leveraging a pretrained protein language model, ESMFold⁵³ and the graph isomorphism network (GIN)⁵⁴, the embeddings of protein sequences and chemical structures are combined via attention pooling. The generated vectors from the attention pooling are used to predict binding affinities and train eMOSAIC for uncertainty quantification. For simplicity, whenever eMOSAIC_T is mentioned, it specifically refers to the application of eMOSAIC to the TrustAffinityNet model.

We compare eMOSAIC with state-of-the-art methods in three tasks: (1) binding affinity prediction, (2) polypharmacology compound screening and (3) uncertainty quantification. For binding affinity prediction and polypharmacology screening tasks, we compare eMOSAIC_T with sequence-based deep learning models, including BIND⁴⁹, BACPI⁵⁰, DeepDTA⁵¹ and DeepPurpose⁵², as well as a typical protein–ligand docking method, AutoDock Vina⁵⁵, a deep learning-based docking model, KarmaDock⁵⁶, and a structure-based deep learning model, EHIGN⁵⁷. Along with this, we also compare the performance of eMOSAIC with other uncertainty quantification methods, including the GP-based RIO framework³⁷, Monte Carlo dropout⁵⁸ and CP⁴¹.

We evaluate the performance of eMOSAIC and baseline models for various tasks and corresponding metrics, as detailed in ‘Results’ and Methods. For a fair comparison, the same training, validation and testing sets, training and evaluation procedures, and binding affinity thresholds are applied to all baseline models in the comparisons for binding affinity predictions and uncertainty quantification.

eMOSAIC improves OOD binding affinity prediction

When we apply eMOSAIC to base models, including TrustAffinityNet, BIND, BACPI, DeepDTA and DeepPurpose, in the OOD setting, as seen in Fig. 2a, we observe significant performance improvements across all base models and all four regression metrics: root mean square error (RMSE), mean absolute error (MAE), Pearson correlation and Spearman correlation. These results demonstrate that eMOSAIC is model-agnostic and has the potential to enhance any deep learning models.

Fig. 2 |. Performance comparisons in the OOD setting.

eMOSAIC-enhanced models (eMOSAIC + TrustAffinityNet (eMOSAIC_T), eMOSAIC + BIND, eMOSAIC + DeepDTA, eMOSAIC + DeepPurpose and eMOSAIC + BACPI) are compared against the base models (TrustAffinityNet, BIND, DeepDTA, DeepPurpose and BACPI). a, Barplots comparing RMSE, MAE, Pearson correlation and Spearman correlation between actual and predicted binding affinities. Darker and lighter shaded bars indicate the results obtained from base models and applying eMOSAIC to the base model. The bars represent the mean values across all test pairs (n = 50,944 protein–ligand samples, unit of study = individual protein–ligand pairs). Error bars indicate the s.d. across three independent splits (obtained via resampling replicates). Statistical significance of the differences between baseline and eMOSAIC-augmented versions as determined by P values from the Mann–Whitney U-test is shown in the bottom of the barplot, and the exact P values are present in source data. b, Scatterplot of actual binding affinities versus predicted binding affinities. c, Receiver operating characteristic curves for binding classification. d, Precision–recall curves for binding classification. e, EF for virtual screening based on the binding affinity threshold of 1 μM. f, EF for virtual screening based on the binding affinity threshold of 10 μM.

In the OOD setting, using the structure-informed ESMFold model⁵³ for protein sequence pretraining, TrustAffinityNet outperforms state-of-the-art methods for the protein–ligand binding affinity predictions, as shown in Fig. 2a. eMOSAIC further boosts the performance of TrustAffinityNet by quantifying the uncertainty of predictions from it. eMOSAIC_T significantly outperforms all baseline models for the binding affinity predictions, as shown in Fig. 2a,b.

Although BIND, BACPI, DeepPurpose and DeepDTA have acceptable performance in the random-split setting (Supplementary Table 1), their performance significantly drops in the scaffold-split setting. In contrast, the correlation between predicted binding affinities by eMOSAIC_T and actual binding affinities remains high when testing chemicals have scaffolds different from those in the training set. In addition, the performance difference is not significant when using scaffold and random splitting for eMOSAIC_T (Supplementary Fig. 2). These findings clearly demonstrate the superior generalization power of eMOSAIC_T when predicting the binding affinity in the OOD setting. The predictions by eMOSAIC_T not only have higher correlation but also have significantly lower deviation as recorded by the RMSE on average 19.39%, 16.92%, 21.54% and 18.97% lower and MAE on average 20.03%, 16.98%, 22.55% and 19.95% lower when compared to BACPI, BIND, DeepPurpose and DeepDTA, respectively. As shown in Fig. 2b, the prediction errors mainly come from the low- and high-affinity regions where there are few training data. Although eMOSAIC_T can alleviate the problem compared to other methods, further improvement is needed.

When using a binding affinity threshold of pKi > −2, eMOSAIC_T also significantly outperforms all baseline models in a classification task. As shown in Fig. 2c,d, the receiver operating characteristic area under the curve and precision–recall area under the curve of eMOSAIC_T improve 4.6% and 4.7% over TrustAffinityNet, respectively. The improvements over other state-of-the-art models are more significant. They are at least 7.1% and 7.2%, respectively. The performance improvement in the low false positive rate is more obvious. When the false positive rate is 0.05, the positive rate of eMOSAIC_T is more than 50% higher than the existing methods, demonstrating the potential of eMOSAIC_T for compound screening. As seen in Fig. 2e,f, the enrichment factor (EF) of eMOSAIC_T is indeed significantly higher than baseline models in the early ranking for virtual screening.

eMOSAIC significantly outperforms structure-based methods for binding affinity predictions and compound screening

We further compare the performance of eMOSAIC_T with structure-based methods widely applied in compound virtual screening. Besides the commonly used baseline Autodock Vina, the state-of-the-art method KarmaDock⁵⁶ was included because it has been reported to outperform widely used docking packages such as GLIDE⁵⁹ on standard docking and screening benchmarks while also offering substantial speed advantages, thereby serving as a rigorous and practically relevant baseline. We also include EHIGN⁵⁷, a structure-based binding affinity prediction method that can be applied to virtual screening when binding conformations are computed via docking. When tested on the scaffold-split set, structure-based methods show a notably poor correlation between predicted and actual binding affinities, as demonstrated in Fig. 3a,b (also Supplementary Table 2). The structure-based deep learning models, EHIGN and KarmaDock, offer modest improvements over AutoDock Vina. In contrast, eMOSAIC_T significantly enhances the correlation of predicted binding affinities, showing an eightfold improvement over KarmaDock, as shown in Fig. 3. For the structure-based experiment, the definition of binding pockets is critical. However, in a real-world polypharmacology application, the binding pocket often is not clearly defined due to a lack of experimentally determined ligand-bound structures. Two methods are applied to determining binding pockets: the AlphaFill approach based on the alignment with cocrystallized complex structures⁶⁰ and predicted binding pockets from a machine learning method, P2Rank⁶¹. Although both AutoDock Vina and EHIGN perform better on these proteins with AlphaFill-predicted binding pockets, KarmaDock works better on the proteins with P2rank-predicted binding pockets.

Fig. 3 |. Performance comparison of eMOSAICT with structure-based methods.

Results on the Pearson correlation (a), Spearman correlation (b), EF with 1 μM threshold (c) and EF with 10 μM threshold (d).

For high-throughput screening applications, we compare the EF for thresholds of 1.0 μM and 10.0 μM. In this context, eMOSAIC_T significantly outperforms all structure-based baselines, as shown in Fig. 3c,d. The EF from eMOSAIC_T is approximately eight times higher than that of structure-based methods for the threshold of 1.0 μM and three-and-a-half times higher for the threshold of 10.0 μM for top 0.1% ranked compounds. Note that a transformation is applied to align the docking score distribution with the mean and s.d. of the ground truth pKi before computing the performance metrics. Our findings suggest that for a large-scale polypharmacology screening where understudied proteins are often involved, eMOSAIC has a clear advantage over structure-based methods.

Furthermore, eMOSAIC_T is several orders of magnitude faster than structure-based methods (Supplementary Fig. 3). It takes less than 0.01 s for eMOSAIC_T, around 1 s for KarmaDock and 30 s for AutoDock Vina to predict the binding affinity of a protein–ligand pair. As a result, screening 1 million compounds against a target with eMOSAIC_T can be completed in just a few days.

We also evaluated eMOSAIC_T along with sequence-based and structure-based deep learning baselines on LIT-PCBA⁶², a benchmark curated for compound virtual screening on a limited number of well-characterized ligand-bound conformations with well-defined binding pockets, an unrealistic scenario for polypharmacology. Because LIT-PCBA provides only separate ligand information and protein structures, an additional time-consuming docking step is required to generate protein–ligand binding complexes before structure-based scoring can be applied⁵⁷. As shown in Supplementary Fig. 4 and consistent with previous results, eMOSAIC_T outperforms all sequence-based models and achieves performance comparable to the structure-based method EHIGN that is trained using holo structures.

eMOSAIC outperforms state-of-the-art uncertainty quantification methods

We compared eMOSAIC with three state-of-the-art methods for uncertainty quantification: RIO (which is based on the GP), Monte Carlo dropout and CP. Figure 4a presents sparsification curves, which assess how effectively uncertainty estimates identify unreliable predictions. Each curve is generated by progressively removing the most uncertain predictions and plotting the model’s performance (RMSE, in this case) on the remaining data as a function of the fraction removed. eMOSAIC achieves the lowest area under the sparsification error (AUSE), indicating the most effective uncertainty ranking. Its RMSE decreases sharply as uncertain samples are excluded, a desirable property for uncertainty quantification. In contrast, the RMSE of all baseline models shows little improvement. In addition, eMOSAIC has the highest Pearson correlation between predicted residual and actual residual, as shown in Fig. 4b.

Fig. 4 |. Performance comparison of eMOSAIC with state-of-the-art uncertainty quantification methods.

a, Sparsification curves: RMSE on remaining samples as an increasing fraction of highest-uncertainty cases is removed. AUSE is reported in the legend (lower is better). b, Barplot of Pearson correlation between predicted and true residuals across three independent random train/test splits of the scaffold-OOD dataset. Bars indicate mean values over independent protein–ligand pairs (n = 50,944 per split; unit of study = independent protein–ligand pairs), with s.d. across the three splits shown as error bars. Statistical significance of correlation improvement was evaluated between eMOSAIC and CP using a two-sided Mann–Whitney U-test. c, Errors and correlations of actual versus predicted binding affinities for high-confidence predictions selected by each method. Statistical significance between eMOSAIC and CP was assessed using a two-sided Mann–Whitney U-test. Exact P values for all comparisons are provided in the source data.

When evaluating methods to enhance absolute binding affinity predictions, applying either eMOSAIC or CP to the base model TrustAffinityNet leads to a statistically significant performance improvement, as shown in Fig. 4c. In contrast, RIO and Monte Carlo dropout do not provide noticeable gains. Between the two best methods, eMOSAIC outperforms CP across all evaluation metrics. A key distinction is that eMOSAIC offers instance-specific error estimation, whereas CP provides only global confidence intervals without per-sample error estimates, as demonstrated by examples in Supplemental Table 3. Instance-specific uncertainty quantification is particularly critical for high-stakes applications such as drug discovery, where decisions must be made on a compound-by-compound basis. Interestingly, the high-confidence predictions generated by eMOSAIC and those identified by CP are largely complementary, as illustrated in Supplementary Fig. 5. The Pearson correlation between their predicted values is modest (r = 0.38), suggesting that the two methods capture different aspects of the predictive uncertainty. These observations point to the potential value of integrating eMOSAIC and CP to develop a more robust and comprehensive framework for individual uncertainty quantification.

The Monte Carlo dropout method is a crude approximation of uncertainty estimation. Its performance is worse than the base model without uncertainty quantification because the Bayesian posterior distribution estimated by dropout might be more complex⁶³. Second, Monte Carlo dropout explores limited configurations of potential weights, not all of them, resulting in incomplete uncertainty estimation⁶⁴. The GP-based RIO framework performs even worse than Monte Carlo dropout, due to the limited capability of the current kernels that fail to effectively model the dependency between the PLI embeddings and residuals. It results in less accurate residual correction and uncertainty estimation and hence the poorest performance among all the methods.

The number of clusters in the embedding space may affect the performance of eMOSAIC. We evaluated the impact of this parameter (Supplementary Table 4). Our findings indicate that fine-grained clustering improves the capture of detailed information about the model’s uncertainty. However, beyond 50 clusters, performance plateaus, suggesting that further increasing the number of clusters does not yield significant improvements.

eMOSAIC enhances polypharmacology screening

G-protein-coupled receptors (GPCRs) constitute one of the most important target classes in pharmacology, accounting for approximately 35% of all approved therapeutic drugs. GPCRs’ druggability and accessibility make them central targets for therapeutic interventions⁶⁵. Along with these, protein kinase inhibitors have become a critical class of drugs, especially in oncology. Up to 33% of the drug development process targets these kinases⁶⁶. Furthermore, many clinically successful therapeutics targeting GPCRs and kinases are proven to be polypharmacological drugs¹⁰. Therefore, we assess the models’ performance in a polypharmacological context by evaluating the ability of ligands to interact with multiple GPCRs and kinases. Predicted pKi values for protein–ligand pairs are used to identify meaningful interactions. We apply a threshold of less than 100 μM of the predicted pKi values because weak bindings may contribute to polypharmacology¹⁰. The evaluation is based on a multilabel classification to determine the potential of ligands to selectively act across different therapeutic targets. We evaluate the models using multilabel accuracy, microprecision, microrecall, micro-F₁ score and Hamming loss. For the docking scores, a transformation was applied to align their distribution with the mean and s.d. of the ground truth pKi before computing the performance metrics.

As shown in Table 1, eMOSAIC_T demonstrates significantly improved performance over other methods across all evaluation metrics for GPCR polypharmacology. Specifically, accuracy, precision, recall and F₁ score are improved by 18.4%, 15.5%, 16% and 15.6%, respectively, and the Hamming loss is reduced by 166.7% compared with the second-best method. For kinase polypharmacology, docking-based and structure-based methods including AutoDock, KarmaDock and EHIGN, have the best sensitivity but poor specificity, aligning with the observations that protein–ligand docking has a high false positive rate in general. In contrast, eMOSAIC_T has balanced sensitivity and specificity, and thus it has the best accuracy and F₁ score. Among all the evaluation metrics, the Hamming loss is the best for evaluating multilabel classification performance. eMOSAIC_T reduces the Hamming loss by 8.3% compared with the second-best method.

Table 1 |.

Performance comparison of polypharmacology compound screening on GPCR and kinase proteins

Protein type	Model	Accuracy ↑	Precision ↑	Recall ↑	F1 score ↑	Hamming loss ↓
GPCR	AutoDock	0.50	0.45	0.62	0.47	0.50
	KarmaDock	0.52	0.52	0.40	0.35	0.48
	EHIGN	0.50	0.50	0.57	0.44	0.50
	DeepPurpose	0.72	0.70	0.68	0.67	0.28
	DeepDTA	0.74	0.72	0.70	0.69	0.26
	BIND	0.76	0.77	0.75	0.76	0.24
	BACPI	0.74	0.74	0.72	0.65	0.26
	eMOSAIC T	0.90	0.89	0.87	0.87	0.09
Kinase	AutoDock	0.42	0.42	0.79	0.41	0.58
	KarmaDock	0.51	0.51	0.45	0.28	0.48
	EHIGN	0.46	0.46	0.67	0.38	0.54
	DeepPurpose	0.70	0.61	0.64	0.59	0.30
	DeepDTA	0.71	0.63	0.63	0.61	0.29
	BIND	0.72	0.63	0.63	0.63	0.26
	BACPI	0.68	0.68	0.41	0.31	0.32
	eMOSAIC T	0.76	0.66	0.65	0.62	0.24

Bold, best method; underlined, second-best model.

Discussion

In this work, we propose eMOSAIC_T, a framework for accurate, reliable and scalable prediction of binding affinity along with an estimation of the associated uncertainty. We have successfully demonstrated the robust OOD generalization capabilities of eMOSAIC_T, yielding reliable (high-confidence) binding affinity with high accuracy. Furthermore, we highlight the framework’s notable advantage in terms of rapid inference speed, in contrast to protein–ligand docking, thereby rendering it well-suited for deployment in automated polypharmacology processes to leverage uncertainty-based methodologies.

Despite eMOSAIC_T’s superior performance, it has certain limitations that can be further addressed. First, the partitioning of the embedding space affects eMOSAIC’s performance. Using improved clustering methods, such as supervised contrastive learning, could yield better results. Second, in classification-based single-target compound screening, decoys are commonly used to augment the dataset with inactive compounds. However, in regression tasks aimed at predicting binding affinities, it is not straightforward to assign ‘pseudo’ binding affinity values to decoys. Incorporating semisupervised techniques⁶⁷ during training may facilitate data augmentation and boost its generalization to OOD data. Third, the available bioactive chemicals used for training and testing are biased towards the existing target space and established medicinal chemistry practices. Further studies are needed to assess the generalization ability of machine learning methods in unexplored chemical spaces. Finally, the interpretability of trained machine learning models in OOD scenarios remains an open challenge. This issue warrants further investigation, such as by incorporating sequence-based ligand binding pose predictions⁶⁸ or exploring the combination of uncertainty quantification and interpretability within a transformer-based evidential learning framework⁶⁹.

Methods

eMOSAIC for uncertainty quantification

eMOSAIC algorithm.

Figure 1c–e provides an overview of eMOSAIC. We utilize embedding clustering and MD to identify anomalies and quantify uncertainties. P. C. Mahalanobis introduced the MD metric in 1936 as a measure of anomaly and for detecting outliers⁷⁰. For a given point in a distribution, the MD is defined as follows:

$$\text{MD}\left(x\right)=\sqrt{\left(x-\mu \right){\Sigma }^{-1}{\left(x-\mu \right)}^{\text{T}}}$$ (1)

Here, μ and Σ refer to the mean and variance of the distribution, and T denotes the transpose. It considers the variance between the various variables in a multivariate distribution because real-life data often contain many correlated variables. MD has normalization through division by the covariance matrix, ensuring that variables with different scales are suitably handled. Because of these properties, it can effectively identify outliers or anomalies from the main distribution. MD has previously been used in deep learning frameworks for anomaly, OOD and adversarial detection in computer vision, time series and natural language processing tasks^71–74. These instances have clearly motivated the usage of MD for anomaly detection tasks to improve the reliability of predictions by deep learning frameworks.

After training a deep learning model, which serves as the protein– ligand binding affinity prediction module in this paper, we extract the PLI embeddings that represent the protein–ligand pairs. We then perform k-means clustering on the PLI embeddings from the deep learning model’s training set and obtain the mean (μ) and variance matrix (Σ) for each cluster to facilitate MD calculation. Given the embedding of an unseen PLI, we calculate the MD to each cluster and use these distances as features to train a simple MLP to predict the absolute residual, defined as the absolute difference between the predicted and true binding affinity. This approach enables the MLP to identify the PLI anomalies predicted by the model. For outlier detection, we select protein–ligand pairs with predicted residuals below 0.5, ensuring that only high-confidence predictions are retained.

Training of eMOSAIC_T.

As shown in Fig. 1c, we first train the protein–ligand binding affinity prediction module. After training, we use the best model, selected based on validation set performance, to extract PLI embeddings and generate clusters. Finally, we train eMOSAIC, the OOD uncertainty quantification module, to identify outliers. More details about the hyperparameters and configuration of the eMOSAIC_T and TrustAffinityNet model are present in Supplementary Tables 5 and 6.

Binding affinity module TrustAffinityNet

The binding affinity module, TrustAffinityNet, consists of three submodules: the protein sequence module, the ligand processing module and the PLI module. All of these modules collaboratively predict the binding affinity associated with the PLI.

Protein sequence module.

Protein sequence representation is one of the most vital components in the machine learning frameworks for predicting not only PLIs^34,50,51,75 but also their 3D structure^23,53. Protein sequences contain information that can be used to infer protein structure, function and family⁵³, making them a rich source of data for machine learning models^23,53. Large datasets of protein sequences are available^23,76,77, enabling machine learning frameworks to learn high-level, general representations of proteins. We utilize ESMFold⁵³ to obtain the protein sequence embeddings, which deploys a large language model, ESM-2, alongside a folding module and a structure module for modelling the protein structure. The ESM-2 protein language model, which is able to capture the protein structures at the fine resolution of the atomic level, consists of variable parameters ranging from 8 million to 15 billion. We use the 650-million-parameter model to obtain the refined protein sequence representation. We observed that the sequence representations obtained from the structure module of the ESMFold model performed better than the protein embeddings obtained from the ESM-2 model directly as well as the sequence embeddings obtained from the folding block, possibly because the structure block refines the protein sequence obtained from the ESM-2 model. We remove protein sequences greater than 700 in length as they are very few in number, and due to constraints on time and memory. Because the embeddings obtained are variable in size corresponding to the protein sequence length, to make them consistent for the next steps, we perform padding to pad the sequences with lengths less than 700 and define masks associated with the sequences that track the padding. Because convolutional neural networks are known to work well with processing sequence representations, we use the ResNet model⁷⁸ with five layers, and each layer has four convolutional layers to obtain a refined protein sequence embedding. Finally, adaptive masking (using interpolation) is used based on the changes to the embeddings to avoid the loss of information.

Ligand module.

We represent each ligand as a two-dimensional graph, where the nodes symbolize atoms and the edges are bonds. Embeddings for both node and edge are learned using GIN⁵⁴. For atom or node attributes, we used atom types, hybridization types, atom degrees, atom chirality, atom formal charges and atom aromatic, all converted to one-hot encoding before being utilized by GIN. We use a five-layer GIN architecture, which aggregates and updates node embedding for each atom/node. To obtain a graph-level or a ligand-level embedding that remains permutation invariant, a final sum pooling operation is used.

PLI module.

After obtaining both the protein and ligand embeddings, we use the attentive pooling network⁷⁹ such that the model is aware of both protein and ligand and that the interaction isn’t solely dependent on either protein or ligand. This network gives us the attention-weighted embeddings for both, which are then concatenated and fed to an MLP that predicts the final binding affinity.

Experiments

Dataset.

We train TrustAffinityNet and eMOSAIC on the ChEMBL31 database³⁰, which consists of 350,400 PLI pairs, along with their binding affinity values in nanomolar (nM), denoted as Ki. ChEMBL has been extensively used as a benchmark to develop and evaluate molecular property prediction models^80,81. In the experiments, we split the dataset into training, testing and validation sets by 7:2:1. Negative log (base 10) transformation was performed on Ki (binding affinity) to obtain pKi values. The data were split using the following methods. (1) Random split uses a random selection of protein–ligand pairs. (2) Random scaffold split uses a random selection of scaffolds of chemical structures⁸² such that the chemicals in the testing set have different scaffolds from those in the training/validation set. Scaffold split ensures that there is no overlap of the scaffold in the training, validation and testing sets. This was done to validate the model’s generalization power in a real-world OOD setting. Moreover, considering the vast chemical space for drug discovery, it is very likely that the model will encounter unknown and new scaffolds. In addition, no chemicals in the testing set are similar to those in the training/validation set with a Tanimoto coefficient larger than 0.4. (3) UMAP-based clustering split, introduced by Guo et al.⁴⁸, performs dimensionality reduction of Morgan fingerprints using UMAP followed by clustering. In addition to ChEMBL31, we utilize LIT-PCBA⁶², which is a benchmark curated for unbiased compound virtual screening. Specifically, we include data for seven targets: ALDH1 (Protein Data Bank (PDB) 4×4l), FEN1 (PDB 5fv7), GBA (PDB 2v3e), KAT2A (PDB 5h84), MAPK1 (PDB 2ojg), PKM2 (PDB 3gr4) and VDR (PDB 3a2j). The protein–ligand complex structures of these seven targets were obtained from ref. 83.

Baseline models.

We test eMOSAIC against baselines in two different objectives: (1) binding affinity prediction and (2) uncertainty quantification.

For the binding affinity prediction task, we compare eMOSAIC_T with sequence-based deep learning models, including BIND⁴⁹, BACPI⁵⁰, DeepDTA⁵¹ and DeepPurpose⁵², as well as a typical protein– ligand docking method, AutoDock Vina⁵⁵; a deep learning-based docking model, KarmaDock⁵⁶; and a structure-based deep learning model, EHIGN⁵⁷, on the OOD test set. Specifically, when applying eMOSAIC to sequence-based deep learning models, for each case, we attempt to select the most meaningful embeddings. For BACPI, we select the embeddings after the bidirectional interaction of both the ligand and protein representations. For BIND, we apply eMOSAIC to learnable commutative monoid aggregated embeddings, which already have ligand-influenced protein representation. For DeepPurpose and DeepDTA, we extract the protein and ligand representation once they have been encoded by their respective architectures and then use that as input to eMOSAIC. Finally, for TrustAffinityNet, we extract the attentive pooling-based representation. To compare with docking programs, AutoDock and KarmaDock were used to predict the docking scores for ligand–protein pairs for the scaffold splitting dataset. Alphafill⁶⁰ annotated binding pockets were used to define the searching space for AutoDock and KarmaDock. After removing the binding pockets for 38 different irons, such as FE, NA, 1,246 binding pockets on 429 proteins were used as the predefined binding pockets to set up docking. If there were multiple pockets on one protein, the ligand was docked into all pockets, and the one with the best docking score was selected for this ligand–protein pair. For comparison with EHIGN on polypharmacology screenings, we obtained protein–ligand complex structures for the pairs in the OOD test set using AutoDock Vina. Binding affinity predictions were then obtained using the pretrained EHIGN model provided by the original authors, which was trained on the PDBBind dataset based on experimentally determined protein–ligand complex. The same pretrained model was also used to predict binding affinities for the seven targets from the LIT-PCBA dataset with the corresponding protein–ligand complex structures obtained from additional docking⁸³. For other proteins not in the Alphafill dataset, P2Rank⁶¹ was used to predict the binding pockets on their Alphafold predicted model structures. The ligands were then docked into these binding pockets by AutoDock and KarmaDock, and the best docking scores were selected for these ligand–protein pairs. To evaluate how many of the top-ranked pairs have high binding affinities, the top 1,000 hit rates were calculated, which measured how many of the top 1,000 ranked pairs had pKi > −log(1) or −log(10). A higher top₁₀₀₀ hit rate correlates with better docking rank performance.

Along with this, we also compare the performance of our uncertainty quantification module with other uncertainty quantification methods, including the GP-based RIO framework³⁷, Monte Carlo dropout⁵⁸ and CP⁴¹. CP relies on obtaining a non-conformity score to measure the confidence interval for the predictions. This non-conformity score is crucial, and several different methods exist to compute it⁴², and several methods are available that describe various variants of CP⁸⁴. The most common non-conformity measures for regression-based models are based on absolute error, including using the calibration set’s absolute errors to provide intervals for the new predictions and using predicted absolute errors⁴². In our case, we use the attentive pooling embeddings to train an MLP for predicting the absolute error in the case of the CP baseline. We then select the points that have a confidence interval length lower than the average on the test set. We do the same for obtaining the high-confidence points for the Monte Carlo dropout method.

Evaluation.

For binding affinity prediction, we evaluate model performance using both regression-based metrics (including RMSE, MAE, Pearson’s correlation coefficient and Spearman’s correlation coefficient) and screening-based metrics (including EFs at different top-ranked percentages). Compounds are labelled as active or inactive based on binding affinity thresholds of 1.0 μM and 10.0 μM.

For polypharmacology compound screening, formulated as a multilabel classification task, we report accuracy, microprecision, microrecall, micro-F₁ score and Hamming loss. A binding affinity threshold of 100 μM is used to determine active targets, as weak binding interactions can still contribute to polypharmacological effects^6,10.

For uncertainty quantification, we evaluate the methods based on their impact on binding affinity prediction performance, the correlation between predicted and experimentally determined binding residues and the distribution of predicted binding residues.

Hardware and software.

Models were implemented in Python 3.10.9 using PyTorch (version 1.12.1, CUDA 11.3.1 and cuDNN 8) and PyTorch Geometric 2.2.0, on an NVIDIA Tesla V100 32-GB GPU. Additionally, analysis was performed using the following libraries: (1) NumPy (version 1.24.2), (2) Pandas (version 1.5.3), (3) SciPy (version 1.10.0) and (4) scikit-learn (version 2.1.1).

Supplementary Material

The online version contains supplementary material available at https://doi.org/10.1038/s42256-025-01151-2.

Data availability

The original binding affinity data were obtained from the publicly available ChEMBL database (https://www.ebi.ac.uk/chembl/). The LIT-PCBA dataset was downloaded from the official LIT-PCBA project website (https://drugdesign.unistra.fr/LIT-PCBA/), and the associated protein complexes were obtained via Zenodo at https://doi.org/10.5281/zenodo.4291724 (ref. 85). The scaffold-based split data used in our experiments, along with the UMAP-based split referenced in the Supplementary Information, are available via Zenodo at https://doi.org/10.5281/zenodo.17409085 (ref. 86.). Source data are provided with this paper.

Code availability

The code for this work is available via GitHub at https://github.com/XieResearchGroup/eMOSAIC and via Zenodo at https://doi.org/10.5281/zenodo.15313879 (ref. 87).