Toward Conditional Distribution Calibration in Survival Prediction

A survival dataset $\mathcal{D}=\{(\boldsymbol{x}_{i},t_{i},\delta_{i})\}_{i=1}^{n}$ contains $n$ tuples, each containing covariates $\boldsymbol{x}_{i}\in\mathbb{R}^{d}$, an observed time $t_{i}\in\mathbb{R}_{+}$, and an event indicator $\delta_{i}\in\{0,1\}$. For each subject, there are two potential times of interest: the event time $e_{i}$ and the censoring time $c_{i}$. However, only the earlier of the two is observable. We assign $t_{i}\triangleq\min\{e_{i},c_{i}\}$ and $\delta_{i}\triangleq\mathbbm{1}[e_{i}\leq c_{i}]$, so $\delta_{i}=0$ means the event has not happened by $t_{i}$ (right-censored) and $\delta_{i}=1$ indicates the event occurred at $t_{i}$ (uncensored). Let $\mathcal{I}$ denote the set of indices in dataset $\mathcal{D}$, then we can use $i\in\mathcal{I}$ to represent $(\boldsymbol{x}_{i},t_{i},\delta_{i})\in\mathcal{D}$.

Our objective is to estimate the Individualized Survival Distribution (ISD), $S(t\mid\boldsymbol{x}_{i})=\mathbb{P}(e_{i}>t\mid\mathbf{X}=\boldsymbol{x}_{i})$, which represents the survival probabilities of the $i$-th subject for any time $t\geq 0$.

Calibration measures the alignment between the predictions against observations. Consider distribution calibration at the individual level: if an oracle knows the true ISD $S(t\mid\boldsymbol{x}_{i})$, and draws realizations of $e_{i}\mid\boldsymbol{x}_{i}$ (call them $e_{i}^{(1)},e_{i}^{(2)},\ldots$), then the survival probability at observed time $\{S(e_{i}^{(m)}\mid\boldsymbol{x}_{i})\}_{m}$ should be distributed across a standard uniform distribution $\mathcal{U}_{[0,1]}$ (probability integral theorem (Angus, 1994)). However, in practice, for each unique $\boldsymbol{x}_{i}$, there is only one realization of $e_{i}\mid\boldsymbol{x}_{i}$, meaning we cannot check the calibration in this individual manner.

To solve this, Haider et al. (2020) proposed marginal calibration, which holds if the predicted survival probabilities at event times $e_{i}$ over the $\boldsymbol{x}_{i}$ in the dataset, $\{\hat{S}(e_{i}\mid\boldsymbol{x}_{i})\}_{i\in\mathcal{I}}$, matches $\mathcal{U}_{[0,1]}$.

We can “blur” each censored subject uniformly over the probability intervals after the survival probability at censored time $\hat{S}(c_{i}\mid\boldsymbol{x}_{i})$ Haider et al. (2020) (see the derivation in Appendix A):

Figure 1(a) illustrates how marginal calibration is assessed using 6 uncensored subjects. Figure 1(b, c) show the histograms and P-P plots, showing the predictions are marginally calibrated, over the predefined 3 bins, as we can see there are 6/3 = 2 instances in each of the 3 bins. However, if we divide the datasets into two groups (orange vs. blue – think men vs. women), we can see that this is not the case, as there is no orange instance in the $[\frac{2}{3},1]$ bin, and 2 orange instances in the $[\frac{1}{3},1]$ bin.

In summary, individual calibration is ideal but impractical. Conversely, marginal calibration is more feasible but fails to assess calibration relative to certain subsets of the population by features. This discrepancy motivates us to explore a middle ground – conditional calibration. A conditionally calibrated prediction, which ensures that the predicted survival probabilities are uniformly distributed in each of these groups, as shown in Figure 1(d, e, f), is more effective in real-world scenarios. Consider predicting employee attrition within a company: while a marginal calibration using a Kaplan-Meier (KM) Kaplan and Meier (1958) curve might reflect overall population trends, it fails to account for variations such as the tendency of lower-salaried employees to leave earlier. A model that is calibrated for both high and low salary levels would be more helpful for predicting the actual quitting times and facilitate planning. Similarly, when predicting the timing of death from cardiovascular diseases, models calibrated for older populations, who exhibit more predictable and less varied outcomes Visseren et al. (2021), may not apply to younger individuals with higher outcome variability. Using age-inappropriate models could lead to inaccurate predictions, adversely affecting treatment plans.

