PaperlayerView exact revision on arXiv

Source-linked AI summary

Environment Inference for Invariant Learning

Elliot Creager, Jörn-Henrik Jacobsen, Richard Zemel

arXiv:2010.07249v5cs.LGcs.AI

TL;DR

Distribution shifts create a need for models that rely on domain-invariant rather than domain-specific features, but invariant learning commonly assumes environment labels are available. EIIL infers informative environment partitions from aggregated training data and uses them for downstream invariant learning. The paper reports improvements over labeled-environment invariant learning on CMNIST, over ERM on Waterbirds and CivilComments worst-group performance, and in accuracy and calibration for fair prediction.

  • Problem

    Invariant learning commonly requires environment partitions, but such labels may be unavailable, difficult to obtain, or unclear to define.

  • Method

    EIIL infers environment assignments from the bias of a reference classifier and uses them to train an invariant learner.

  • Results

    EIIL can outperform or approach invariant learning methods requiring environment labels and can outperform ERM when spurious features or subgroup performance discrepancies affect standard learning.

  • Takeaways & Limitations

    EIIL provides inferred environments for invariant learning in domain-generalization and fairness settings where meaningful environments or sensitive subgroups are difficult to specify.

  • Takeaways & Limitations

    Invariant-learning guarantees and performance remain constrained by restrictive assumptions, nonlinear objectives with multiple optima, and inconsistent empirical advantages over ERM under careful model selection.

Abstract

from arXiv · show

Learning models that gracefully handle distribution shifts is central to research on domain generalization, robust optimization, and fairness. A promising formulation is domain-invariant learning, which identifies the key issue of learning which features are domain-specific versus domain-invariant. An important assumption in this area is that the training examples are partitioned into "domains" or "environments". Our focus is on the more common setting where such partitions are not provided. We propose EIIL, a general framework for domain-invariant learning that incorporates Environment Inference to directly infer partitions that are maximally informative for downstream Invariant Learning. We show that EIIL outperforms invariant learning methods on the CMNIST benchmark without using environment labels, and significantly outperforms ERM on worst-group performance in the Waterbirds and CivilComments datasets. Finally, we establish connections between EIIL and algorithmic fairness, which enables EIIL to improve accuracy and calibration in a fair prediction problem.

1. Introduction

Distribution shifts can sharply degrade supervised models, while invariant learning often requires environment labels that are unavailable or difficult to define. EIIL infers informative environments from training data using a reference classifier, enabling invariant learning without provided environment assignments.

  • Distribution shifts can severely degrade model performance, including under small perturbations and newly collected test sets.The paper motivates methods that generalize beyond the training distribution.
  • Invariant learning seeks features that generalize across domains but typically requires training examples to be partitioned into environments.Environment assignments specify the variation to which the model should become invariant or robust.
  • Environment labels may be unavailable because they are difficult to obtain, restricted by privacy limitations, or unclear to specify from available metadata.The same challenge arises when defining sensitive attributes for fairness applications.
  • EIIL infers environment partitions from the bias of an ERM-trained reference model and then trains an invariant learner using those inferred environments.Its environment-inference objective is differentiable with respect to a distribution over environment assignments.
  • EIIL discovers environments that capture spurious correlations and can be used with off-the-shelf invariant learning algorithms when environment labels are unavailable.The method can also improve invariant learning when manually specified environments exist.
  • EIIL connects invariant learning with fairness and improves accuracy and calibration in a fair prediction problem.The paper also reports that its inferred-environment approach outperforms invariant learning with environment labels on CMNIST and ERM on Waterbirds.

2. Invariant Learning

