Toward Conditional Distribution Calibration in Survival Prediction

Summary

This paper proposes a new postprocessing method, CSD-iPOT, for survival analysis based on conformal prediction.

Strengths

The proposed method tries to achieve the conditional calibration, which is known to be hard.

Weaknesses

This paper puts emphasis on achieving conditional calibration (not marginal calibration), but the conditional calibration is known to be hard to achieve even for datasets without any censored data point. See, e.g., + Lei, J. and Wasserman, L. (2014). Distribution-free prediction bands for non-parametric regression. Journal of the Royal Statistical Society: Series B (Statistical Methodology), 76(1):71–96. + Rina Foygel Barber, Emmanuel J Candes, Aaditya Ramdas, Ryan J Tibshirani, The limits of distribution-free conditional predictive inference, Information and Inference: A Journal of the IMA, Volume 10, Issue 2, June 2021, Pages 455–482. + Zhao et al., Individual Calibration with Randomized Forecasting, ICML 2020. This paper does not cite these papers, and this paper does not discuss the hardness of conditional calibration. Lack of extensive discussion on the hardness of conditional calibration is a serious problem of this paper. Despite the hardness, this paper claims that the conditional calibration can be achieved in Theorem 3.2. However, this theorem actually shows nothing: it shows only that, if we have an estimator that achieves the conditional calibration (i.e., Eq. (10) is satisfied), the output of the proposed method achieves the conditional calibration, too. I think that if we have an estimator that achieves the conditional calibration, the proposed method is not required. Furthermore, there are many other problems in this paper: + While $\Gamma_{M}$ should be a set of scalar values according to line 9 of Algorithm 1, $\Gamma_{M}$ is a set of pairs of scalar values according to line 11 of Algorithm 1. + According to lines 180-183, $R$ copies of comformity scores $\gamma_{i,M}$ are generated for each uncensored data point. If so, the equation between lines 537 and 538 is incorrect. + According to lines 180-183, $R$ copies of comformity scores $\gamma_{i,M}$ are generated for each uncensored data point. Neverthless, this paper assumes that the set $\Gamma_{M}$ does not have any tie in line 538. + The proof of Theorem C.1 are not fully described. Even though the goal of the proof is to prove equation (9) on $\rho_{1}$ and $\rho_{2}$, the proof for censored data points does not argue $\rho_{1}$ and $\rho_{2}$. An equation for censored data points analogous to the equation for uncensored data points (between lines 537-538) must be presented. + The proof of Theorem C.1 completely ignores $R$, even though Theorem C.1 does not hold if $R=1$. (I think that the main idea of the proposed method, CSD-iPOT, is to use a sufficiently large $R$ to "blur" the censored data points.) + The assumptions used in the proof of Theorem C.2 are not clearly stated before the proof: many implicit assumptions are used in lines 563 and 564. In particular, during the proof of (i), a statement (i.e., an implicit assumption) similar to (i) is used in line 563. + In the proof of Theorem C.2, the alleged proof of (i) on $x_{n+1}$ does not include any discussion on $x_{n+1}$. ------------------------ # Additional Comments I submitted the following comments in the Author-Reviewer discussion period, but it accidentally did not visible to the authors (probably due to the complex system on comments visibility of OpenReview). I noticed this fact during the Reviewer-AC discussion period, and the AC allowed me to post the comments here. ==comments begin== Thank you for your comments. I will keep my score. This paper has several critial problems. + The authors' comments did not give any evidence that the proposed algorithm achieves conditional calibration (claimed in lines 47-48). Since the hardness of the conditional calibration is an important topic in machine learning as already studied by many researchers, the authors must pay careful attention when they discuss on this topic. + Huge discrepancy between the implemented algorithm (with R=1000) in the experimental section and the provided proof (valid only for R=1), even though the key idea of the proposed algorithm CSD-iPOT is to "blur" a censored subject with a large . Regarding the presentation, all the assumptions must be clearly presented before the proof. The assumption is stated at the end of the proof (in lines 546-548). Minor thing: The authors violated the rule on the 1-page pdf in the rebuttal phase: "Please use this PDF only for figures (including tables) and captions that describe the figure." ==comments end==

Questions

Nothing.

Rating

3

Confidence

4

Soundness

1

Presentation

2

Contribution

2

Limitations

This paper does not discuss the hardness of conditional calibration.

Summary

