Simpler Calibration for Survival Analysis

1 INTRODUCTION . Survival analysis , also known as time-to-event analysis , is the problem to predict the time of the occurrence of an event . In healthcare applications , the event typically corresponds to a death or the onset of disease in a patient . The time between a well-defined starting point and the occurrence of the event is called the survival time or failure time . In survival analysis , we usually estimate the distribution of the survival times of patients . Survival analysis has important applications in healthcare as well as various other fields ( e.g. , credit scoring ( Dirick et al. , 2017 ) and fraud detection ( Zheng et al. , 2019 ) ) . The recent progress of prediction models for survival analysis has been summarized in a survey paper ( Wang et al. , 2019 ) . In survival analysis , datasets are often censored , which means that events of interest might not be observed for some instances . This may be due to either the limited observation time window or missing traces caused by other irrelevant events . Typical censored data are right censored data . These are the data points whose exact times of the events are unknown ; we know only that the events had not happened up to a certain time . In this paper , we focus on the uncensored data and the right censored data , as shown in Figure 1 . Here , the event for data point x1 is observed during the period of study and hence this data is categorized as uncensored data . The data points x2 and x3 are categorized as right censored data because we did not observe the events during the period of study . The time between a well-defined starting point and the last observation time ( e.g. , the time of the end of study ) is called the censoring time . One of the classical methods to solve the survival analysis problem is the Kaplan-Meier estimator ( Kaplan & Meier , 1958 ) . This is a non-parametric method to estimate the distribution of the survival times as a survival function S ( t ) , where the value S ( t∗ ) for a specific time t∗ represents the survival rate at time t∗ ( i.e. , the ratio of the patients who survived at time t∗ ) . It is easy to estimate the survival function S ( t ) if the dataset contains only uncensored data points , but the Kaplan-Meier estimator is designed to work for datasets that include censored data . Here , we briefly explain the algorithm of the Kaplan-Meier estimator . Let { ti } ki=1 be the set of distinct times when at least one uncensored event was observed in the dataset . Let di be the number of ( uncensored ) events that happened exactly at time ti , and let ni be the number of data points that are known to have survived at time ti . Then , the Kaplan-Meier estimator outputs the survival function S ( t ) = ∏ i : ti≤t ( 1− di ni ) . Figure 2 shows an example of the survival function S ( t ) estimated by the Kaplan-Meier estimator for the flchain dataset ( Dispenzieri et al. , 2012 ) . Here , we can see that the survival rate at time t = 2500 is approximately 80 % . Note that the true survival rate S ( t ) at t =∞ must be zero , but the Kaplan-Meier estimator outputs the survival function S ( t ) only for time t ∈ [ 0 , tmax ] , where tmax is the maximum survival time of the uncensored data points in a dataset . This is because we can not estimate the survival rate S ( t ) for the time t > tmax . A drawback of the Kaplan-Meier estimator is that it outputs a survival function S ( t ) for the entire population and not for a specific patient . Therefore , there have been many algorithms to estimate the survival rate S ( t|x ) for each patient x so as to enable personalized medicine ( Wang et al. , 2019 ) . In particular , many neural network models that predict the survival function S ( t|x ) have been proposed ( Lee et al. , 2018 ; Ren et al. , 2019 ; Zheng et al. , 2019 ; Tjandra et al. , 2021 ) . 1.1 CALIBRATION . When we use a prediction model , it should be calibrated . In binary classification , this means that a prediction model that outputs a confidence of a true label for an input is expected to satisfy the condition that the average of the confidence values over all inputs are equal to the ratio of the data points with a true label in the dataset . For example , if a dataset contains 40 % of the data points with a true label , then we expect that a calibrated prediction model will output a confidence of 0.4 on average . Even though calibration is important for prediction models , and neural network models have been widely used for prediction models , ( Guo et al. , 2017 ) showed that neural network models are often miscalibrated . In regression analysis , quantile-based calibration is widely used as the definition of calibration ( Kuleshov et al. , 2018 ; Song et al. , 2019 ; Cui et al. , 2020 ; Zhao et al. , 2020 ) . In survival analysis , the definition of calibration is based on this definition for regression analysis . Goldstein et al . ( 2020 ) showed a method to train a neural network model to achieve calibration by adding a new regularizer , X-CAL , to the loss function of the neural network to achieve calibration during the training . This method is in contrast to the widely-used calibration methods for regression analysis such as Platt scaling and isotonic regression , which are post-training methods . Note that a calibrated model is not necessarily useful . For example , in binary classification we can construct a trivial calibrated model that always outputs the ratio of true label data points in a dataset as a confidence value . Therefore , a useful prediction model also needs to be sharp , which means that it outputs a confidence value close to one or zero for each input . To obtain a prediction model that is both calibrated and sharp , one approach is to train a neural network with a loss function consisting of one loss term aimed at sharp prediction and another aimed at calibration . X-CAL utilizes this approach in that it is a regularizer for calibration and is used in combination with another loss term that aims at sharp prediction . 1.2 OUR CONTRIBUTIONS . In this paper , we propose a Kaplan-Meier regularizer as a simpler alternative to X-CAL . An advantage of our regularizer is that the obtained prediction model is calibrated for any time t ∈ [ 0 , tmax ] . Another advantage is that our regularizer does not require any hyperparameter , whereas X-CAL requires hyperparameters . The idea behind our new regularizer is simple : just to reduce the difference between the average predicted survival function and the Kaplan-Meier survival function ( see Figure 3 ) . Our new regularizer is based on a new definition of calibration for survival analysis , and we discuss its advantages in Section 4 . 2 SURVIVAL ANALYSIS . We formally describe the problem settings of regression analysis and survival analysis . In regression analysis , we assume that there is an unknown probability distribution P on the sample space X ×Y , where X is the domain of features and Y is the interval [ −∞ , ∞ ] . We refer to the associated random variables with capital letters ( i.e. , X and Y ) and realizations with lower case letters ( x , y ) ∼ P . We assume that we can obtain independent samples { ( xi , yi ) } ni=1 from X ×Y according to distribution P . In survival analysis , we assume that there is an unknown probability distribution Q on the sample space X × T × C , where X is the domain of features and T and C are time intervals [ 0 , ∞ ] . We refer to the associated random variables with capital letters ( i.e. , X , T , and C ) and realizations with lower case letters ( x , t , c ) ∼ Q . The random variable T corresponds to the survival time ( i.e. , the time of the occurrence of an event ) , which might not be observable due to the censoring , and the random variable C corresponds to the censoring time . We assume that the censoring time is∞ for an uncensored event . The random variable defined by Z = min { T , C } corresponds to the time of the last observation ( i.e. , the survival time or the censoring time ) . Different from the regression analysis , we can not obtain samples { ( xi , ti , ci ) } ni=1 from X ×T ×C according to distributionQ due to the censoring in survival analysis . However , we assume that we can obtain independent samples D = { ( xi , zi , δi ) } ni=1 instead , where δi is a binary value indicating if the i-th data point is censored or not . Here an uncensored data point ( xi , zi , δi ) ∈ D satisfies δi = 1 and zi = ti and a censored data point ( xi , zi , δi ) ∈ D satisfies δi = 0 and zi = ci . Hence δi = 0 means that we know only the fact that ti ≥ ci and the exact survival time ti is unknown . The task of survival analysis is to predict the probability of an event of interest occurring at time t for x ∈ X as a probability distribution function f ( t|x ) . This function f ( t|x ) is often represented in other equivalent forms . For example , it can be represented as its cumulative distribution function ( CDF ) F ( t|x ) = ∫ t 0 f ( τ |x ) dτ or as the survival function S ( t|x ) defined by S ( t|x ) = 1− F ( t|x ) . Intuitively , F ( t|x ) represents the probability of observing the event by time t for x and the survival function S ( t|x ) represents the probability of not-observing the event until time t for x . Many neural network models have been proposed for survival analysis ( e.g. , ( Lee et al. , 2018 ; Ren et al. , 2019 ; Zheng et al. , 2019 ; Tjandra et al. , 2021 ) ) . They output the probability distribution in the form of fθ ( t|x ) , Fθ ( t|x ) , Sθ ( T |x ) , or something equivalent to these forms , where θ is the parameters of the neural network . 3 QUANTILE-BASED CALIBRATION . In this section , we review the definitions of calibration for distribution regression and survival analysis in the literature . First , we consider the distribution regression whose task is to predict the distribution of the target variable as a CDF Fθ ( y|x ) for x ∈ X , where θ is the parameters of the prediction model . In distribution regression , quantile-based calibration ( Kuleshov et al. , 2018 ) is widely used as the definition of calibration ( e.g. , ( Song et al. , 2019 ; Cui et al. , 2020 ; Zhao et al. , 2020 ) ) . Definition 3.1 ( Quantile-based calibration . ) A prediction modelFθ ( y|x ) for distribution regression is quantile-calibrated if this equation holds for any quantile level τ ∈ [ 0 , 1 ] : Pr ( X , Y ) ∼P ( Fθ ( Y |X ) ≤ τ ) = τ . ( 1 ) If we can compute the inverse of Fθ ( y|x ) , we can rewrite Eq . ( 1 ) as Pr ( X , Y ) ∼P ( Y ≤ F−1θ ( τ |X ) ) = τ . This equation means that , for a random sample ( x , y ) ∼ P , y must be at most the τ -th quantile of the predicted CDF ( i.e. , F−1θ ( τ |x ) ) exactly with probability τ . We can rewrite Definition 3.1 in another equivalent formulation : the following equation holds for any subinterval [ τ1 , τ2 ] ⊆ [ 0 , 1 ] : Pr ( X , Y ) ∼P ( Fθ ( Y |X ) ∈ [ τ1 , τ2 ] ) = τ2 − τ1 . This equation means that the quantile level Fθ ( y|x ) predicted for a random sample ( x , y ) ∼ P is contained in a subinterval [ τ1 , τ2 ] ⊆ [ 0 , 1 ] exactly with probability τ2 − τ1 . On the basis of Definition 3.1 , Goldstein et al . ( 2020 ) define calibration for survival analysis . Definition 3.2 ( Quantile-based calibration for survival analysis . ) A prediction model Fθ ( t|x ) for survival analysis is quantile-calibrated if this equation holds for any quantile level τ ∈ [ 0 , 1 ] : Pr ( X , T , C ) ∼Q ( Fθ ( T |X ) ≤ τ ) = τ . ( 2 ) We can rewrite Definition 3.2 into another equivalent formulation : the following equation holds for any subinterval I = [ τ1 , τ2 ] ⊆ [ 0 , 1 ] : Pr ( X , T , C ) ∼Q ( Fθ ( T |X ) ∈ I ) = |I| = τ2 − τ1 . ( 3 ) A problem when using Definition 3.2 is that we can not get samples directly from T due to the censoring . As such , we can not verify Eq . ( 3 ) for datasets that include censored data . Goldstein et al . ( 2020 ) resolved this problem by showing how to estimate the probability Pr ( Fθ ( t|x ) ∈ I ) for any t ∈ [ c , ∞ ] from Pr ( Fθ ( c|x ) ∈ I ) for a randomly sampled censored data point ( x , c , 0 ) . Under the assumption that T and C are independent ( i.e. , T ⊥ C |X ) , they show Pr ( Fθ ( t|x ) ∈ I ) = ( τ2 − v ) 1 [ v ∈ I ] 1− v + ( τ2 − τ1 ) 1 [ v < τ1 ] 1− v , ( 4 ) where v = Fθ ( c|x ) and 1 [ · ] denotes the step function . By using this estimation , we can compute the left-hand side of Eq . ( 3 ) from dataset D that include censored data points . On the basis of Eq . ( 3 ) and the approaches described in ( Andres et al. , 2018 ; Haider et al. , 2020 ) , Goldstein et al . ( 2020 ) proposed a metric called distributional calibration ( D-CAL ) , which is defined as ` D−CAL ( θ ) = ∑ I∈I ( E ( X , T , C ) ∼Q1 [ Fθ ( T |X ) ∈ I ] − |I| ) 2 , where the collection I is chosen to contain disjoint contiguous subintervals of C ⊆ [ 0 , 1 ] that cover the whole interval [ 0 , 1 ] . A prediction model Fθ ( t|x ) with a lower ` D−CAL ( θ ) is said to be more calibrated , but we can not construct a neural network model that directly minimizes D-CAL because ` D−CAL ( θ ) is not a differentiable function due to its step function . Therefore , Goldstein et al . ( 2020 ) defined the explicit calibration ( X-CAL ) , an approximation of D-CAL , by replacing the step function with a sigmoid function so that X-CAL becomes a differentiable function . Moreover , X-CAL is designed to handle a set of data points B = { ( xi , zi , δi ) } bi=1 as a mini-batch , which makes it possible to integrate XCAL into the loss function of neural network models because most of those models use a mini-batch training rather than the full batch training . Formally , X-CAL is defined as RX−CAL ( θ ) = EB∼Q ∑ I∈I ( E ( X , T , C ) ∼Bζ ( Fθ ( T |X ) ; I , γ ) − |I| ) 2 , where ζ ( z ; I , γ ) is a sigmoid function to approximate the step function in D-CAL . ( See ( Goldstein et al. , 2020 ) for the precise definition of the function ζ . ) Here , we abuse notation ( X , T , C ) ∼ B to indicate that we obtain sample data points from mini-batchB ( rather than the probability distribution Q ) . Note that Goldstein et al . ( 2020 ) proposed using X-CAL in the loss function of a neural network model as a regularizer , which means that it is intended to be combined with other loss functions . This is because a prediction model that aims only at calibration is useless ( as discussed in Section 1 ) and the balance between sharpness and calibration must be considered for obtaining a useful prediction model .

Summary: In this work, the authors propose a novel approach for learning calibrated predictions for survival analysis and similar tasks. The intuition of the approach is that the average probability that a prediction has a value less than or equal to $t$ should approximately equal the number of observations with value less than or equal to $t$. The authors then derive a differentiable regularization term using this idea and Kaplan-Meier curves. A small set of empirical evaluations suggest the proposed approach modestly outperforms another recent survival analysis calibration method according to some metrics.