Methods based on the objective function Avati et al. (2020); Goldstein et al. (2020); Chapfuwa et al. (2020) have been developed to enhance the marginal calibration of ISDs, involving the addition of a calibration loss to the model’s original objective function (e.g., likelihood loss). However, while those methods are effective in improving the marginal calibration performance of the model, their model often significantly harms the discrimination performance Goldstein et al. (2020); Chapfuwa et al. (2020); Kamran and Wiens (2021), a phenomenon known as the discrimination-calibration trade-off Kamran and Wiens (2021).

Post-processed methods Candès et al. (2023); Qi et al. (2024) have been proposed to solve this trade-off by disentangling calibration from discrimination in the optimization process. Candès et al. (2023) uses the individual censoring probability as the weighting in addition to the regular Conformalized Quantile Regression (CQR) Romano et al. (2019) method. However, their weighting method is only applicable to Type-I censoring settings where each subject must have a known censored time Klein and Moeschberger (2006) – which is not applicable to most of the right-censoring datasets.

Qi et al. (2024) developed Conformalized Survival Distribution (CSD) by first discretizing the ISD curves into percentile times (via predefined percentile levels), and then applying CQR Romano et al. (2019) for each percentile level (see the visual illustration of CSD in Figure 6 in Appendix B). Their method handles right-censoring using KM-sampling, which employs a conditional KM curve to simulate multiple event times for a censored subject, offering a calibrated approximation for the ISD based on the subject’s censored time. However, their method struggles with some inherent problems of KM Kaplan and Meier (1958) – e.g., KM can be inaccurate when the dataset contains a high proportion of censoring Liu et al. (2021). Furthermore, we also observed that the KM estimation often concludes at high probabilities (as seen in datasets like HFCR, FLCHAIN, and Employee in Figure 9). This poses a challenge in extrapolating beyond the last KM time point, which hinders the accuracy of KM-sampling, thereby constraining the efficacy of CSD (see our results in Figure 3).

Our work is inspired by CSD (Qi et al., 2024), and can be seen as a percentile-based refinement of their regression-based approach. Specifically, our CiPOT effectively addresses and resolves issues prevalent in the KM-sampling, significantly outperforming existing methods in terms of improving the marginal distribution calibration performance. Furthermore, to our best knowledge, this is the first approach that optimizes conditional calibration within the survival analysis that can deal with censorship.

However, achieving conditional calibration (also known as conditional coverage in some literature Vovk et al. (2005)) is challenging because it cannot be attained in a distribution-free manner for non-trivial predictions. In fact, guarantees of finite sample for conditional calibration are impossible to achieve even for standard regression datasets without censorship Vovk et al. (2005); Lei and Wasserman (2014); Foygel Barber et al. (2021). This limitation is an important topic in the statistical learning and conformal prediction literature Vovk et al. (2005); Lei and Wasserman (2014); Foygel Barber et al. (2021). Therefore, our paper does not attempt to provide finite sample guarantees. Instead, following the approach of many other researchers Romano et al. (2019); Lei et al. (2018); Sesia and Candès (2020); Izbicki et al. (2020); Chernozhukov et al. (2021); Izbicki et al. (2022), we provide only asymptotic guarantees as the sample size approaches infinity. The key idea behind this asymptotic conditional guarantee is that the construction of post-processing predictions relies on the quality of the original predictions. Thus, we aim for conditional calibration only within the class of predictions that can be learned well – that is, consistent estimators.

We acknowledge that this assumption may not hold in practice; however, (i) reliance on consistent estimators is a standard (albeit strong) assumption in the field of conformal prediction Sesia and Candès (2020); Izbicki et al. (2020); Chernozhukov et al. (2021), (ii) to the best of our knowledge, no previous results have proven conditional calibration under more relaxed conditions, and (iii) we provide empirical evidence of conditional calibration using extensive experiments (see Section 5) .

