Leveraging uncertainty estimates for drug response prediction in cancer cell lines

Machine learning models for drug response prediction in cancer cell lines could advance precision oncology by tailoring treatments to the molecular tumor profile. Their application is challenged by variability in prediction quality and distribution shifts between training and application. Yet, the most effective uncertainty estimator is domain-specific. We benchmark uncertainty-aware models for drug response prediction. We focus on epistemic uncertainty via ensemble agreement, and aleatoric uncertainty via distributional modeling, or both.

We find that ensemble-based estimates are more sensitive to distribution shift and can flag out-of-distribution examples. In contrast, distributional models yield stronger prediction error reductions among high-confidence subsets. The combination can provide both advantages: an ensemble of neural networks that estimate a Gaussian predictive distribution can reduce the mean squared error by 64% when restricting predictions to the 10% most confident drug-cell line pairs, and reliably indicates distribution shifts and platform differences. We show that uncertainty estimates enable a new dimension of model interpretability: by attributing predicted uncertainty to input features, we identify genes that signal unpredictability of drug response rather than sensitivity or resistance. We further demonstrate uncertainty-guided selection of measurements for active learning.

The code is available at https://github.com/PascalIversen/LUDRP.