PaperlayerView exact revision on arXiv

Source-linked AI summary

Fishr: Invariant Gradient Variances for Out-of-Distribution Generalization

Alexandre Rame, Corentin Dancette, Matthieu Cord

arXiv:2109.02934v3cs.LGcs.AIcs.CV

TL;DR

OOD generalization remains difficult because existing domain-invariance methods have not shown systematic benefits over ERM under controlled evaluation. Fishr matches domain-level variances of individual loss gradients, connecting this regularization to risks and Hessians. It achieves state-of-the-art DomainBed performance and systematically outperforms ERM on the evaluated real datasets.

  • Problem

    Existing domain-invariance approaches have not shown systematic benefits over ERM under controlled evaluation, despite the need for models robust to distribution shifts.

  • Method

    Fishr is a scalable regularizer that matches domain-level variances of per-sample loss gradients across training domains.

  • Results

    Fishr reaches state-of-the-art performance on DomainBed and is systematically better than ERM across PACS, VLCS, OfficeHome, TerraIncognita, and DomainNet.

  • Takeaways & Limitations

    Matching gradient variances provides a proxy for matching domain-level risks and Hessians in OOD classification.

  • Takeaways & Limitations

    The empirical-Fisher approximation is theoretically justified only under strong assumptions, so the paper trades theoretical guarantees for efficiency.

Abstract

from arXiv · show

Learning robust models that generalize well under changes in the data distribution is critical for real-world applications. To this end, there has been a growing surge of interest to learn simultaneously from multiple training domains - while enforcing different types of invariance across those domains. Yet, all existing approaches fail to show systematic benefits under controlled evaluation protocols. In this paper, we introduce a new regularization - named Fishr - that enforces domain invariance in the space of the gradients of the loss: specifically, the domain-level variances of gradients are matched across training domains. Our approach is based on the close relations between the gradient covariance, the Fisher Information and the Hessian of the loss: in particular, we show that Fishr eventually aligns the domain-level loss landscapes locally around the final weights. Extensive experiments demonstrate the effectiveness of Fishr for out-of-distribution generalization. Notably, Fishr improves the state of the art on the DomainBed benchmark and performs consistently better than Empirical Risk Minimization. Our code is available at https://github.com/alexrame/fishr.

1. Introduction

Fishr addresses OOD generalization by matching domain-level variances of per-sample loss gradients, motivated by the need for invariance under distribution shifts. The method theoretically connects gradient variance to domain-level risks and Hessians, and empirically improves performance across benchmark settings.

  • Motivation: Distribution shifts can make neural networks’ reliance on simple training features harmful when more complex features are needed.Examples include medical models exploiting patient age or body position rather than pathology.
  • Motivation: Existing domain-invariance methods have not shown significant advantages over ERM under controlled evaluation protocols.Prior approaches align features or seek classifiers that are simultaneously optimal across domains.
  • Method: Fishr matches the second moment—domain-level variances—of per-sample loss gradients across training domains.Unlike Fish, which matches gradient means, Fishr operates on gradient distributions’ variances.
  • Method: Fishr is theoretically motivated by relations among gradient variance, Fisher Information, and Hessians, linking the regularizer to locally aligned domain loss landscapes.The paper states that Fishr matches domain-level risks and Hessians and reduces cross-domain inconsistencies.
  • Implementation: The regularizer is made scalable by applying the gradient approximation in the classifier, with implementation support from BackPACK.The paper presents this as a computational-cost reduction.
  • Results: Fishr performs best on DomainBed and is systematically better than ERM across all evaluated real datasets.The reported datasets are PACS, VLCS, OfficeHome, TerraIncognita, and DomainNet.

2. Context and Related Work

The paper situates Fishr within OOD classification methods that seek invariant mechanisms across multiple domains. It distinguishes gradient-variance matching from feature, risk, and gradient-mean approaches, while retaining within-domain individual-gradient information discarded by batch averaging.

  • Problem definition: The task is classification under diversity and correlation shifts, with the goal of predicting well on an unseen test domain.The model is a feature extractor followed by a dense linear classifier, trained on multiple domains.
  • Problem definition: OOD generalization methods seek domain-invariant mechanisms because training and test distributions may differ in domain diversity or covariate correlations.The paper frames invariant mechanisms as desirable under these shifts.
  • Feature-based methods: Feature-alignment approaches include adversarial methods and covariance matching such as CORAL, but related work reports insufficient guarantees for good generalization.CORAL minimizes the distance between feature covariance matrices across domains.
  • Risk-based methods: Invariant Risk Minimization seeks predictors that are simultaneously optimal across domains, but recent work identifies pitfalls for nonlinear data.The paper presents this as a limitation of the related approach.
  • Gradient-based methods: Earlier gradient-based methods align domain gradient means, directions, or agreements, including IGA, Fish, and AND-mask.These methods differ in how they constrain gradients across domains.
  • Fishr’s distinction: Per-domain batch averaging discards granular pairwise interactions among same-domain sample gradients, whereas Fishr preserves individual-gradient information and matches gradient variances.Fishr is described as covariance matching in gradient space rather than feature space.

