# Subgroup Discovery with the Cox Model

arXiv: https://arxiv.org/abs/2512.20762 · HTML: https://mlsanity.com/paper/2512.20762/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 16 of its 68 references that are in the map.
- First posted on arXiv 23 Dec 2025.

## Builds on (most cited first, 16 of 16)

- [DeepSurv: Personalized Treatment Recommender System Using A Cox Proportional Hazards Deep Neural Network](https://mlsanity.com/paper/1606.00931/) (2016)
- [Time-to-Event Prediction with Neural Networks and Cox Regression](https://mlsanity.com/paper/1907.00825/) (2019)
- [Machine Learning for Survival Analysis: A Survey](https://mlsanity.com/paper/1708.04649/) (2017)
- [Effective Ways to Build and Evaluate Individual Survival Distributions](https://mlsanity.com/paper/1811.11347/) (2018)
- [Survival Regression with Proper Scoring Rules and Monotonic Neural Networks](https://mlsanity.com/paper/2103.14755/) (2021)
- [An Effective Meaningful Way to Evaluate Survival Models](https://mlsanity.com/paper/2306.01196/) (2023)
- [Optimal Survival Trees](https://mlsanity.com/paper/2012.04284/) (2020)
- [Survival Kernets: Scalable and Interpretable Deep Kernel Survival Analysis with an Accuracy Guarantee](https://mlsanity.com/paper/2206.10477/) (2022)
- [Proper Scoring Rules for Survival Analysis](https://mlsanity.com/paper/2305.00621/) (2023)
- [Conditional Distribution Function Estimation Using Neural Networks for Censored and Uncensored Data](https://mlsanity.com/paper/2207.02384/) (2022)
- [Dynamical Survival Analysis with Controlled Latent States](https://mlsanity.com/paper/2401.17077/) (2024)
- [Time-to-Event Pretraining for 3D Medical Imaging](https://mlsanity.com/paper/2411.09361/) (2024)
- [Neural Frailty Machine: Beyond proportional hazard assumption in neural survival regressions](https://mlsanity.com/paper/2303.10358/) (2023)
- [Data-Driven Subgroup Identification for Linear Regression](https://mlsanity.com/paper/2305.00195/) (2023)
- [Learning Exceptional Subgroups by End-to-End Maximizing KL-divergence](https://mlsanity.com/paper/2402.12930/) (2024)
- [Modern approaches for evaluating treatment effect heterogeneity from clinical trials and observational data](https://mlsanity.com/paper/2311.14889/) (2023)

## Abstract (arXiv)

We study the problem of subgroup discovery for survival analysis, where the goal is to find an interpretable subset of the data on which a Cox model is highly accurate. Our work is the first to study this particular subgroup problem, for which we make several contributions. Subgroup discovery methods generally require a "quality function" in order to sift through and select the most advantageous subgroups. We first examine why existing natural choices for quality functions are insufficient to solve the subgroup discovery problem for the Cox model. To address the shortcomings of existing metrics, we introduce two technical innovations: the *expected prediction entropy (EPE)*, a novel metric for evaluating survival models which predict a hazard function; and the *conditional rank statistics (CRS)*, a statistical object which quantifies the deviation of an individual point to the distribution of survival times in an existing subgroup. We study the EPE and CRS theoretically and show that they can solve many of the problems with existing metrics. We introduce a total of eight algorithms for the Cox subgroup discovery problem. The main algorithm is able to take advantage of both the EPE and the CRS, allowing us to give theoretical correctness results for this algorithm in a well-specified setting. We evaluate all of the proposed methods empirically on both synthetic and real data. The experiments confirm our theory, showing that our contributions allow for the recovery of a ground-truth subgroup in well-specified cases, as well as leading to better model fit compared to naively fitting the Cox model to the whole dataset in practical settings. Lastly, we conduct a case study on jet engine simulation data from NASA. The discovered subgroups uncover known nonlinearities/homogeneity in the data, and which suggest design choices which have been mirrored in practice.