Heterogeneous Risk Minimization
This addresses a critical issue for machine learning practitioners dealing with real-world datasets that are often merged from multiple sources without clear labels, enabling more robust models in scenarios like healthcare or finance where distribution shifts are common.
The paper tackles the problem of out-of-distribution generalization in machine learning when datasets lack explicit source labels, proposing the Heterogeneous Risk Minimization (HRM) framework to jointly learn latent heterogeneity and invariant relationships, resulting in stable predictions under distributional shifts.
Machine learning algorithms with empirical risk minimization usually suffer from poor generalization performance due to the greedy exploitation of correlations among the training data, which are not stable under distributional shifts. Recently, some invariant learning methods for out-of-distribution (OOD) generalization have been proposed by leveraging multiple training environments to find invariant relationships. However, modern datasets are frequently assembled by merging data from multiple sources without explicit source labels. The resultant unobserved heterogeneity renders many invariant learning methods inapplicable. In this paper, we propose Heterogeneous Risk Minimization (HRM) framework to achieve joint learning of latent heterogeneity among the data and invariant relationship, which leads to stable prediction despite distributional shifts. We theoretically characterize the roles of the environment labels in invariant learning and justify our newly proposed HRM framework. Extensive experimental results validate the effectiveness of our HRM framework.