Uncertainty-Aware Optimal Treatment Selection for Clinical Time Series

TL;DR

The findings indicate that the method has robust performance across different counterfactual estimation baselines, showing that introducing uncertainty quantification in these settings helps the current baselines in finding more reliable and accurate treatment selection.

Abstract

In personalized medicine, the ability to predict and optimize treatment outcomes across various time frames is essential. Additionally, the ability to select cost-effective treatments within specific budget constraints is critical. Despite recent advancements in estimating counterfactual trajectories, a direct link to optimal treatment selection based on these estimates is missing. This paper introduces a novel method integrating counterfactual estimation techniques and uncertainty quantification to recommend personalized treatment plans adhering to predefined cost constraints. Our approach is distinctive in its handling of continuous treatment variables and its incorporation of uncertainty quantification to improve prediction reliability. We validate our method using two simulated datasets, one focused on the cardiovascular system and the other on COVID-19. Our findings indicate that our method has robust performance across different counterfactual estimation baselines, showing that introducing uncertainty quantification in these settings helps the current baselines in finding more reliable and accurate treatment selection. The robustness of our method across various settings highlights its potential for broad applicability in personalized healthcare solutions.

Figures (6)

• Figure 1: Performance of the baselines CRN, CT, CF-ODE, G-NET, BNCDE when selecting treatments for patients from the cardiovascular (left) and COVID-19 (right) datasets. For most baselines, the $\text{RMSE}_{\text{selection}}$ for the potential outcome compared to the desired outcome decreases, the higher the uncertainty weight in the optimization.
• Figure 2: Comparison of clamping constraints for treatment selection, using the CRN as a counterfactual estimator under a range of uncertainty weights. On the left we show results on the cardiovascular dataset and on the right on the COVID-19 dataset. The performance of treatment selection is robust to different treatment constraints choices, with an overall improvement of the performance in treatment selection for higher uncertainty weights.
• Figure 3: Comparison of the model-agnostic uncertainty quantification methods: MC dropout, ensemble and geometric ensemble methods applied to the CRN on the cardiovascular dataset (upper) and COVID-19 dataset (lower). Performance of treatment selection when varying the uncertainty estimate of the selected samples (left) and varying the uncertainty weight on the whole dataset (right). Selecting the least uncertain samples yields more reliable predictions of the outcome. On the left, we evaluate the models on an increasing percentage of the least uncertain samples (solid line) and compare this to a random subset of the validation data (dashed line). In all uncertainty-quantification methods, treatment selection yields more reliable treatments by increasing the weight on the uncertainty objective.
• Figure 4: Encouraging a balancing representations with increasing weight on the HSIC objective has a negligible effect on reliable counterfactual estimation compared to the uncertainty objective
• Figure 5: Performance comparison of the baseline models CRN, CT, CF-ODE, G-NET, BNCDE in treatments selection for cardiovascular (left) and COVID-19 (right) datasets. This comparison includes varying percentages of the most uncertain samples alongside a random subset of the validation data. The trends observed across different baselines are relatively stable, showing overall improved performance with increased certainty in the sample selection.