Distributionally Robust Learning in Survival Analysis

2025-06-02

AI Summary

Cox regression suffers from sensitivity to baseline distribution assumptions, model misspecification, and poor robustness against data perturbations. To address this, we introduce Wasserstein distributionally robust learning (DRL) to survival analysis for the first time, proposing a robust Cox model built upon a Wasserstein ambiguity set. Leveraging Wasserstein duality theory, we reformulate the original distributionally robust optimization problem into an equivalent, efficiently solvable exponential-cone regularized program, and establish finite-sample theoretical guarantees. Extensive simulations and experiments on real-world clinical datasets demonstrate that our method significantly outperforms conventional Cox regression and state-of-the-art survival models in both predictive accuracy and robustness under distributional shifts and noise perturbations. The proposed framework offers a statistically principled and computationally tractable paradigm for survival analysis in high-uncertainty medical settings.

Technology Category

• Machine Learning: Adversarial Learning & Robustness
• Constraint Satisfaction and Optimization: Distributed CSP/Optimization
• Reasoning under Uncertainty: Stochastic Optimization

Application Category

• Web Mining and Content Analysis: Robustness and generalizability of Web mining methods
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Attacks and countermeasures in recommendation systems

Abstract

We introduce an innovative approach that incorporates a Distributionally Robust Learning (DRL) approach into Cox regression to enhance the robustness and accuracy of survival predictions. By formulating a DRL framework with a Wasserstein distance-based ambiguity set, we develop a variant Cox model that is less sensitive to assumptions about the underlying data distribution and more resilient to model misspecification and data perturbations. By leveraging Wasserstein duality, we reformulate the original min-max DRL problem into a tractable regularized empirical risk minimization problem, which can be computed by exponential conic programming. We provide guarantees on the finite sample behavior of our DRL-Cox model. Moreover, through extensive simulations and real world case studies, we demonstrate that our regression model achieves superior performance in terms of prediction accuracy and robustness compared with traditional methods.

Problem

• Enhances robustness and accuracy in survival predictions
• Reduces sensitivity to data distribution assumptions
• Improves resilience to model misspecification and data perturbations

Innovation

• Incorporates DRL into Cox regression
• Uses Wasserstein distance-based ambiguity set
• Reformulates problem via exponential conic programming