Predictive Heterogeneity: Measures and Applications
This work addresses the challenge of data heterogeneity for improving generalization and fairness in machine learning applications such as precision medicine and autonomous driving, though it appears incremental as it builds on existing concepts of heterogeneity.
The paper tackles the problem of data heterogeneity in machine learning by proposing a measure called 'usable predictive heterogeneity' that accounts for model capacity and computational constraints, and demonstrates its utility in improving out-of-distribution generalization in tasks like income prediction, crop yield prediction, and image classification.
As an intrinsic and fundamental property of big data, data heterogeneity exists in a variety of real-world applications, such as precision medicine, autonomous driving, financial applications, etc. For machine learning algorithms, the ignorance of data heterogeneity will greatly hurt the generalization performance and the algorithmic fairness, since the prediction mechanisms among different sub-populations are likely to differ from each other. In this work, we focus on the data heterogeneity that affects the prediction of machine learning models, and firstly propose the \emph{usable predictive heterogeneity}, which takes into account the model capacity and computational constraints. We prove that it can be reliably estimated from finite data with probably approximately correct (PAC) bounds. Additionally, we design a bi-level optimization algorithm to explore the usable predictive heterogeneity from data. Empirically, the explored heterogeneity provides insights for sub-population divisions in income prediction, crop yield prediction and image classification tasks, and leveraging such heterogeneity benefits the out-of-distribution generalization performance.