For simplicity, our method is motivated by the split conformal prediction Papadopoulos et al. (2002); Romano et al. (2019). We start the process by splitting the instances of the training data into a proper training set $\mathcal{D}^{\text{train}}$ and a conformal set $\mathcal{D}^{\text{con}}$. Then, we can use any survival algorithm or quantile regression algorithm (with the capability of handling censorship) to train a model $\mathcal{M}$ using $\mathcal{D}^{\text{train}}$ that can make ISD predictions for $\mathcal{D}^{\text{con}}$ – see Figure 2(a).

With little loss of generality, we assume that the ISD predicted by the model, $\hat{S}_{\mathcal{M}}(t\mid\boldsymbol{x}_{i})$, are right-continuous and have unbounded range, i.e., $\hat{S}_{\mathcal{M}}(t\mid\boldsymbol{x}_{i})>0$ for all $t\geq 0$. For survival algorithms that can only generate piecewise constant survival probabilities (e.g., Cox-based methods Cox (1972); Kvamme et al. (2019), discrete-time methods Yu et al. (2011); Lee et al. (2018), etc.), the continuous issue can be fixed by applying some interpolation algorithms (e.g., linear or spline).

We start by sketching how CiPOT deals with only uncensored subjects. Within the conformal set, for each subject $i\in\mathcal{I}^{\text{con}}$, we define a distributional conformity score, wrt the model $\mathcal{M}$, termed the predicted Individual survival Probability at Observed Time (iPOT):

Here, for uncensored subjects, the observed time corresponds to the event time, $t_{i}=e_{i}$. Recall from Section 2.2 that predictions from model $\mathcal{M}$ are marginally calibrated if the iPOT values follow $\mathcal{U}_{[0,1]}$ – i.e., if we collect the distributional conformity scores for every subject in the conformal set $\Gamma_{\mathcal{M}}=\{\gamma_{i,\mathcal{M}}\}_{i\in\mathcal{I}^{\text{con}}}$, the $\rho$-th percentile value in this set should be equal to exactly $\rho$. If so, no post processing adjustments are necessary.

In general, of course, the estimated Individualized Survival Distributions (ISDs) $\hat{S}(t\mid\boldsymbol{x}_{i})$ may not perfectly align with the true distributions $S(t\mid\boldsymbol{x}_{i})$ from the oracle. Therefore, for a testing subject with index $n+1$, we can simply apply the following adjustment to its estimated ISD:

Here, $\text{Percentile}(\rho;\,\Gamma_{\mathcal{M}})$ calculates the $\frac{\left\lceil\rho(|\mathcal{D}^{\text{con}}|+1)\right\rceil}{|\mathcal{D}^{\text{con}}|}$-th empirical percentile of $\Gamma_{\mathcal{M}}$. This adjustment aims to re-calibrate the estimated ISD based on the empirical distribution of the conformity scores.

Visually, this adjustment involves three procedures:

1. (i)

   It first identifies the empirical percentiles of the conformity scores – $\text{Percentile}(\frac{1}{3};\,\Gamma_{\mathcal{M}})$ and $\text{Percentile}(\frac{2}{3};\,\Gamma_{\mathcal{M}})$, illustrated by the two grey lines at 0.28 and 0.47 in Figure 2(d), respectively – which uniformly divide the stars according to their vertical locations;

2. (ii)

   It then determines the corresponding times on the predicted ISDs that match these empirical percentiles (the hollow circles, where each ISD crosses the horizontal line);

3. (iii)

   Finally, the procedure shifts the empirical percentiles (grey lines) to the appropriate height of desired percentiles ($\frac{1}{3}$ and $\frac{2}{3}$), along with all the circles. This operation is indicated by the vertical shifts of the hollow points, depicted with curved red arrows in Figure 2(d).

This adjustment results in the post-processed curves depicted in Figure 2(e). It shifts the vertical position of the green star $\bigstar$ from the interval $[\frac{1}{3},\frac{2}{3}]$ to $[\frac{2}{3},1]$, and the pink star $\bigstar$ from $[0,\frac{1}{3}]$ to $[\frac{1}{3},\frac{2}{3}]$. These shifts ensure that the calibration histograms and P-P plots achieve a uniform distribution (both marginally on the whole population and conditionally on each group) across the defined intervals.

After generating the post-processed curves, we can apply a final step, which involves transforming the inverse ISD function back into the ISD function for the testing subject:

The simple visual example in Figure 2 shows only two percentiles created at $\frac{1}{3}$ and $\frac{2}{3}$. In practical applications, the user provides a predefined set of percentiles, $\mathcal{P}$, to adjust the ISDs. The choice of $\mathcal{P}$ can slightly affect the resulting survival distributions, each capable of achieving provable distribution calibration; see ablation study #2 in Appendix E.6 for how $\mathcal{P}$ affects the performance.

It is challenging to incorporate censored instances into the analysis as we do not observe their true event times, $e_{i}$, which means we cannot directly apply conformity score in (3) and the subsequent conformal steps. Instead, we only observe the censoring times, which serve as lower bounds of the event times.

Given the monotonic decreasing property of the ISD curves, the iPOT value for a censored subject, i.e., $\hat{S}_{\mathcal{M}}(t_{i}\mid\boldsymbol{x}_{i})=\hat{S}_{\mathcal{M}}(c_{i}\mid\boldsymbol{x}_{i})$, now serves as the upper bound of $\hat{S}_{\mathcal{M}}(e_{i}\mid\boldsymbol{x}_{i})$. Therefore, given the prior knowledge that $\hat{S}_{\mathcal{M}}(e_{i}\mid\boldsymbol{x}_{i})\sim\mathcal{U}_{[0,1]}$, the observation of the censoring time updates the possible range of this distribution. Given that $\hat{S}_{\mathcal{M}}(e_{i}\mid\boldsymbol{x}_{i})$ must be less than or equal to $\hat{S}_{\mathcal{M}}(c_{i}\mid\boldsymbol{x}_{i})$, the updated posterior distribution follows $\hat{S}_{\mathcal{M}}(e_{i}\mid\boldsymbol{x}_{i})\sim\mathcal{U}_{[0,\hat{S}_{\mathcal{M}}(c_{i}\mid\boldsymbol{x}_{i})]}$.

Following the calibration calculation in Haider et al. (2020), where censored patients are evenly "blurred" across subsequent bins of $\hat{S}(c_{i}\mid\boldsymbol{x}_{i})$, our approach uses the above posterior distribution to uniformly draw $R$ potential conformity scores for a censored subject, for some constant $R\in\mathbbm{Z}^{+}$. Specifically, for a censored subject, we calculate the conformity scores as:

Here, $\boldsymbol{u}_{R}$ is a pseudo-uniform vector to mimic the uniform sampling operation, significantly reducing computational overhead compared to actual uniform distribution sampling. For uncensored subjects, we also need to apply a similar sampling strategy to maintain a balanced censoring rate within the conformal set. Because the exact iPOT value is known and deterministic for uncensored subjects, sampling involves directly drawing from a degenerate distribution centered at $\hat{S}_{\mathcal{M}}(e_{i}\mid\boldsymbol{x}_{i})$ – i.e., just drawing $\hat{S}_{\mathcal{M}}(e_{i}\mid\boldsymbol{x}_{i})$ $R$ times. The pseudo-code for implementing the CiPOT process with censoring is outlined in Algorithm 1 in Appendix B.

Note that the primary computational demand of this method stems from the optional interpolation and extrapolation of the piecewise constant ISD predictions. Calculating the conformity scores and estimating their percentiles incur negligible costs in terms of both time and space, once the right-continuous survival distributions are established. We provide computational analysis in Appendix E.5.

Here we discuss the theoretical properties of CiPOT. Unlike CSD Qi et al. (2024), which adjusts the ISD curves horizontally (changing the times, for a fixed percentile), our refined version scales the ISD curves vertically. This vertical adjustment leads to several advantageous properties. In particular, we highlight why our method is expected to yield superior performance in terms of marginal and conditional calibration compared to CSD Qi et al. (2024). Table 1 summarizes the properties of the two methods.

Due to space constraints, the main text presents partial results for datasets with high censoring rates and high KM ending probabilities (HFCR, FLCHAIN, Employee, MIMIC-IV). Notably, CQRNN did not converge on the MIMIC-IV dataset. Thus, we conducted a total of 104 method comparisons ($15\ \text{datasets}\times 7\ \text{baselines}-1$). Table 2 offers a detailed performance summary of CiPOT versus baselines and CSD method across these comparisons. Appendix E.4 presents the complete results.