3. Fishr

Fishr regularizes domain-level gradient variances to reduce inconsistencies between domains and align their local loss landscapes. Its theory connects this variance matching to similar risks and Hessians, while acknowledging approximation and computational trade-offs.

  • Gradient-variance regularization: Unlike gradient-mean methods, Fishr matches the second moment of gradient distributions, capturing disagreements that average learning directions cannot.Previous methods match first moments, whereas Fishr targets domain-level variances.
  • Gradient-variance regularization: Fishr matches domain-level gradient variance vectors, complementing the empirical risk minimization objective with a distance-to-mean variance penalty.The penalty is balanced by λ > 0; variances may be computed on all weights or a subset to reduce memory and training costs.
  • Inconsistency and theory: Fishr minimizes a generalized inconsistency bound by simultaneously matching domain-level risks and Hessians under the paper’s stated assumptions.The analysis uses local loss-landscape geometry and assumes simultaneous convergence; Proposition 1 further assumes quadratic bowls with positive definite Hessians.
  • Consequences: Fishr’s gradient-variance matching induces similar domain-level risks and Hessian diagonals, thereby aligning local loss landscapes and reducing domain inconsistencies.The paper reports these effects at convergence and concludes that they increase domain generalization.
  • Hessian alignment: Matching gradient variances is motivated as an efficient proxy for aligning Hessian diagonals, reducing the targeted components from |θ|^2 to |θ|.The proxy relies on empirical relations among gradient covariance, the empirical Fisher Information Matrix, and the Hessian.
  • Limitations: The empirical-Fisher approximation is theoretically justified only under strong assumptions, so Fishr trades theoretical guarantees for computational efficiency.The paper also notes that Hessian computation is impractical as a training objective because it is substantially more computationally intensive.

4. Experiments

Experiments evaluate Fishr on Colored MNIST and DomainBed, showing improved out-of-distribution generalization, favorable computational cost, and consistent gains over ERM.

  • 4.1. Proof of concept on Colored MNIST: Colored MNIST tests whether models ignore spuriously correlated colors when the correlation reverses at test time.The two training domains have color-label correlations of 90% and 80%; a shape-based oracle reaches 75% test accuracy, while a color-only model reaches 10%.
  • 4.1. Proof of concept on Colored MNIST: 71.2% test accuracy is achieved by Fishrθ, while Fishrω reaches 69.5% with substantially lower computational cost.Fishrφ performs best in test accuracy but has lower train accuracy; the results support applying gradient variance matching at different network layers.
  • 4.1. Proof of concept on Colored MNIST: After activation at epoch 190, Fishrθ aligns domain-level gradient variances, risks, and Hessians while reducing train accuracy and increasing test accuracy.The synthetic dynamics support the theoretical link between gradient variance matching and aligned local loss landscapes.
  • 4.2. DomainBed benchmark: DomainBed evaluates methods across synthetic and real multi-domain image-classification datasets under controlled hyperparameter selection.The benchmark includes Colored MNIST, Rotated MNIST, VLCS, PACS, OfficeHome, TerraIncognita, and DomainNet.
  • 4.2.2. Implementation details: Fishr adds modest overhead relative to ERM: 0.2% memory and 2.7% training time on PACS, versus 7.0% and 6.5% on DomainNet.The classifier-only approximation is used on DomainBed because storing individual feature-extractor gradients is computationally infeasible.
  • 4.2.3. Results: Fishr is the only method reported to perform systematically better than ERM across all five real DomainBed datasets.Reported comparisons include VLCS at 78.2% versus 77.6%, OfficeHome at 68.2% versus 66.4%, and DomainNet at 41.8% versus 41.3%.
  • 4.2.3. Results: Fishr achieves the best average DomainBed accuracy at 70.8%, compared with 69.2% for CORAL, and has a median ranking of second.Excluding Colored MNIST, Fishr averages 71.1 versus 70.5 for ERM across the six remaining datasets.

5. Conclusion

The paper proposes Fishr as a simple gradient-variance regularizer for computer-vision OOD classification and reports state-of-the-art DomainBed performance.

  • 5. Conclusion: Fishr matches gradient variances across domains as a proxy for matching domain-level risks and Hessians.The paper proves that this reduces cross-domain inconsistencies and provides an open-source implementation.
  • 5. Conclusion: Fishr reaches state-of-the-art performance on DomainBed when samples from the test domain are available for model selection.The authors suggest consistent improvements for deep classifiers trained on data from multiple domains.
  • 5. Conclusion: The paper targets out-of-distribution generalization for classification in computer vision.The conclusion frames Fishr as a gradient-based regularization approach for improving deep-network generalization.

