Subgroup Discovery with the Cox Model

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. We examine why existing quality functions are insufficient for this problem and introduce two technical innovations: the expected prediction entropy (EPE), a novel metric for evaluating survival models that predict hazard functions, and the conditional rank statistics (CRS), which quantifies individual point deviation from a subgroup's survival time distribution. We study the EPE and CRS theoretically and show they address problems with existing metrics. We then introduce eight algorithms for Cox subgroup discovery. Our main algorithm is based on the DDGroup framework of Izzo et al. (2023) and leverages both the EPE and CRS, allowing theoretical correctness guarantees in well-specified settings. Empirical evaluation on synthetic and real data confirms our theory, showing our methods recover ground-truth subgroups in well-specified cases and achieve better model fit than naively fitting the Cox model to the entire dataset. A case study on NASA jet engine simulation data demonstrates that discovered subgroups uncover known nonlinearities in the data and suggest design choices mirrored in practice.

We study a problem in medical and reliability prediction called survival analysis, where the goal is to predict how long it will take before an event occurs, such as equipment failure or disease progression. In many real-world datasets, a single prediction model does not work equally well for everyone. Instead, there may be smaller groups of individuals for whom the model is especially accurate. Our work focuses on automatically finding these groups in a way that remains understandable to humans. Existing methods for identifying such subgroups often rely on evaluation criteria that are poorly suited for modern survival prediction models. To address this, we introduce two new statistical tools that better measure how well a model captures patterns in survival data and how unusual an individual’s outcome is relative to others in the same subgroup. Using these ideas, we develop several new algorithms for subgroup discovery, including one with mathematical guarantees under ideal conditions. We test our methods on both simulated and real datasets. The results show that our approach can successfully recover meaningful hidden groups and produce more accurate predictions than applying a single model to the entire dataset. In a case study involving NASA jet engine simulations, the discovered subgroups revealed known nonlinear behaviors in engine performance and highlighted design patterns already used in practice.