Data-Driven Subgroup Identification for Linear Regression

Zachary Izzo,Ruishan Liu,James Zou

2023-04-29

Abstract:Medical studies frequently require to extract the relationship between each covariate and the outcome with statistical confidence measures. To do this, simple parametric models are frequently used (e.g. coefficients of linear regression) but usually fitted on the whole dataset. However, it is common that the covariates may not have a uniform effect over the whole population and thus a unified simple model can miss the heterogeneous signal. For example, a linear model may be able to explain a subset of the data but fail on the rest due to the nonlinearity and heterogeneity in the data. In this paper, we propose DDGroup (data-driven group discovery), a data-driven method to effectively identify subgroups in the data with a uniform linear relationship between the features and the label. DDGroup outputs an interpretable region in which the linear model is expected to hold. It is simple to implement and computationally tractable for use. We show theoretically that, given a large enough sample, DDGroup recovers a region where a single linear model with low variance is well-specified (if one exists), and experiments on real-world medical datasets confirm that it can discover regions where a local linear model has improved performance. Our experiments also show that DDGroup can uncover subgroups with qualitatively different relationships which are missed by simply applying parametric approaches to the whole dataset.

Machine Learning

What problem does this paper attempt to address?

The problem that this paper attempts to solve is that in medical research, simple parametric models (such as linear regression) are usually fitted on the entire data set, but there may be different trends or relationships among different sub - groups, resulting in a single model being unable to capture these heterogeneous signals. For example, in a data set, a certain drug may be effective for some patients, while being ineffective or even harmful for other patients. Therefore, the paper proposes a data - driven method - DDGroup, which is used to effectively identify sub - groups in the data, and these sub - groups have a unified linear relationship between features and labels. DDGroup can output an interpretable region, within which the linear model is expected to fit well.

Specifically, the main contributions of the paper include:
1. **Proposing a flexible formalization method** to define the sub - group selection problem, and proposing a general algorithm framework and its specific implementation - DDGroup.
2. **Proving the theoretical properties of DDGroup**, and the experimental results on synthetic data and real - world medical data sets show that DDGroup is effective in practice.
3. **Proving through experiments** that DDGroup can discover those sub - groups that would be missed when using simple parametric models to fit on the entire data set, and these sub - groups have qualitatively different relationships.

In conclusion, this paper aims to identify specific sub - groups in the data through a data - driven method, so that within these sub - groups, the linear model can better fit the data, thereby improving prediction accuracy and scientific interpretability.