A. Additional Theoretical Analysis

The additional analysis connects domain inconsistency to risks and local curvature, and interprets gradient-variance matching through feature-adaptive error weighting and kernel geometry.

  • A.1. Demonstration of Proposition 1 from Section 3.2.1: Under the quadratic-bowl assumption around θ∗, local domain risks are represented by Hessian quadratic forms.The analysis treats each domain risk as remaining within an ϵ interval around its value at θ∗.
  • A.1. Demonstration of Proposition 1 from Section 3.2.1: The proof shows that the maximum inconsistency can be expressed through R(A, B) + Hϵ(A, B) under the stated assumptions.The argument handles separately the cases where R(A, B) is nonnegative or negative.
  • A.1. Demonstration of Proposition 1 from Section 3.2.1: The inconsistency Iϵ(A, B) decomposes into a risk difference R(A, B) and a curvature term Hϵ(A, B).The curvature term compares quadratic loss changes under the Hessians of domains A and B.
  • A.2. Fishr as a feature-adaptive version of V-REx: For a linear classifier, matching bias-gradient variances matches domain mean squared errors and therefore domain risks.This connects Fishr to V-REx, with squared error replacing negative log likelihood.
  • A.2. Fishr as a feature-adaptive version of V-REx: Weight-gradient variance matching corresponds to matching feature-weighted squared errors across domains.Each feature induces a distinct weighting scheme, making Fishr a feature-adaptive version of V-REx; binary features yield subgroup-specific risk matching.
  • A.2. Fishr as a feature-adaptive version of V-REx: Fishr matches centered gradient variances, which are equivalent to uncentered variances at convergence when domain mean gradients are approximately zero.Experiments report similar performance with centered and uncentered variances.
  • A.3. Gradient covariance and NTK: Similar NTK spectra across domains would align learned mechanisms and facilitate simultaneous convergence during optimization.Similar top eigenvectors discourage domain-dependent shortcuts, while similar top eigenvalues support choosing a common learning rate.
  • A.4. Gradient covariance and Hessian: The analysis uses gradient covariance as a proxy for Hessian regularization, while noting that formal approximation bounds are outside the paper’s scope.The correspondence is empirically validated at convergence and during training.

B. Experiments on a Linear Example

Fishr shifts a linear classifier away from domain-specific features toward the invariant feature, improving performance under changing correlations.

  • 93% test accuracy with Fishr matches its train accuracy, compared with 57% test accuracy for ERM.ERM assigns large weights of 3.3 to the more predictive but spurious features f2 and f3, whereas Fishr reduces them to 1.2.
  • Fishr makes the linear model rely mostly on invariant feature f1 when spurious feature-label correlations vary across domains.
  • The example uses four static features, with f1 invariant and f2–f4 spurious across the two training domains and test domain.

C.1. Description of the Colored MNIST experiment

Colored MNIST tests whether models can ignore a training correlation between digit color and labels when that correlation reverses at test time.

  • Colored MNIST collapses digits into two classes with 25% label flipping and colors digits red or green.
  • Training domains associate green color with labels at 90% and 80%, whereas the test-domain association reverses to 10%.
  • Ignoring color and using digit shape would yield 75% test accuracy under the shifted correlation.
  • The experiment uses a three-layer ReLU MLP optimized with Adam and retains IRM’s specified hyperparameters.The setup includes 501 epochs, batch size 25,000, and a penalty warmup at iteration 190.

C.2.1. HESSIAN MATCHING

Fishr aligns domain-level Hessians and gradient variances during training, whereas ERM leaves cross-domain discrepancies and color reliance largely intact.

  • C.2.1. HESSIAN MATCHING: After λ increases at epoch 190, Fishr brings domain-level Hessians close in distance, norms, and directions.The Hessian norms remain high, indicating sharp loss landscapes at convergence.
  • C.2.1. HESSIAN MATCHING: Under ERM, the distance between domain-level Hessians increases, their cosine similarity decreases, and converged loss landscapes are flat.
  • C.2.1. HESSIAN MATCHING: ERM’s domain-level gradient-variance and Hessian distances keep increasing, while risk distance decreases more slowly than with Fishr.
  • C.2.2. COLORED MNIST WITHOUT LABEL FLIPPING: Without label flipping, Fishr and V-REx reach 95.3%, exceeding ERM at 91.8% and IRM at 82.2%.This experiment supports Fishr’s effectiveness when digit shape fully predicts the label.
  • C.2.3. GRADIENT VARIANCE OR COVARIANCE ?: Matching only diagonal covariance components produces similar or slightly better results than using the full covariance matrix.The full-covariance comparison is feasible only when gradients are taken in the classifier because of memory constraints.
  • C.2.3. GRADIENT VARIANCE OR COVARIANCE ?: Centered gradient covariance and uncentered empirical Fisher diagonals perform almost equivalently in the reported comparison.

