Heterogeneous Risk Minimization

Jiashuo Liu; Zheyuan Hu; Peng Cui; Bo Li; Zheyan Shen

Source-linked AI summary

Heterogeneous Risk Minimization

Jiashuo Liu, Zheyuan Hu, Peng Cui, Bo Li, Zheyan Shen

arXiv:2105.03818v3May 9, 2021cs.LG

TL;DR

ERM and existing invariant-learning methods face OOD generalization challenges when correlations shift and datasets contain unlabeled heterogeneity. HRM jointly identifies latent heterogeneity and invariant predictors, supported by theoretical analysis and experiments reporting stronger average, stability, and worst-case performance under distributional shifts.

Problem
ERM relies on IID data and can exploit correlations that are unstable under distributional shifts, while invariant learning requires training environments that may be unavailable in merged datasets.
Method
HRM jointly learns latent heterogeneity and invariant prediction through heterogeneity identification and invariant-learning modules.
Results
HRM shows superiority in average, stability, and worst-case performance under different settings of distributional shifts.
Takeaways & Limitations
HRM enables OOD generalization without explicit environment labels by jointly optimizing heterogeneity identification and invariant prediction.
Takeaways & Limitations
The analysis focuses on raw variables with X = [Φ∗, Ψ∗]T, and the clustering algorithm lacks theoretical convergence guarantees.

Abstract

from arXiv · show

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.

1. Introduction

OOD generalization is difficult because ERM exploits correlations that can fail under distributional shifts, while invariant learning often depends on environment labels unavailable in merged datasets. HRM addresses this by jointly identifying latent heterogeneity and learning invariant predictors, with theoretical and empirical support.

Motivation: ERM can generalize poorly under distributional shifts because it exploits unstable correlations in training data.The IID assumption underlying ERM is often violated by selection biases or unobserved confounders.
Limitations of existing methods: Invariant learning seeks stable predictors across training environments but becomes inapplicable when merged datasets lack explicit source labels.Its effectiveness also depends heavily on training-environment quality.
HRM framework: HRM jointly learns latent data heterogeneity and an invariant predictor to improve generalization despite distributional shifts.The framework includes heterogeneity identification and invariant-learning modules.
Empirical contribution: Experiments on synthetic and real-world datasets report superior average, stability, and worst-case performance under different distributional shifts.The reported evaluation covers multiple performance dimensions rather than a single outcome.
Theoretical contribution: HRM theoretically characterizes environment-label roles and justifies mutual promotion between heterogeneity identification and invariant learning.This reciprocal relationship motivates the framework’s joint optimization process.

2. Problem Formulation

The paper formulates OOD generalization through invariant predictors and shows why observed or unavailable environment labels may fail to identify the invariances needed for unseen environments. HRM addresses this formulation by exploiting latent heterogeneity to generate refined environments and learn an invariant model.

2.1. OOD and Maximal Invariant Predictor: OOD generalization seeks a predictor that minimizes worst-case risk across possible environments whose data distributions may differ from training environments.The paper formalizes risk using environment-conditional loss and assumes the set of all environments contains the training environments.
2.1. OOD and Maximal Invariant Predictor: Invariant learning assumes a representation Φ∗ whose conditional relationship with Y is invariant across environments and sufficient for prediction.The paper states that, without prior knowledge or structural assumptions, unseen latent environments cannot be characterized for OOD generalization.
2.1. OOD and Maximal Invariant Predictor: The maximal invariant predictor is defined by maximizing mutual information with Y over invariant representations, and it solves the stated OOD problem under the sufficiency assumption.The corresponding optimal prediction is EY[Y|Φ∗].
2.1. OOD and Maximal Invariant Predictor: Environment labels are not guaranteed to identify invariances shared by all possible environments, because training environments may constrain only a larger training-specific invariance set.The paper states IE ⊆ IEtr, so a predictor invariant on training environments need not be invariant on unseen environments.
2.1. OOD and Maximal Invariant Predictor: In the toy example, training environments e5 and e6 permit X2 in the invariance set, whereas subdividing them into e1–e4 yields the desired set based on X1.The example illustrates that finer environment partitions can exclude additional variant features.
2.2. Problem of Heterogeneous Risk Minimization: The paper therefore poses HRM as learning environments with minimal |IEtr| from heterogeneous data without environment labels, then learning an invariant model under them.The formulation targets latent environments whose varying P(Y|Ψ∗(X)) can exclude variant features, while acknowledging a circular dependency between heterogeneity and invariant learning.
2.2. Problem of Heterogeneous Risk Minimization: HRM resolves the circular dependency by jointly exploiting heterogeneity identification and invariant learning, whose interaction is theoretically supported as mutually promoting.The framework seeks tighter invariance constraints by generating environments that vary the relationship between the variant component and the target.

