T-SCI: A Two-Stage Conformal Inference Algorithm with Guaranteed Coverage for Cox-MLP

TL;DR

This paper introduces T-SCI, a two-stage conformal inference algorithm that guarantees coverage for survival time predictions with neural network models, addressing the challenge of censored data without linear assumptions.

Contribution

It proposes a novel two-stage conformal inference method, T-SCI, that guarantees coverage for Cox-MLP models under milder assumptions than existing methods.

Findings

T-SCI achieves guaranteed coverage in experiments.

The method outperforms existing approaches on synthetic and real data.

Theoretical analysis confirms coverage guarantees.

Abstract

It is challenging to deal with censored data, where we only have access to the incomplete information of survival time instead of its exact value. Fortunately, under linear predictor assumption, people can obtain guaranteed coverage for the confidence band of survival time using methods like Cox Regression. However, when relaxing the linear assumption with neural networks (e.g., Cox-MLP (Katzman et al., 2018; Kvamme et al., 2019)), we lose the guaranteed coverage. To recover the guaranteed coverage without linear assumption, we propose two algorithms based on conformal inference. In the first algorithm WCCI, we revisit weighted conformal inference and introduce a new non-conformity score based on partial likelihood. We then propose a two-stage algorithm T-SCI, where we run WCCI in the first stage and apply quantile conformal inference to calibrate the results in the second stage. Theoretical analysis shows that T-SCI returns guaranteed coverage under milder assumptions than WCCI. We conduct extensive experiments on synthetic data and real data using different methods, which validate our analysis.