Distributionally Robust Learning in Survival Analysis

Yeping Jin, Lauren Wise, Ioannis Paschalidis
Proceedings of the sixth Conference on Health, Inference, and Learning, PMLR 287:369-380, 2025.

Abstract

We introduce an innovative approach that incorporates a $\textit{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.

Cite this Paper

BibTeX
@InProceedings{pmlr-v287-jin25a, title = {Distributionally Robust Learning in Survival Analysis}, author = {Jin, Yeping and Wise, Lauren and Paschalidis, Ioannis}, booktitle = {Proceedings of the sixth Conference on Health, Inference, and Learning}, pages = {369--380}, year = {2025}, editor = {Xu, Xuhai Orson and Choi, Edward and Singhal, Pankhuri and Gerych, Walter and Tang, Shengpu and Agrawal, Monica and Subbaswamy, Adarsh and Sizikova, Elena and Dunn, Jessilyn and Daneshjou, Roxana and Sarker, Tasmie and McDermott, Matthew and Chen, Irene}, volume = {287}, series = {Proceedings of Machine Learning Research}, month = {25--27 Jun}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v287/main/assets/jin25a/jin25a.pdf}, url = {https://proceedings.mlr.press/v287/jin25a.html}, abstract = {We introduce an innovative approach that incorporates a $\textit{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.} }
Endnote
%0 Conference Paper %T Distributionally Robust Learning in Survival Analysis %A Yeping Jin %A Lauren Wise %A Ioannis Paschalidis %B Proceedings of the sixth Conference on Health, Inference, and Learning %C Proceedings of Machine Learning Research %D 2025 %E Xuhai Orson Xu %E Edward Choi %E Pankhuri Singhal %E Walter Gerych %E Shengpu Tang %E Monica Agrawal %E Adarsh Subbaswamy %E Elena Sizikova %E Jessilyn Dunn %E Roxana Daneshjou %E Tasmie Sarker %E Matthew McDermott %E Irene Chen %F pmlr-v287-jin25a %I PMLR %P 369--380 %U https://proceedings.mlr.press/v287/jin25a.html %V 287 %X We introduce an innovative approach that incorporates a $\textit{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.
APA
Jin, Y., Wise, L. & Paschalidis, I.. (2025). Distributionally Robust Learning in Survival Analysis. Proceedings of the sixth Conference on Health, Inference, and Learning, in Proceedings of Machine Learning Research 287:369-380 Available from https://proceedings.mlr.press/v287/jin25a.html.

Related Material