Bayesian Nonparametric Survival Analysis via Deep Dirichlet Process

Jiacheng Xu, Tsai Hor Chan, Fuying Wang, Guosheng Yin, Lequan Yu

Keywords: Bayesian nonparametric methods, Survival Analysis, Variational Inference

TL;DR: We develop a deep Dirichlet process mixture model for scalable survival analysis

Abstract: The analysis of time-to-event data has received increasing attention in many application fields. The key challenge is that the data are mostly incomplete, with the right censoring mechanism being the most popular form. While Cox's proportional hazards assumption has shown adaptivity to traditional time-to-event datasets, challenges are observed when generalizing this assumption to modern survival analysis --- the proportional hazards assumption is often violated when covariates are high-dimensional. Moreover, traditional parametric assumptions on the survival distribution mostly belong to the exponential family and thus the assumption is strong and their exponential decay rate leads to poor long-tail approximations. To overcome these challenges, we propose a novel deep learning framework for survival analysis, named **DDPSurv**, which adopts a deeply parameterized Dirichlet process (DP) mixture model on survival distribution. Different from previous deep parametric approaches which rely on strong statistical assumptions, our framework can model the survival distribution with greater flexibility by adopting a DP mixture model. With the DP mixture model, we can improve the flexibility in modelling the survival distributions and achieve better tail behaviour by including the heavy-tail distributions in the mixture. We theoretically show that the proposed model can approximate the true survival distribution at a tight concentration rate. Empirical evaluations on standard survival benchmarks validate the satisfactory performance of the proposed method. Extensive experiments on large-scale clinical datasets --- MIMIC-III and MIMIC-IV --- highlight the scalability and clinical significance of our method. Codes are anonymously available at https://anonymous.4open.science/r/DeepSurv-net-2215

Primary Area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.

Reciprocal Reviewing: I understand the reciprocal reviewing requirement as described on https://iclr.cc/Conferences/2025/CallForPapers. If none of the authors are registered as a reviewer, it may result in a desk rejection at the discretion of the program chairs. To request an exception, please complete this form at https://forms.gle/Huojr6VjkFxiQsUp6.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 3394