Subgroup Discovery with the Cox Model

2025-12-23

AI Summary

This study addresses interpretable subgroup discovery in survival analysis, aiming to identify data subsets that substantially improve prediction accuracy of the Cox proportional hazards model. We propose a novel paradigm that embeds subgroup mining directly into the Cox framework. Two theoretically grounded metrics are introduced: Expected Prediction Entropy (EPE), quantifying uncertainty in risk predictions within a subgroup, and Conditional Rank Statistic (CRS), measuring local association between covariates and event ordering. We provide formal theoretical guarantees for algorithmic correctness. Integrating information-theoretic principles, rank-based statistical inference, and eight domain-adapted search strategies, our method consistently outperforms the global Cox model on both synthetic benchmarks and real-world datasets—including NASA’s jet engine simulation data—successfully recovering known nonlinear patterns and engineering design principles.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Search and Optimization: Mixed Discrete/Continuous Search
• Reasoning under Uncertainty: Graphical Models

Application Category

• Graph Algorithms and Modeling for the Web: Representation, reconstruction, and subgraph or motif discovery in Web-related graphs
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

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.

Problem

• Develops subgroup discovery methods for survival analysis using Cox model.
• Introduces new metrics to evaluate survival models and subgroup deviations.
• Proposes algorithms to find interpretable subsets where Cox model performs accurately.

Innovation

• Introduces expected prediction entropy for hazard function evaluation
• Develops conditional rank statistics to quantify individual deviation
• Proposes eight algorithms combining both metrics for subgroup discovery