Budget-constrained Active Learning to Effectively De-censor Survival Data

TL;DR

This work addresses learning from survival data under labeling budget constraints by extending BatchBALD to handle right-censoring and partial de-censoring. It introduces BB_surv, which adjusts probability estimates for censored instances and uses a budgeted, probe-depth aware acquisition strategy grounded in a maximum-coverage formulation with a $(1-\frac{1}{e})$ guarantee. The authors prove theoretical bounds, implement the method with real-world survival datasets (MIMIC-IV, NACD, SUPPORT), and demonstrate that BB_surv consistently yields superior MAE-PO performance compared to several baselines, including BatchBALD, across uniform and non-uniform costs. The approach enables cost-effective, information-rich data collection in survival analysis, with practical implications for clinical follow-up planning and reliability/survival forecasting in other domains.

Abstract

Standard supervised learners attempt to learn a model from a labeled dataset. Given a small set of labeled instances, and a pool of unlabeled instances, a budgeted learner can use its given budget to pay to acquire the labels of some unlabeled instances, which it can then use to produce a model. Here, we explore budgeted learning in the context of survival datasets, which include (right) censored instances, where we know only a lower bound on an instance's time-to-event. Here, that learner can pay to (partially) label a censored instance -- e.g., to acquire the actual time for an instance [perhaps go from (3 yr, censored) to (7.2 yr, uncensored)], or other variants [e.g., learn about one more year, so go from (3 yr, censored) to either (4 yr, censored) or perhaps (3.2 yr, uncensored)]. This serves as a model of real world data collection, where follow-up with censored patients does not always lead to uncensoring, and how much information is given to the learner model during data collection is a function of the budget and the nature of the data itself. We provide both experimental and theoretical results for how to apply state-of-the-art budgeted learning algorithms to survival data and the respective limitations that exist in doing so. Our approach provides bounds and time complexity asymptotically equivalent to the standard active learning method BatchBALD. Moreover, empirical analysis on several survival tasks show that our model performs better than other potential approaches on several benchmarks.

Figures (25)

• Figure 1: Showing how we modify the probabilities assigned to each of 5 time bins, for a person censored at a time 1.3, when the Probe Depth $k=1$. (left) Original probabilities for the 5 bins [ 0.20, 0.25, 0.20, 0.05, 0.30]. (middle) Showing $p_{cens}$, which corrects as the patient is alive for the first bin, so has 0 probability of dying. (right) Collecting all of the "un-knowable bins" into one: $p_{final}$.
• Figure 2: Plot of MAE-PO evaluation of different acquisition functions with error bars as average of 40 predictions by the model. Both settings include a probe depth of 10. The left plot uses NACD dataset with 100 uncensored and 900 censored points in the uniform setting. The right plot uses MIMIC dataset starting with 300 uncensored and 2700 censored points in the non uniform setting.
• Figure 3: A Venn Diagram of the different areas of study presented in this work.
• Figure 4: Representation of the maximum coverage problem visually. Each unique intersection between two sets is an intersection we denote as an element part of a set.
• Figure 5: PCA visualization of $BB_{surv}$ choosing 20 instances from NACD with probe depth = 10, and budget = 20. Uniform setting.

Theorems & Definitions (5)

• Definition A.1
• Theorem A.2
• proof
• Theorem A.3
• proof