The authors study the problem of how to create individual survival distribution (ISD) models which are well-calibrated, both in a marginal and conditional sense, without negatively affecting the discriminative performance. They refine the "Conformalized Survival Distribution" (CSD) approach and propose "Conformalized Survival Distribution using Individual survival Probability at Observed Time" (CSD-iPOT). Both are post-processing methods which utilize conformal prediction, and can be applied on top of various survival analysis models to improve their calibration. CSD-iPOT is however designed to not only improve marginal calibration, but also conditional calibration. The method is evaluated on 15 datasets, comparing baseline (without neither CSD nor CSD-iPOT), CSD and CSD-iPOT versions of 7 survival models. The results are quite promising overall.

Strengths

- The paper studies an interesting and important problem, how to build well-calibrated and discriminative survival analysis models. The special focus on conditional calibration is also important. - The paper is well-written and very solid overall, the authors definitely seem knowledgeable. - The general idea of the proposed method, to utilize the conformal prediction framework in order to adjust predicted survival distribution curves on calibration data, makes sense. - The evaluation is quite extensive with 15 datasets and 7 baseline survival analysis models, and a direct comparison with CSD. - The results are quite promising overall, CSD-iPOT improves the calibration of CSD in most cases without negatively affecting the discrimination very often.

Weaknesses

- I found parts of Section 3 and 4 a bit difficult to follow, the proposed method/metric could perhaps be described in a more intuitive way. Figure 2 is neat, but would be helpful to have a similar visualization to illustrate the difference between CSD and CSD-iPOT. - The technical novelty/innovation compared to CSD is perhaps somewhat limited, both methods employ the same general approach (post-processing methods which utilize conformal prediction). - The experimental results could be more convincing. CSD-iPOT usually improves the calibration compared to CSD, but far from always, and the gains also seem to be relatively small quite often. Overall, I find it quite difficult to judge how much added benefit CSD-iPOT actually would have compared to CSD in practice. Are there concrete examples where CSD leads to significant mis-calibration within a certain patient subgroup, which is addressed by CSD-iPOT? The current results/metrics are a bit abstract / difficult to interpret. Summary: - Well-written and very solid paper overall that studies an important/interesting problem. The technical novelty compared to CSD, and how much better the proposed method actually would be in practice, is however a bit unclear. I am leaning towards accept. *** *** *** *** **Update after the rebuttal:** I have read the other reviews and all rebuttals. 2/3 other reviews are also positive, and I think the authors respond well to the third. All my questions were addressed, Figure 1 and 2 in the provided pdf are neat. I will increase my score to "7: Accept", this is a very well-written and solid paper that I think should be accepted.

Questions

- 189: "Unlike CSD [8], which adjusts the ISD curves horizontally (changing the times, for a fixed percentile), our refined version scales the ISD curves vertically", could you perhaps visualize how both CSD and CSD-iPOT modify the ISD curves in an example? I think that could help illustrate how these two methods differ. - I thought that CSD-iPOT shouldn't affect the relative ordering of subjects, but when comparing the ISD curves before and after in Figure 2(d) and (e), the top orange curve is always above the top blue curve in (e), whereas it is partially below the blue curve in (d)? - 236: "Qi et al. [8] demonstrated that CSD theoretically guarantees the preservation of the original model’s discrimination performance in terms of Harrell’s concordance index (C-index) [1]. However, CSD-iPOT lacks this property", but both methods have a check mark for "Discrimination guarantee Harrell’s" in Table 1? Minor things: - 77: "In summary, individual calibration is ideal but not impractical", impractical --> practical? - 219: "with adequate modifications to the accommodate our method", remove "the"? - 297: "Compared CSD" --> "Compared to CSD"?

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

Yes.

Summary

The paper enhances the Conformalized Survival Distribution (CSD) post-processing framework to account conditional calibration. The proposed framework, CSD-iPOT, utilizes a conformal set to adjust survival curves vertically, aligning them with predetermined percentiles at test time. Unlike CSD, which relies on Kaplan-Meier curves, CSD-iPOT leverages Individualized Survival Distributions (ISD) for censored events when constructing the conformal set. The paper provides theoretical guarantees for marginal calibration, conditional calibration, and the monotonicity of ISD. Comprehensive experimental results from 15 datasets demonstrate that CSD-iPOT enhances both the marginal and conditional calibration of baseline models with minimal impact on the concordance index (C-index).

Strengths

- The paper is well-written and easy to follow. - The reviewer appreciates that the visual plots provided offer great intuitive illustrations of the proposed approach. - This is the first paper to address conditional calibration in survival analysis, an important yet under-explored problem. - The paper provides theoretical guarantees for the proposed approach in terms of marginal calibration, conditional calibration, and monotonicity. - CSD-iPOT is more computationally efficient in terms of storage than CSD. - Extensive experimental results across 15 datasets and 7 baselines demonstrate that CSD-iPOT significantly improves the calibration (both marginal and conditional) of baseline models with minimal loss in the concordance index (C-index).

Weaknesses

