Enhancing Uncertainty Quantification in Drug Discovery with Censored Regression Labels

TL;DR

This work adapt ensemble-based, Bayesian, and Gaussian models with tools to learn from censored labels by using the Tobit model from survival analysis, demonstrating that despite the partial information available in censored labels, they are essential to accurately and reliably model the real pharmaceutical setting.

Abstract

In the early stages of drug discovery, decisions regarding which experiments to pursue can be influenced by computational models. These decisions are critical due to the time-consuming and expensive nature of the experiments. Therefore, it is becoming essential to accurately quantify the uncertainty in machine learning predictions, such that resources can be used optimally and trust in the models improves. While computational methods for drug discovery often suffer from limited data and sparse experimental observations, additional information can exist in the form of censored labels that provide thresholds rather than precise values of observations. However, the standard approaches that quantify uncertainty in machine learning cannot fully utilize censored labels. In this work, we adapt ensemble-based, Bayesian, and Gaussian models with tools to learn from censored labels by using the Tobit model from survival analysis. Our results demonstrate that despite the partial information available in censored labels, they are essential to accurately and reliably model the real pharmaceutical setting.

Figures (16)

• Figure 1: Five-fold Temporal Split. Illustrating the five folds created based on the date of the experiments and how they are used to create three temporal settings, each with more training data. For each setting, the first subsequent fold is used for validation, and the second subsequent fold is used for testing. Note that the folds were created to have roughly equal size, not based on fixed time intervals.
• Figure 2: Overview of Models. Illustrations of all models used in this study. The top row shows ensemble-based and Bayesian methods for which epistemic uncertainty can be obtained from the standard deviation in sampled predictions. All models in the bottom row produce aleatoric estimates of uncertainty. Additionally, epistemic uncertainty can be derived from the Gaussian Ensemble and the Evidential model.
• Figure 3: Ablation Study. Summary of the difference between NLL for models trained without censored labels and with censored labels. All three temporal settings are combined for readability. Stars above the bars indicate that the censored model was significantly better ($p<0.05$) for a majority of temporal settings, whereas stars below the bars indicate that the model trained without censored labels was significantly better for a majority of temporal settings.
• Figure 4: Predictive Accuracy. Comparing the predictive accuracy of all models in terms of MSE, aggregated over 10 experiments. For each dataset, the best model in terms of mean MSE is marked with a start together with any other models not statistically worse based on a one-sided Mann-Whitney-Wilcoxon test. Apart from the Random Forest and Evidential models, all other models are trained with censored labels.
• Figure 5: Calibration Curves. Full calibration curves for all uncertainty estimates on the third temporal setting containing three folds in the training set, aggregated over 10 experiments. The black line in each panel illustrates what a perfectly calibrated model would look like. Apart from the Random Forest and Evidential models, all other models are trained with censored labels.