PaperlayerView exact revision on arXiv

Source-linked AI summary

Invariant Information Bottleneck for Domain Generalization

Bo Li, Yifei Shen, Yezhen Wang, Wenzhen Zhu, Colorado J. Reed, Jun Zhang, Dongsheng Li, Kurt Keutzer, Han Zhao

arXiv:2106.06333v6cs.LGstat.ML

TL;DR

IRM can be difficult to optimize for nonlinear classifiers and can fail with pseudo-invariant features or geometric skews. The paper proposes IIB, using mutual-information formulations and a variational objective to constrain input–representation information. IIB improves over IRM in synthetic failure-mode settings and achieves a 0.7% average improvement across seven DomainBed datasets.

  • Problem

    IRM is difficult to optimize for nonlinear classifiers and can fail when pseudo-invariant features or geometric skews exist.

  • Method

    IIB formulates invariant causal prediction through mutual information, uses variational approximation for nonlinear classifiers, and constrains mutual information between inputs and representations.

  • Results

    0.7% average improvement across 7 DomainBed datasets, while synthetic experiments show IIB significantly improves performance over IRM in failure-mode settings.

  • Takeaways & Limitations

    IIB provides an information-theoretic domain-generalization objective that mitigates pseudo-invariant and spurious information while supporting nonlinear classifiers.

  • Takeaways & Limitations

    IRM may retain pseudo-invariant features when the number of environments is smaller than the number of spurious features, limiting feature identification.

Abstract

from arXiv · show

Invariant risk minimization (IRM) has recently emerged as a promising alternative for domain generalization. Nevertheless, the loss function is difficult to optimize for nonlinear classifiers and the original optimization objective could fail when pseudo-invariant features and geometric skews exist. Inspired by IRM, in this paper we propose a novel formulation for domain generalization, dubbed invariant information bottleneck (IIB). IIB aims at minimizing invariant risks for nonlinear classifiers and simultaneously mitigating the impact of pseudo-invariant features and geometric skews. Specifically, we first present a novel formulation for invariant causal prediction via mutual information. Then we adopt the variational formulation of the mutual information to develop a tractable loss function for nonlinear classifiers. To overcome the failure modes of IRM, we propose to minimize the mutual information between the inputs and the corresponding representations. IIB significantly outperforms IRM on synthetic datasets, where the pseudo-invariant features and geometric skews occur, showing the effectiveness of proposed formulation in overcoming failure modes of IRM. Furthermore, experiments on DomainBed show that IIB outperforms $13$ baselines by $0.9\%$ on average across $7$ real datasets.

Introduction

Domain generalization addresses failures caused by distribution shifts, while IRM remains difficult to optimize for nonlinear classifiers and can fail with pseudo-invariant features or geometric skews. The paper proposes IIB, an information-theoretic formulation that constrains input–representation mutual information and improves performance in these settings.

  • Distributional shifts violate the i.i.d. assumption and can cause models to deteriorate on samples from different domains.
  • IRM is difficult to optimize for nonlinear classifiers and can fail when pseudo-invariant features or geometric skews exist.
  • IIB is a novel information-theoretic formulation for domain generalization that minimizes invariant risks while mitigating pseudo-invariant features and geometric skews.
  • IIB formulates invariant causal prediction through mutual information and uses variational approximation to obtain tractable loss functions for nonlinear classifiers.
  • Constraining mutual information between inputs and representations targets the impact of pseudo-invariant features and geometric skews.
  • 0.7% consistent improvement on 7 datasets is reported for IIB on DomainBed, alongside better elimination of spurious information than existing domain-generalization methods.

Related Work

Related work spans data manipulation, domain adaptation, invariant representations, and invariant risk minimization. These approaches differ in whether they diversify data, align representations, or enforce invariant feature-conditioned label distributions.

  • Data-manipulation methods increase training-data diversity through augmentation, while domain randomization applies more complex image transformations.
  • Meta-Reg learns a meta-regularizer for the classifier, whereas Meta-VIB regularizes representation distributions across domains within a meta-learning framework.
  • Mutual-information approaches in domain adaptation seek representations informative about labels and invariant across domains.
  • IRM seeks invariant causal prediction by requiring feature-conditioned label distributions to remain consistent across environments.
  • The IRM objective is bilevel and difficult to optimize, motivating a first-order approximation with a dummy classifier.

Preliminaries

IRM has documented failure modes involving pseudo-invariant features and geometric skews, while Figure 1 contrasts how ERM, IRM, and IIB use features for OOD generalization. A structural causal model distinguishes causal from environmental parts of the input.

  • Feature selection in OOD generalization: In the bear illustration, ERM uses all five features, IRM uses the first three, and IIB selects only the shape or outline.The figure presents invariant features as yielding correct invariant predictions, while pseudo-invariant features can mislead test-domain generalization.
  • IRM failure modes: IRM can fail when pseudo-invariant features or geometric skews cause classifiers to use spurious features.Pseudo-invariant features may leak into recovered representations, while geometric skews can make spurious shortcuts attractive to max-margin classifiers.
  • Pseudo-invariant features: With fewer environments than spurious features, IRM may combine causal features with transformed spurious features and achieve lower training error.Those spurious features can vary arbitrarily at test time, harming OOD generalization.
  • Geometric skews: Geometric skews can make a spurious-feature classifier have smaller norm than an invariant classifier, leading to OOD generalization failure.The setup includes majority and minority groups whose spurious-feature correlations differ.
  • Structural causal model: The structural causal model divides input X into causal feature Zc and environmental feature Za, with both correlated with Y but only Zc treated as causal.The model is used to distinguish different causal relationships between input parts and the output.

