Inverse-Weighted Survival Games

1 Introduction . Survival analysis is the modeling of time-to-event distributions and is widely used in healthcare to predict time from diagnosis to death , risk of disease recurrence , and changes in level of care . In survival data , events , known as failures , are often right-censored , i.e. , only a lower bound on the time is observed , for instance , when a patient leaves a study before failing . Under certain assumptions , maximum likelihood estimators are consistent for survival modeling [ Kalbfleisch and Prentice , 2002 ] . Recently , deep survival models have obtained state-of-the-art results [ Ranganath et al. , 2016 , Alaa and van der Schaar , 2017 , Katzman et al. , 2018 , Kvamme et al. , 2019 , Zhong and Tibshirani , 2019 ] . Common among these are discrete-time models [ Yu et al. , 2011 , Lee et al. , 2018 , Fotso , 2018 , Lee et al. , 2019 , Ren et al. , 2019 , Kvamme and Borgan , 2019b , Kamran and Wiens , 2021 , Goldstein et al. , 2020 , Sloma et al. , 2021 ] even when data are continuous because they can borrow classification architectures and flexibly approximate continuous densities [ Miscouridou et al. , 2018 ] . Though training is often based on maximum likelihood , criteria such as Brier score ( BS ) and Bernoulli log likelihood ( BLL ) have been used to evaluate survival models [ Haider et al. , 2020 ] . The BS and BLL are classification losses adapted for survival by treating the model as a binary classifier at various time ∗Equal Contribution . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . horizons ( will the event occur before or after 5 years ? ) [ Kvamme and Borgan , 2019b , Lee et al. , 2019 , Steingrimsson and Morrison , 2020 ] . BS can also be motivated by calibration ( section 3 ) which is valuable because survival probabilities are used to communicate risk [ Sullivan et al. , 2004 ] . However BS and BLL are challenging to estimate because they require inverse probability of censor-weighting ( IPCW ) , which depends on the true censoring distribution [ Van der Laan et al. , 2003 ] . Though consistent , due to finite data , maximum likelihood may lead to models with poor BS and BLL . But directly optimizing these criteria is challenging because IPCW estimation requires solving an additional survival modeling problem to estimate the unknown censoring distribution . This poses a re-weighting dilemma : each model is required for training the other under these criteria but neither are known . To resolve the dilemma , we introduce Inverse-Weighted Survival Games for training with respect to criteria such as BS and BLL . We pose survival analysis as a game with the failure and censoring models as players . Each model ’ s loss is built from IPCW estimates featuring the other model . Inspired by game theory [ Neumann and Morgenstern , 2007 , Letcher et al. , 2019 ] , we ask : should the censoring model ’ s re-weighting role in the failure objective be considered part of the censoring objective ? We find the answer to be no . In each step of training , each model follows gradients of its loss with the other model held fixed to compute weights . When the loss is proper ( e.g . BS , BLL ) [ Gneiting and Raftery , 2007 ] , we show that games have the true failure and censoring distributions as a stationary point . This means the models in the game do not leave the correct distributions once reached . We then describe one case where this stationary point is unique . Finally , we show that inverse-weighted game training achieves better BS and BLL than maximum likelihood methods on simulations and real world cancer and ill-patient data.2 2 Notation and background on IPCW . Notation . Let T be a failure time with CDF F ( t ) = P ( T ≤ t ) , density f , survival function F = 1 − F , and model FθT . Let C be a censoring time with CDF G , density g , G = 1 − G , and model GθC . This means G ( t ) = P ( C > t ) . Let G ( t − ) denote P ( C ≥ t ) . We observe features X , time U = min ( T , C ) and ∆ = 1 [ T ≤ C ] . For discrete models over K times , let θTt = Pθ ( T = t ) and θCt = Pθ ( C = t ) . Models . We focus on deep discrete models like those studied in Lee et al . [ 2018 ] , Kvamme and Borgan [ 2019b ] . The model maps inputs X to a categorical distribution over times . When the observations are continuous , a discretization scheme is necessary . Following Kvamme and Borgan [ 2019b ] , Goldstein et al . [ 2020 ] , we set bins to correspond to quantiles of observed times . We represent all times by the lower boundary of their respective interval . Assumptions . We assume i.i.d . data and random censoring : T ⊥ C | X [ Kalbfleisch and Prentice , 2002 ] . We also require the censoring positivity assumption [ Gerds et al. , 2013 ] . Let f = dF . Then : ∃ s.t . ∀x∀t ∈ { t ≤ tmax | f ( t|x ) > 0 } , G ( t−|x ) ≥ > 0 , ( 1 ) i.e . it is possible that censoring events occur late-enough for us to observe failures up until a maximum time tmax . Truncating at a maximum time is necessary in practice for continuous distributions because datasets may have no samples in the tails , leading to practical positivity violations [ Gerds et al. , 2013 ] . This truncation happens implicitly for categorical models by choosing the last bin . In this work , we model the censoring distribution . This task is dual to the original survival problem : the roles of censoring and failure times are reversed . Therefore , to observe censoring events properly , we also require a version of eq . ( 1 ) to hold with the roles of F and G reversed ( appendix A ) . IPCW estimators . Inverse probability of censor-weighting ( IPCW ) is a method for estimation under censoring [ Van der Laan et al. , 2003 , Bang and Robins , 2005 ] . Consider the marginal mean E [ T ] . IPCW reformulates such expectations in terms of observed data . Using IPCW , we can show that : E [ T ] = E X E T |X [ E [ 1 [ T ≤ C|X ] ] E [ 1 [ T ≤ C|X ] ] T ] = E X E T |X E C|X [ 1 [ T ≤ C ] G ( T−|X ) T ] = E T , C , X [ ∆U G ( U−|X ) ] 2Code is available at https : //github.com/rajesh-lab/Inverse-Weighted-Survival-Games We derive this fully in appendix B . The second equality holds because ∆ = 1 =⇒ U = T and means we can identify E [ T |X ] provided that we know G and that random censoring and positivity hold . 3 Time-dependent survival evaluations . Brier score ( BS ) [ Brier and Allen , 1951 ] is proper for classification , meaning that it has a minimum at the true data distribution [ Gneiting and Raftery , 2007 ] . The BS is often adapted for survival evaluations [ Lee et al. , 2019 , Kvamme et al. , 2019 , Haider et al. , 2020 ] . For time t , it computes differences between the CDF and true event status at t , turning survival analysis into a classification problem at a given time horizon : BS ( t ; θ ) = E [ ( FθT ( t | X ) − 1 [ T ≤ t ] ) 2 ] ( 2 ) BS is often used as a proxy for marginal calibration error [ Kumar et al. , 2018 , Lee et al. , 2019 ] , which measures differences between CDF levels α ∈ [ 0 , 1 ] and observed proportions of datapoints with Fθ ( T |X ) < α [ Demler et al. , 2015 ] . This usage of BS stems from its decomposition into calibration plus a refinement ( discriminative ) term [ DeGroot and Fienberg , 1983 ] . Unfortunately one can not compute BS unmodified since 1 [ T ≤ t ] is unobserved for a point censored before t. IPCW BS [ Graf et al. , 1999 , Gerds and Schumacher , 2006 ] estimates BS ( t ) under censoring : BS ( t ; θ ) = E [ F θT ( t | X ) 2∆1 [ U ≤ t ] G ( U− | X ) + FθT ( t | X ) 21 [ U > t ] G ( t | X ) ] . ( 3 ) eq . ( 3 ) is equivalent to eq . ( 2 ) ( appendix C ) . Negative Bernoulli log likelihood ( BLL ) is similar , but with log loss ( appendix D ) . BS and BLL are proper for classification at each time t , so their sum or integral over t is still proper ( appendix H ) . Proper objectives differ . Though negative log likelihood ( NLL ) , BS and BLL all have the same true distribution at optimum with infinite data , they may yield significantly different solutions in practice . For example , NLL-trained models may not achieve good BS [ Kvamme and Borgan , 2019a ] . In fig . 1 , we show test set NLL and BS for 5 models trained with NLL at different learning rates on Gamma-simulated data ( described in section 5.1 ) . NLL does not align with BS : models that have low NLL may not have low BS . Model 4 has the lowest NLL but not the lowest BS . When a practitioner requires good performance under BS or BLL , they should optimize directly for those metrics . Re-weighting dilemma . Censoring introduces challenges because we must use IPCW to estimate BS and BLL . Crucially , the G in eq . ( 3 ) is the true censoring distribution rather than a model , but during training , we only have access to models . This poses a dilemma : can the models be used in re-weighting estimates during training to successfully optimize these criteria under censoring ? 4 Inverse-Weighted Survival Games . A reasonable attempt to solve the dilemma is to jointly optimize the sum of Fθ and Gθ ’ s losses where each model re-weights the other ’ s loss . The expectation is that both models will improve over training and yield reliable IPCW estimates for each other . Concretely , consider this for eq . ( 3 ) plus the same objective with the roles of Fθ and Gθ reversed . Unfortunately , there exist solutions to this optimization problem with smaller loss than for the distributions from which the data was generated , making this summed objective improper for the pair of distributions . In fig . 2 , we plot this for IPCW BS ( t = 1 ) for models over two timesteps3 as a function of each model ’ s single parameter . To address this phenomenon , we introduce Inverse-Weighted Survival Games . In these games , a failure player and censor player simultaneously minimize their own loss function . The failure and censoring model are featured in both loss functions and playing the game results in a trained failure and censoring model . We show in experiments that these games produce models with good BS , BLL , and concordance relative to those trained with maximum likelihood . For simplicity , we present the games for marginal categorical models . The analysis can be extended to conditional parameterizations with the usual caveats shared by maximum likelihood . Our experiments explore the conditional setting .

Summary: The authors are interested in the setting of predicting discrete outcomes within time intervals of interest (ex. 5-year survival, 10-year survival) from time-to-event data, i.e. the setting in which some observations may be right-censored. The conventional approach relies on inverse probability of censorship weighting (IPCW), which first estimates the distribution over time-to-censorship $p(c|x)$ (for example, via a Kaplan-Meier estimate), then re-weights the data to estimate the distribution over time-to-event p(t|x)$.   The authors instead propose *Inverse-Weighted Survival Games*, in which they learn *both* the time-to-censorship and the time-to-event distributions in a non-adversarial game, using each to re-weight the other's loss. On the theoretical side, they show that when the loss function is proper, the resulting game has the true distributions as a stationary point. Moreover, when the loss function is the Brier Score, this is the only stationary point. On the experimental side, they run experiments on simulations and public datasets (Survival-MNIST, METABRIC, SUPPORT) and show that IPCW games improve upon NLL-trained models in terms of concordance, KM-weighted log-likelihood, and KM-weighted Brier Score in the data-limited regime.