Invariant learning targets representations whose predictive relationships remain stable across environments, complementing robust optimization's focus on worst-case perturbations. Its effectiveness is constrained by restrictive assumptions, nonlinear optimization issues, and inconsistent empirical advantages over ERM.

  • Domain generalization seeks low error on unseen test distributions, while invariant learning aims to discover relationships that remain stable across environments.Invariant representations encode features inducing the same conditional label distribution across environments.
  • Invariant representations require an additional predictive objective because trivial constant representations can satisfy environment invariance.IRM introduces a regularized objective requiring the same classifier to be optimal in every environment.
  • Robust optimization minimizes worst-case loss over specified perturbation sets or subsets of the training data rather than directly optimizing invariance.GroupDRO uses environment labels to share importance weights across examples in relevant groups.
  • Environment inference enables GroupDRO to improve over standard learning without requiring group labels.This extends robust optimization to settings where relevant groups are not directly provided.
  • IRM can generalize out of distribution only under restrictive assumptions, including linear data generation and access to many environments.Most benchmark datasets are described as nonlinear, limiting the direct scope of those guarantees.
  • For some nonlinear datasets, the IRMv1 penalty has multiple optima, and invariant learners have not consistently outperformed ERM under careful model selection.The paper emphasizes model selection as important for domain-generalization approaches.

3. Invariance Without Environment Labels

EIIL infers environment assignments without a priori environment knowledge by maximizing violations of invariance for a fixed reference classifier, then trains an invariant learner on the inferred environments. The framework can identify spurious-feature partitions and improve invariant learning, although some guarantees depend on the reference model.

  • Environment Inference and Invariant Learning: EIIL uses a two-stage process: infer environment assignments, then run invariant learning using those assignments.Environment inference maximizes violations of the invariant-learning principle before invariant learning is performed.
  • Environment Inference and Invariant Learning: Soft environment assignments q(e′|x_i, y_i) are optimized for worst-case environments while the reference classifier Φ̃ remains fixed.The soft per-environment risk is differentiable with respect to q, allowing gradients to update environment assignments.
  • Limitations: Jointly training the representation and environment assignments showed no empirical benefit and introduced optimization and conceptual difficulties.The formulation must ensure that invariances apply to all environments discovered throughout learning.
  • Spurious Features: The inferred environments can separate examples according to spurious-feature behavior, enabling invariant representations that ignore those features.When environments differ in agreement between labels y and spurious features z, a representation ignoring z satisfies the invariance constraint.
  • Empirical and Theoretical Scope: EIIL improves IRM test accuracy over CMNIST’s hand-crafted environments, while Waterbirds results show effective invariant learning without a guarantee of maximal violation.The Waterbirds case involves a reference model focused on a mixture of spurious and valuable features.
  • Binned Environment Invariance: Binary-class binning maximizes the invariance violation by assigning y = 1 and y = −1 examples to different environments within each representation bin.This procedure yields expected labels of ±1 per bin and reaches the maximum possible ΔEIC per bin.

4. Related Work

Related work connects invariant learning and fairness through statistics evaluated across environments or sensitive groups, while contrasting representation-space invariance with loss-space robustness. EIIL also relates to methods that exploit reference-model errors or inferred groups, but differs in its downstream invariant-learning objective.

  • Invariant Learning: IRM addresses test-time label shift, while REx prefers equalized risk across environments by penalizing risk variance.REx’s practical V-REx algorithm is reported to improve generalization in several settings.
  • Algorithmic Fairness: Recent work also extends invariance methods to class-conditioned predictive distributions and relates them to equalized odds in fair classification.The cited approach trains on environments inferred by EIIL for systematic generalization.
  • Reference Models and Inferred Groups: Several methods exploit a pretrained or ERM reference model’s mistakes or confidence to infer groups, assign importance weights, or train GroupDRO models.These approaches overlap with EIIL’s use of a reference classifier but differ in how inferred information enters the primary objective.
  • Algorithmic Fairness: Domain generalization and fairness methods can be viewed as matching or optimizing statistics across a conditioning variable representing environments or sensitive-group membership.The comparison uses learned vector and scalar functions of the inputs, denoted Φ and S.
  • Algorithmic Fairness: EIIL is related to fairness methods that seek good performance for worst-case group assignments under a fairness criterion.The paper describes environment inference as serving a similar purpose in domain generalization and fairness.
  • Robustness and Invariance: Adversarially Reweighted Learning prioritizes stability in loss space, whereas EIIL encourages invariance in learned representation space.The paper presents these as complementary approaches suited to different types of distribution shift.