Our Proposed Method

IIB formulates invariant causal prediction through mutual information, then adds an information-bottleneck penalty to select compressed representations and address IRM failure modes. Variational approximations yield tractable objectives for nonlinear classifiers, implemented with invariant, domain-dependent, and encoder components.

  • Invariant Causal Prediction via Mutual Information: IIB seeks representations satisfying conditional invariance while remaining informative about labels, formalizing invariant causal prediction through I(Y,D∣Z)=0.This implies that label predictions conditioned on the representation do not depend on the domain.
  • On the Failure Modes of IRM: Among representations satisfying invariant causal prediction, IIB selects the least-capacity solution by minimizing mutual information I(X,Z).The paper motivates this criterion as choosing the representation with the largest compression.
  • On the Failure Modes of IRM: The information bottleneck is motivated by IRM failure with pseudo-invariant features and geometric skews, which can cause spurious features to enter the learned classifier.Without the sparsity constraint, the analyzed IRM setting uses invariant, pseudo-invariant, and skew-related features for classification.
  • Our Proposed Method: The IIB objective combines the information bottleneck term I(Z,Y) − βI(X,Z) with the invariant-risk term I(Y,D∣Z).The first term encourages informative compression, while the second implements the IRM principle.
  • Loss Function Design: Variational approximation converts the mutual-information objective into tractable classification and encoder losses for nonlinear models.The construction uses an approximation r(z) to p(z), q(y∣z) to p(y∣z), and a stochastic encoder p(z∣x).
  • Loss Function Design: IIB optimizes an invariant predictor fi(Z), a domain-dependent predictor fd(Z,D), and an encoder g(X), with three losses corresponding to the three mutual-information terms.The predictors support the variational treatment of label information and conditional domain information.

Synthetic Experiments

Synthetic experiments evaluate IIB under pseudo-invariant features and geometric skews, where spurious correlations can undermine domain-generalization methods. Across CS-CMNIST and geometric-skew CIFAR10 settings, IIB improves robustness relative to competing approaches, including IRM and IB-IRM.

  • CS-CMNIST: CS-CMNIST tests whether models rely on colors that predict labels during training but become random in the test environment.The task uses two training environments with p_e values of 1 and 0.9, while the test environment has p_e = 0.
  • Geometric Skew CIFAR10: Geometric-skew CIFAR10 includes Cross Lines and Vertical Line experiments that introduce class-correlated or scale-varying spurious lines.Cross Lines uses ten-valued line configurations, while Vertical Line varies a colored line’s spurious feature scale across domains.
  • Methods: IIB is compared with ERM, IRM, and IB-IRM on CS-CMNIST, with IB-IRM combining information bottleneck and IRM to address geometric skews.The comparison distinguishes IIB’s objective optimization from the direct IB-IRM combination.
  • Results: IIB shows observable improvements over other algorithms across the two synthetic datasets.In Vertical Line experiments, IIB retains good results despite large offsets in spurious feature scale, whereas training and testing accuracy otherwise decreases substantially.
  • Results: In Cross Lines experiments, information bottleneck structure helps mitigate the geometric-skew failure observed with IRM.The comparison uses colored cross-line training images and normal test images.

DomainBed Experiments

DomainBed evaluates IIB across seven datasets using training-domain validation and leave-one-domain-out model selection. IIB achieves the best average performance, with stronger results on larger datasets despite no method consistently dominating on small datasets.

  • Experimental Setup: The experiments cover 7 different datasets of different sizes in DomainBed.The study uses DomainBed to empirically evaluate IIB across datasets with varying scale.
  • Model Selection Strategy: The evaluation uses validation subsets of training domains rather than test-domain validation during training.This model-selection strategy is intended to characterize in-distribution generalization capability.
  • Results: IIB achieves the best average performance across 7 datasets under leave-one-domain-out model selection.The comparison reports accuracy for IIB and state-of-the-art domain generalization methods, along with the average across datasets.
  • Results: IIB performs better than other methods on larger datasets, including PACS, Office-Home, and DomainNet.These gains produce its stronger average performance across the seven datasets.
  • Results: No domain generalization method consistently dominates on the small Colored-MNIST and Rotated-MNIST datasets.This observation is reported as consistent with prior DomainBed findings.

Conclusion

The paper develops invariant information bottleneck (IIB) to address IRM’s limitations in domain generalization. Its variational objective targets nonlinear classifiers and seeks minimum-capacity features, with superior performance demonstrated on synthetic and real datasets.

  • Conclusion: IIB addresses IRM limitations by seeking minimum-capacity features among those satisfying the original invariant-risk objective.The approach is motivated by insufficient training domains for identifying every potential spurious feature.
  • Conclusion: A variational approach makes the IIB objective applicable beyond IRM’s gradient-penalty formulation for linear classifiers.The paper contrasts the proposed optimization approach with the earlier formulation’s linear-classifier restriction.
  • Conclusion: IIB’s superior performance is demonstrated on both synthetic and real datasets.The conclusion summarizes evidence across both types of evaluation settings.
Loading 2106.06333v6…