*The paper highlights that CSD-iPOT lacks theoretical guarantees for preserving Harrell's concordance index. This is a major limitation of this work; nevertheless, experimental results demonstrate that the impact is minimal* *The description in lines 142-162 requires some improvements:* - Eqn. 4: Needs to be adjusted to something like $\tilde{S}^{-1}(p | x_{n+1}) = T(\hat{S}^{-1}(\text{Percentile}(\cdot) | x_{n+1}))) $, where $T$ is the proposed vertical transformation. Then provide all the necessary numbered steps to finally obtain the calibrated ISD (Eqn. 5). *For completeness, I encourage the author(s) to also benchmark competitive baseline models that directly model event times, e.g., Chapfuwa et al. 2018 and Miscouridou et al. 2018* *Minor:* - In Table 1, CSD-iPOT should be marked with an 'X' under Harrell’s concordance index category. - Line 297: Typo -> "Compare to CSD"

Questions

- Given that CSD-iPOT lacks theoretical guarantees for preserving Harrell's concordance index, is there a hyperparameter that explicitly controls this trade-off? - It seems that CSD-iPOT achieves calibration lower than the empirical lower limit set by Kaplan-Meier. Are these the instances that result in an impact on the C-index? - Why does CSD-iPOT struggle with both calibration and the C-index for models such as DeepHit? - Are the calibration metrics obtained at different percentiles than those used for the post-processed ISD? What happens if these two sets are different? - Why does CSD-iPOT improve marginal and conditional calibration models better than CSD? Is the difference due to how CSD-iPOT handles censored events?

Rating

7

Confidence

4

Soundness

4

Presentation

3

Contribution

3

Limitations

I encourage the authors to discuss the limitations of their work, including any violations of modeling assumptions.

Summary

While previous work focuses on calibration in a marginal sense, this paper proposes a post-processing approach that also imposes conditional calibration on all individual features. Therefore, the proposed method can guarantee equal calibration for different groups, ensuring fairness. Contributions: 1. The paper proposes the post-processing approach to ensure conditional distribution calibration in survival analysis (accommodating censorship) and develops the corresponding metric. 2. The paper provides asymptotic guarantees for both marginal and conditional calibration and demonstrates the computation complexity. 3. The paper validates the method by conducting experiments on 15 datasets.

Strengths

Originality: The paper motivates the use of conditional distribution calibration in survival analysis via post-processing. The idea is intuitive as both the class-specific and the general survival probability at observed time should be uniformly distributed. Quality and significance: Although the calibration method is simple and known, the paper accommodates it to the censored setting and proves its effectiveness empirically. Besides, the paper provides corresponding theoretical guarantees. Lastly, it applies to various SOTA methods in survival analysis and improves the performance in most cases. Clarity: The paper clearly conveys the idea and method of calibrating survival probability curves.

Weaknesses

1. Missing related work in conformal prediction: the paper claims the method is based on conformal prediction but it is not mentioned and illustrated clearly in the main content. Thus, this knowledge gap hampers the understanding. 2. The method seems to be sensitive to the percentiles. As shown in Figures 2 (d) and (e), the post-processed curves are piecewise functions highly related to the percentiles (1/3, 2/3). Thus, the resulting curves would be sensitive to the choice of percentiles. Therefore, the number of percentiles should be included in the complexity and asymptotic behavior analysis. 3. The discrimination properties of the method are unclear. The authors first claim it lacks the discrimination performance preservation property for Harrel's C-index, which conflicts with the results in Table 1. 4. The empirical exploration and validation of the group fairness is missing.

Questions

As mentioned in the weakness: Q1: How does the choice of percentiles affect the performance? Are there any empirical results that can illustrate the effect? Q2: Does the approach keep the Harrel C-index? Q3: There are plenty of survival analysis methods exploring the heterogeneity of groups. How does the proposed method compare to the SOTA heterogeneity methods? Will it improve the performance?

Rating

5

Confidence

3

Soundness

3

Presentation

2

Contribution

2

Limitations

It remains unknown if this method would improve/hurt fairness in survival analysis. The study is missing.

Official Comment by Reviewer kAoa

Thanks for addressing most of my concerns. However, this statement "If a baseline model can generate only a scalar value, it is not immediately clear how to convert this scalar value into a distribution," is not necessarily accurate. Individual survival distributions can be obtained from parametric time-to-event models (e.g., AFT) since $f(t|x) = h(t|x) S(t|x)$. For non-parametric approaches (e.g., DATE), time-to-event distributions are implicitly defined through sampling. After reviewing the rebuttal and reviewer's response, I find that overall, this is a solid paper, and I am still leaning more towards acceptance. I am keeping my score as it is.