D.1. Description of the DomainBed benchmark

The DomainBed benchmark evaluates domain generalization across seven multi-domain vision datasets using standardized hyperparameter searches and multiple baselines.

  • Each method receives a 20-trial random hyperparameter search over a joint distribution, with shared optimization settings across methods.Most methods use Adam; domain data are split into training and held-out portions.
  • The benchmark compares ERM, IRM, GroupDRO, Mixup, MLDG, CORAL, MMD, DANN, CDANN, MTL, SagNet, ARM, V-REx, RSC, AND-mask, SAND-mask, IGA, and Fish.
  • Weight-averaging approaches are omitted because their contribution is complementary, they use a custom hyperparameter search, and they lack Test-domain model-selection scores.
  • DomainBed contains seven multi-domain computer-vision classification datasets, including Colored MNIST, Rotated MNIST, VLCS, PACS, OfficeHome, TerraIncognita, and DomainNet.
  • The real datasets use ImageNet-pretrained ResNet-50 models with dropout before the new dense layer and frozen batch-normalization layers.The MNIST experiments instead use the DomainBed convolutional architecture rather than the proof-of-concept MLP.

D.2. ‘Training-domain’ model selection

Training-domain model selection shows Fishr remains stronger than ERM, but this setup can favor predictive domain correlations over invariant causal mechanisms and is underspecified for test performance.

  • Fishr performs better than ERM on all real datasets under Training-domain selection, except PACS where both reach 85.5%.
  • Fishr averages 67.1%, ranks third, and exceeds V-REx at 65.6% under Training-domain selection.
  • Training-domain selection can select variant models when training correlations are more predictive than causal features.This explains poor results for all methods on Training-domain Colored MNIST, where color is more predictive than shape during training.
  • Training-domain validation is underspecified because predictors with equally strong held-out training performance can behave very differently in test.The text notes that this favors low regularization and low λ values.
  • Test-domain selection is presented as more realistic because users can label samples to validate algorithm efficiency before deployment.DomainBed’s Test-domain setup remains restrictive, allowing one evaluation per hyperparameter choice and no early stopping.
  • Despite differing views on validation protocols, the paper reports that Fishr remains stronger than ERM under Training-domain selection and is useful for real-world applications under Test-domain selection.

D.3.1. FOCUS ON THE EXPONENTIAL MOVING AVERAGE

Fishr uses an exponential moving average to estimate gradient variances, with the largest benefit appearing on Colored MNIST and a smaller, dataset-dependent effect on OfficeHome.

  • Fishr computes DomainBed gradient variances with an exponential moving average parameterized by γ.Higher γ means a batch influences variance estimates over more later steps.
  • On Colored MNIST, ema raises Fishr’s Test-domain score from 34.0% to 58.9% for training domains {90%, 80%} and test dataset 10%.
  • On OfficeHome, ema raises Fishr’s Test-domain score from 67.5% to 68.2%, while worsening V-REx.
  • Fishr outperforms V-REx on OfficeHome both with and without ema.
  • The authors speculate that ema mainly helps when batches are too small to detect slight correlation shifts in training datasets.They note that ema may be less necessary when shifts are more prominent, as in OfficeHome.

D.3.2. COMPONENT ANALYSIS BY COMPARING GRADIENT VARIANCE VERSUS GRADIENT MEAN MATCHING

The component analysis finds gradient variance matching more effective than gradient mean matching, while combining both moments performs best; effects of warmup, ema, and hyperparameter distributions vary by setting.

  • Gradient variance versus gradient mean matching: IGA is approximately (|E| + 1) times longer to train than ERM, making it costly and impractical.
  • Component choices: Applying gradient mean matching only to classifier weights improves IGA to 67.2% on Test-domain OfficeHome while reducing computational overhead.The analysis motivates enforcing invariance in the classifier rather than low-level layers that may adapt to pixel shifts.
  • Component choices: Warmup improves Colored MNIST Test-domain scores from 58.6% to 59.8% for Fishr but is not critical and reduces IGA’s OfficeHome scores.
  • Gradient variance versus gradient mean matching: Fishr’s gradient variance matching outperforms gradient mean matching methods including Fish, AND-mask, and SAND-mask.
  • Gradient variance versus gradient mean matching: Fishr plus IGA, matching gradient means and variances simultaneously, performs best.
  • Hyperparameter distributions: Changing λ sampling from 10^Uniform(1,4) to 10^Uniform(1,5) yields mostly similar results, while performance differences depend on dataset and selection method.
Loading 2109.02934v3…