5. Experiments

The experiments evaluate EIIL across CMNIST, Waterbirds, Adult-Confounded, and CivilComments-WILDS, using inferred environments for invariant learning without pre-specified group labels. Results show improvements over standard baselines, while analysis highlights the importance of the reference model and difficult model-selection conditions.

  • Experimental caveat: Hyperparameter results are optimistic because prior invariant-learning studies allowed limited validation samples for tuning despite differing train and test distributions.The paper describes model selection under distribution shift as a difficult open problem.
  • Datasets and setup: The experiments cover noisy digit classification, high-dimensional bird images, fair prediction, and toxicity prediction under distribution shift.CMNIST reverses color-label correlation at test time; Waterbirds and Adult-Confounded reverse spurious correlations, while CivilComments-WILDS evaluates toxicity prediction with fairness implications.
  • CMNIST: EIIL’s inferred environments outperform standard IRM using provided labels on CMNIST, despite label noise with θy = 0.25.The experiments discard the original environment labels and use ERM for inference followed by IRM; EIIL’s labels are reported as more effective than the dataset’s provided labels.
  • Waterbirds: 18% higher worst-group test accuracy with only a 1% drop in average accuracy distinguishes EIIL from the ERM baseline on Waterbirds.The inferred environments place minority subgroups—landbirds on water and waterbirds on land—mostly together, and GroupDRO uses these inferred labels.
  • Adult-Confounded: EIIL performs favorably against ARL and ERM on the Adult-Confounded audit test set and achieves significantly better calibration without sensitive group information.The method first infers worst-case environments and then enforces invariance across them; the authors associate the result with favorable group sufficiency.
  • CivilComments-WILDS: EIIL improves CivilComments-WILDS worst-group accuracy over ERM with a modest average-accuracy cost, approaching oracle GroupDRO without group labels.The CivilComments-WILDS setup uses an ERM reference classifier and feeds inferred environments to GroupDRO.
  • Method: EIIL uses ERM as a reference model to infer environments, then applies an invariant learner using those inferred labels.The downstream learner is IRM for CMNIST and GroupDRO for Waterbirds and CivilComments-WILDS.
  • Reference model analysis: EIIL generalizes better than IRM when CMNIST label noise satisfies θy > .2, but generalizes poorly under low label noise.The analysis attributes this dependence to whether ERM learns the color or shape feature, and identifies reference-model choice as an open question.

6. Conclusion

EIIL infers environment partitions from aggregated training data, enabling invariant learning when environment labels are unavailable. It can outperform or approach labeled-environment methods and connects domain generalization with fairness.

  • EIIL infers environment partitions of aggregated training data for invariant learning without access to environment labels.
  • EIIL can outperform or approach invariant learning methods that require environment labels.
  • EIIL has implications for domain generalization and fairness because meaningful environments or sensitive subgroups can be difficult to specify.

B.1. Proof of Proposition 1

The proof shows that EIIL can infer environment assignments that maximize the sufficiency gap and therefore maximally violate invariant prediction for spurious-feature-based classifiers. In CMNIST, separating examples by color–label agreement achieves the maximal gap of 1.

  • Proposition 1: A classifier predicting the spurious feature z maximizes the group-sufficiency gap when environments separate agreement from disagreement between z and y.The corresponding environments are e1 = {y = z} and e2 = {y ≠ z}.
  • Proposition 1: Within each spurious-feature value, conditional label expectations differ by 1 across agreement-based environments, yielding the maximum sufficiency gap Δ(S, e) = 1.This follows separately for z = 0 and z = 1.
  • EIIL objective: A global optimum of relaxed EIIL using the IRMv1 regularizer maximally violates the invariance principle.The heuristic is designed to maximize both the softened regularizer violation and the invariance-principle violation.
  • EIIL objective: EIIL maximizes a softened IRMv1 violation by assigning examples with y = 1 to one environment and examples with y = 0 to the other.Within representation bins, this assignment makes the label distributions as different as possible.
  • CMNIST: For CMNIST, a color-based classifier exposes a partition with sufficiency gap 1, showing that the described color-based environment split is maximal.The inferred partition groups examples according to whether their labels agree with the image color.

