Toward Conditional Distribution Calibration in Survival Prediction

2024-10-27

arXiv.org

AI Summary

To address the lack of conditional distribution calibration in survival prediction—which undermines individualized clinical decision support—this paper proposes the first conformal survival prediction framework explicitly designed for conditional calibration. Our method leverages model-predicted individual survival probabilities at observed time points to simultaneously achieve both marginal and conditional calibration, backed by asymptotic theoretical guarantees and practical implementability, without compromising discriminative performance (e.g., C-index). Unlike prior work focusing solely on marginal calibration, we systematically establish and demonstrate the critical role of conditional calibration in personalized prognostic decision-making. Extensive experiments across 15 real-world censored datasets show that our approach significantly improves both marginal and conditional calibration metrics while maintaining high discrimination, thereby validating its robustness and cross-dataset generalizability.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Reasoning under Uncertainty: Relational Probabilistic Models
• Constraint Satisfaction and Optimization: Distributed CSP/Optimization

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs

Abstract

Survival prediction often involves estimating the time-to-event distribution from censored datasets. Previous approaches have focused on enhancing discrimination and marginal calibration. In this paper, we highlight the significance of conditional calibration for real-world applications -- especially its role in individual decision-making. We propose a method based on conformal prediction that uses the model's predicted individual survival probability at that instance's observed time. This method effectively improves the model's marginal and conditional calibration, without compromising discrimination. We provide asymptotic theoretical guarantees for both marginal and conditional calibration and test it extensively across 15 diverse real-world datasets, demonstrating the method's practical effectiveness and versatility in various settings.

Problem

• Incomplete Data Distribution
• Predictive Accuracy
• Personal Decision-Making

Innovation

• Survival Probability Prediction
• Individualized Forecasting Accuracy
• Wide Applicability