3. Method

HRM alternates between identifying heterogeneous environments from learned variant representations and learning invariant predictors from those environments. Its modules use clustering and feature-selection-based invariant learning, with joint optimization exploiting their mutual promotion.

3. Method: HRM combines a heterogeneity-identification module Mc with an invariant-prediction module Mp for pooled heterogeneous data.The framework is described for raw features decomposed into invariant Φ∗ and variant Ψ∗ components.
3. Method: Mc uses the learned variant representation Ψ(X) to generate environments, while Mp uses them to learn invariant features Φ(X) and an invariant predictor.The learned variant representation is then updated to further improve heterogeneity identification.
3. Method: HRM jointly optimizes heterogeneity identification and invariant learning so that improvements in one module can promote the other.The paper identifies this mutual promotion as the basis of the framework’s joint optimization process.
3.1. Implementation of Mp: Mp combines feature selection with variance-penalized invariant learning to select features whose correlations with Y remain stable across the learned environments.The objective averages environment risks and penalizes variance in environment-specific gradients.
3.1. Implementation of Mp: Soft feature selection replaces binary masks with clipped Gaussian gates in [0,1], reducing the high variance associated with hard feature-selection optimization.The gates are parameterized by μ and Gaussian noise, and the objective includes an ℓ0-related sparsity term.
3.1. Implementation of Mp: The prediction module obtains Φ(X) and Ψ(X) from the feature-selection parameters and is theoretically shown to learn the maximal invariant predictor for given environments.The implementation approximates the gradient-variance regularizer for linear models using a masked norm.
3.2. Implementation of Mc: Mc implements heterogeneous clustering by grouping data according to the relationship between Ψ and Y, using P(Y|Ψ) as the cluster centre.The initial variant representation is Ψ(X)=X, and cluster centres are modeled as Gaussian distributions around functions of Ψ.
3.2. Implementation of Mc: The clustering objective fits a mixture distribution to the empirical data, and EM estimates cluster centres and mixture weights before assigning points probabilistically to environments.The mixture weights lie on a K-dimensional simplex.

4. Theoretical Analysis

HRM’s theory analyzes invariant prediction and heterogeneity identification as mutually reinforcing modules. The results characterize when each module recovers invariant structure and explain why clustering learned variant features helps identify heterogeneous environments.

Invariant prediction: The invariant prediction module Mp learns the maximal invariant predictor for the training environments Etr.
Mutual promotion: The positive-feedback analysis assumes stable invariant features Φ∗ and relationships between variant features Ψ∗ and Y that can change across environments.
Heterogeneity identification: Excluding invariant features from clustering pulls cluster centres apart, making heterogeneous environments more likely to be identified.
Heterogeneity identification: Theorem 4.2 bounds the divergence between environment-conditional label distributions given X by the divergence given variant features Ψ∗.
Mutual promotion: Theorem 4.3 states that, under its assumptions, optimality of either Mc or Mp implies optimality of the other and establishes an HRM global optimum.

5. Experiment

Experiments evaluate HRM against ERM, DRO, EIIL, and IRM on simulated and real-world distribution shifts. HRM consistently improves out-of-distribution performance and stability, including settings without environment labels.

Experimental setup: HRM is compared with DRO, EIIL, IRM, and ERM, while an HRM variant without the feedback loop is used for ablation.
Experimental setup: The evaluation reports mean error, standard deviation error across test environments, and maximum error as the worst-case test-environment loss.
Simulation data: Simulation experiments vary selection bias, dimensions nb and d, and anti-causal-effect dimensions, with results averaged across repeated runs.
Simulation data: HRM achieves nearly uniformly good training and testing performance under simulated distribution shifts, while ERM, DRO, and EIIL perform poorly or inconsistently.
Real-world data: The real-world evaluation covers car insurance, personal income, and house price prediction tasks with environment shifts induced by attributes or time periods.
Real-world data: HRM achieves the best performance in 16 of 18 real-world testing environments and significantly outperforms IRM despite IRM receiving perfect environment labels.

6. Related Works

