In-Training Multicalibrated Survival Analysis for Healthcare via Constrained Optimization

TL;DR

This paper introduces GRADUATE, a novel survival analysis model that achieves multicalibration across all subpopulations by framing it as a constrained optimization problem, improving calibration and discrimination in healthcare datasets.

Contribution

The paper proposes a new multicalibrated survival analysis method called GRADUATE that optimizes calibration and discrimination simultaneously through constrained optimization.

Findings

GRADUATE outperforms state-of-the-art baselines on real-world clinical datasets.

It achieves better calibration across diverse subpopulations.

The method is mathematically proven to be near-optimal and feasible.

Abstract

Survival analysis is an important problem in healthcare because it models the relationship between an individual's covariates and the onset time of an event of interest (e.g., death). It is important for survival models to be well-calibrated (i.e., for their predicted probabilities to be close to ground-truth probabilities) because badly calibrated systems can result in erroneous clinical decisions. Existing survival models are typically calibrated at the population level only, and thus run the risk of being poorly calibrated for one or more minority subpopulations. We propose a model called GRADUATE that achieves multicalibration by ensuring that all subpopulations are well-calibrated too. GRADUATE frames multicalibration as a constrained optimization problem, and optimizes both calibration and discrimination in-training to achieve a good balance between them. We mathematically prove that the optimization method used yields a solution that is both near-optimal and feasible with high probability. Empirical comparisons against state-of-the-art baselines on real-world clinical datasets demonstrate GRADUATE's efficacy. In a detailed analysis, we elucidate the shortcomings of the baselines vis-a-vis GRADUATE's strengths.