C. Connections Between Invariant Learning and Algorithmic Fairness

The paper connects invariant learning to algorithmic fairness through shared conditional-statistic constraints across environments or protected groups. In particular, IRMv1 minimizes group sufficiency gaps for scalar representations, extending related fairness analysis beyond convex losses.

  • Shared objectives: Domain generalization and algorithmic fairness both match a chosen statistic across a conditioning variable representing environments or sensitive group membership.The statistic informs the learning objective, although its motivation differs between the fields.
  • Shared objectives: Fairness notions such as demographic parity and equal opportunity correspond to marginal or conditional independence between predictions and group membership.These relationships motivated statistical-independence regularizers in fair-representation learning.
  • Group sufficiency: Group sufficiency studies conditional label expectations given a score and protected group, closely paralleling the invariant conditional expectation constraint.For binary classification, score-based and representation-based formulations are closely related.
  • Group sufficiency: IRMv1 minimizes the group sufficiency gap for scalar representations when the environment variable denotes sensitive group membership.This observation applies to neural-network representations, whereas Liu et al.'s theoretical results rely on convex losses.
  • Scope: The paper leaves empirical study of methods using groups without a specified fairness criterion for future work.This marks a scope boundary in the fairness connection.

D. Dataset details

The dataset section describes CMNIST, Waterbirds, CivilComments-WILDS, and Adult-Confounded as evaluation settings involving spurious correlations or distribution shifts. Adult-Confounded specifically exaggerates subgroup–label correlations and reverses them at test time.

  • CMNIST: CMNIST uses two 25,000-image training environments, a test correlation of 0.1, and label noise with probability θ_y = 0.25.The training environments are described as having color–label correlation 0.8.
  • Waterbirds: Waterbirds uses train, validation, and test splits of 4,795, 1,200, and 5,794 examples, respectively.The validation and test distributions upweight minority groups.
  • CivilComments-WILDS: CivilComments-WILDS uses the train, validation, and test splits from Koh et al. (2021).The paper refers readers to that work for dataset details and differences from the original dataset.
  • Adult-Confounded: Adult-Confounded reverses the subgroup–label relationship in its audit test set after observing that subgroup membership can strongly predict labels.The dataset is constructed as a semi-synthetic variant of UCI Adult.
  • Adult-Confounded: Adult-Confounded exaggerates spurious subgroup–label correlations to emphasize test-time distribution shift and unequal-treatment concerns.Its four subgroups combine binary sex and binarized race attributes.

E. Experimental details

The experiments compare EIIL with invariant and robust-learning baselines across synthetic, image, and text datasets, using dataset-specific training and model-selection procedures. EIIL adds reference-model training and environment inference, with overhead that becomes relatively small as problem size grows.

  • Model selection: Adult-Confounded model selection tunes reference-model and invariant-learner hyperparameters jointly, while other experiments use specified defaults or author-reported settings.The Adult-Confounded strategy includes 20 test evaluations per method; Waterbirds uses reported hyperparameters and validation-based selection.
  • Training times: EIIL requires a pretrained reference model and environment-inference optimization, making training longer than standard invariant learning.The additional time primarily depends on reference-model training and EI optimization steps.
  • Training times: 6.6 minutes versus 2.18 minutes for ERM and 2.20 minutes for IRM: CMNIST EIIL training with 10 restarts on a Tesla P100.The CMNIST reference model used 1,000 steps, while EI optimization used 10,000 steps.
  • Training times: 0.4% relative increase: Waterbirds EIIL took 4.737 hours versus 4.716 hours for GroupDRO on a Tesla P100.The relative overhead is discounted as problem size scales.
  • Environment inference: Environment inference aggregates logits for the full training set before optimizing the EI objective, while minibatched inference and amortized assignments remain future work.Network activations are discarded after logits are aggregated.
  • Dataset procedures: For CivilComments-WILDS, EIIL uses an ERM reference classifier and feeds inferred environments to GroupDRO; error splitting replaced gradient-based inference.The authors found the faster heuristic sufficient and promising for scaling to large datasets.