Related work on OOD generalization mainly follows distributionally robust optimization or invariant learning. DRO optimizes worst-case risk, while invariant learning seeks causally stable predictor–target relationships across environments.

Distributionally robust optimization: DRO methods optimize worst-case risk within an uncertainty set around the observed training distribution.
Invariant learning: Invariant learning assumes causally invariant relationships between predictors Φ(X) and the target Y and learns them from multiple training environments.

7. Discussions

The paper studies HRM as a joint approach to latent heterogeneity identification and invariant prediction, while identifying assumptions and unresolved theoretical limitations.

HRM jointly optimizes heterogeneity identification and invariant prediction without requiring environment labels.
The framework theoretically analyzes environment roles in invariant learning and proposes clustering for identifying heterogeneity in heterogeneous data.
The paper focuses on raw variables X = [Φ*, Ψ*]^T and leaves representation learning from X for future work.
HRM’s mutual promotion mechanism requires variant features Ψ* for heterogeneity identification, although invariant prediction tends to discard them.
The clustering algorithm lacks theoretical convergence guarantees because its relationship-based model dissimilarities are difficult to measure.

A. Additional Simulation Results and Details

Additional simulations define distribution-shift mechanisms through selection bias and anti-causal effects, varying spurious correlations, mixture structure, dimensions, and training conditions.

The simulations assume invariant covariates Φ* preserve P(Y|Φ*) across environments, while variant covariates Ψ* can vary in relation to Y.
Selection Bias: Selection bias experiments use a variable subset Vb ⊂ Ψ* and parameter r to control the strength and direction of spurious correlation with Y.
Selection Bias: Training and testing environments are generated by varying selection-bias conditions, sample composition, and variant dimensions, with repeated averaging across runs.
Anti-Causal Effect: The anti-causal scenario introduces spurious correlations through a relationship from target Y to variant covariates Ψ*.
The data-generation setup represents invariant and variant covariates with dimensions nφ and nψ and constructs Φ* from auxiliary variables Z.
Anti-Causal Effect: Mixture weights over Gaussian components define environments whose varying noise levels produce different Ψ*-Y correlation strengths.

B.1. Proof of Theorem 2.1

The proof establishes that a representation satisfying the invariant-learning assumptions is the maximal invariant predictor and supports a risk-minimization result under Bregman loss.

Theorem B.1 states that a representation Φ* satisfying Assumption 2.1 is the maximal invariant predictor.
The proof compares invariant representations using mutual information to establish the maximality characterization of Φ*.
The argument also proves that the maximal invariant predictor satisfies the sufficiency property required by Assumption 2.1.
Theorem B.2 considers a strictly convex differentiable function, its Bregman loss, and h*(X) = E[Y|Φ*] under Assumption 2.2.

B.3. Proof of Theorem 2.3

Theorem B.4 shows that adding an environment does not tighten the invariance set when its conditional target relationship already matches the existing environments.

If Pnew(Y|Φ̂) equals Pe(Y|Φ̂) for every existing environment, adding the new environment leaves the invariance set unchanged.
The proof explains that every representation invariant over the original environments remains invariant because the new environment excludes no original variables.

B.4. Proof of Theorem 4.1

The proof establishes that the learned representation Φ(X) = M ⊙ X is the maximal invariant predictor, while invariant information reduces environment-dependent conditional variation.

The learned Φ(X) = M ⊙ X is the maximal invariant predictor of IEtr.
For Φ(X), the conditional distribution P^e(Y | Φ(X)) remains invariant across environments.
The objective reaches its minimum when Φ(X) is the maximal invariant predictor, for λ ≥ 0.This combines sufficiency with the vanishing gradient-variance term.
If X contains invariant Φ* and variant Ψ*, conditioning on X yields no greater environment discrepancy than conditioning on Ψ* alone.The theorem bounds the KL divergence between environments using these two conditionings.

B.6. Proof of Theorem 4.3

Theorem B.7 connects the learned environments and representation: when training environments match the true environments, Mp recovers the maximal invariant predictor, and pooled data can recover the full invariance set.

When IE = IEtr, Mp learns the maximal invariant predictor of IE.
When pooled training data contains all environments in supp(E), learned-environment regularization makes IEtr equal to IE.
The mixture-distribution objective reaches its minimum when Q's components correspond to the distributions for environments e ∈ E.
Under the learned maximal invariant predictor, the corresponding variant representation is Ψ*(X), enabling the mixture components to recover the environment distributions.