Training-Set Conditionally Valid Prediction Sets with Right-Censored Data

Abstract: Uncertainty quantification of prediction models through prediction sets is increasingly popular and successful, but most existing methods rely on directly observing the outcome and do not appropriately handle censored outcomes, such as time-to-event outcomes. Recent works have introduced distribution-free conformal prediction methods that construct predictive intervals for right-censored outcomes with marginal coverage guarantees. However, these methods typically assume a restrictive Type I censoring framework, in which censoring times are all observed. In this paper, we leverage a semiparametric one-step estimation framework and propose a novel approach for constructing predictive lower bounds on survival times with training-set conditional validity under right-censoring, where censoring times may be unobserved when the survival time is observed. With slight modification, our method can also provide predictive lower bounds with marginal guarantees. Through extensive simulations and a real-world application dataset tracking users' active times on a mobile application, we demonstrate the effectiveness and practicality of our approach. Compared to existing methods, our technique shows superior efficiency and robustness to model misspecifications, representing a significant advancement in the toolbox for reliable machine learning for time-to-event outcomes.