F.1. Synthetic Data

The synthetic regression study separates causal and non-causal features and evaluates whether methods recover the causally correct solution without handcrafted environments. EIIL uses an ERM reference model to infer environments for downstream IRM, improving the resulting invariant solution.

  • Results: EIIL-discovered environments enable IRM to outperform IRM with handcrafted environments in synthetic regression.Table 7 reports mean squared error plus standard deviation across five runs.
  • Synthetic Data: The data combine causal feature v and non-causal feature z, with noise varying across handcrafted environments.The feature vector is x = [v, z], with each component occupying half the dimensions.
  • Methods: The evaluated methods are ERM, IRM, ICP, and EIIL, with ERM avoiding environment labels and IRM and ICP using handcrafted labels.EIIL infers useful environments from ERM and then applies IRM.
  • Evaluation: The causally correct regression solution is w* = [1, 0], corresponding to prediction from the causal feature rather than the non-causal feature.Performance is assessed through errors in the causal and non-causal dimensions.
  • Results: Because v is marginally noisier than z, ERM focuses on spurious z; EIIL instead learns worst-case environments for invariant training.IRM with handcrafted environments exploits variability in the non-causal feature's noise level.
  • Reference-model bias: EIIL outperforms IRM when its reference classifier emphasizes the spurious feature, including ERM or an α-SPURIOUS model with high α.Figure 5 compares mean squared errors for causal and non-causal features.

F.2. ColorMNIST

The ColorMNIST experiments compare EIIL with label-free and oracle-style baselines, examining whether inferred environments support invariant prediction beyond color. The study also connects environment inference to subgroup sufficiency, calibration, and the necessity of both EIIL stages.

  • Methods: EIIL discovers environments without environment labels and uses them for invariant learning in the ColorMNIST benchmark.The additional baselines include grayscale, color-based, binned-inference, GEORGE, ARL, and LFF methods.
  • Results: Unlike EIIL, the new unsupervised baselines fail to find an invariant classifier that predicts from shape rather than color.The authors hypothesize that GEORGE's hidden-subclass assumption prevents uncovering optimal environments for GroupDRO.
  • Baseline analysis: LFF is highly sensitive to q; at q = 0.7 it performs optimally on the test set by learning an anti-color classifier rather than a shape-based invariant classifier.Below-chance training performance reveals the anti-color behavior.
  • Environment inference: The binned EI heuristic effectively separates errors from correct predictions because the high-confidence reference classifier populates only two bins.This makes binning equivalent to assigning errors and correct predictions to separate environments.
  • Fairness connection: EIIL improves test calibration and accuracy on the Adult-Confounded dataset under a test shift that reverses subgroup-label correlations.The authors interpret this robustness as evidence that EIIL does not rely on subgroup membership for prediction.
  • Subgroup sufficiency: EIIL regularizes inferred environments toward similar calibration profiles, improving subgroup sufficiency relative to ARL without using sensitive group information.ARL shows noticeable calibration discrepancies across groups, whereas EIIL implicitly regularizes the inferred environments.
  • Ablation: Both worst-case environment inference and IRMv1 invariance regularization are required for good test-time performance on Adult-Confounded.The ablation study isolates both ingredients as necessary components of EIIL.
Loading 2010.07249v5…