PaperlayerView exact revision on arXiv

Source-linked AI summary

Gradient Matching for Domain Generalization

Yuge Shi, Jeffrey Seely, Philip H. S. Torr, N. Siddharth, Awni Hannun, Nicolas Usunier, Gabriel Synnaeve

arXiv:2104.09937v3cs.LGstat.ML

TL;DR

Domain generalization requires models to perform on unseen domains despite distribution shift, motivating methods that learn invariant features rather than spurious correlations. The paper proposes IDGM to align inter-domain gradients and Fish as its first-order approximation, achieving strong results across WILDS and DOMAINBED while struggling when domain counts are very large.

  • Problem

    Domain generalization seeks models that perform well on unseen out-of-distribution data, where relying on spurious correlations can fail when domains change.

  • Method

    IDGM maximizes the inner product between gradients from different domains, while Fish approximates this objective with a computationally simpler first-order algorithm.

  • Results

    Fish exhibits state-of-the-art performance on 13 datasets from the WILDS and DOMAINBED benchmarks and performs well across diverse vision and language tasks.

  • Takeaways & Limitations

    Gradient matching provides a simple and efficient alternative for learning invariant features across diverse domain-generalization tasks.

  • Takeaways & Limitations

    Fish struggles to outperform ERM when the number of domains is large, including datasets with 324 and 7,676 domains.

Abstract

from arXiv · show

Machine learning systems typically assume that the distributions of training and test sets match closely. However, a critical requirement of such systems in the real world is their ability to generalize to unseen domains. Here, we propose an inter-domain gradient matching objective that targets domain generalization by maximizing the inner product between gradients from different domains. Since direct optimization of the gradient inner product can be computationally prohibitive -- requires computation of second-order derivatives -- we derive a simpler first-order algorithm named Fish that approximates its optimization. We demonstrate the efficacy of Fish on 6 datasets from the Wilds benchmark, which captures distribution shift across a diverse range of modalities. Our method produces competitive results on these datasets and surpasses all baselines on 4 of them. We perform experiments on both the Wilds benchmark, which captures distribution shift in the real world, as well as datasets in DomainBed benchmark that focuses more on synthetic-to-real transfer. Our method produces competitive results on both benchmarks, demonstrating its effectiveness across a wide range of domain generalization tasks.

1 Introduction

The paper targets domain generalization by aligning gradient directions across domains, encouraging invariant predictions while avoiding the cost of direct second-order optimization through the first-order Fish algorithm. Experiments across WILDS and DOMAINBED show broad performance, with Fish improving inter-domain gradient alignment relative to ERM.

  • Motivation: Domain generalization seeks models that perform well on unseen out-of-distribution data by learning invariant rather than spuriously correlated features.Multiple training domains provide the basis for identifying features that remain stable across domains.
  • Method: IDGM maximizes the inner product between domain-specific gradients while minimizing loss, encouraging shared optimization paths and invariant predictions.Unlike ERM, which can favor spurious features, IDGM is designed to focus on features common across domains.
  • Method: Directly optimizing IDGM is computationally costly because gradient descent requires second-order derivatives.Fish uses a first-order approximation inspired by Reptile and is presented as simple and computationally effective.
  • Results: Fish exhibits state-of-the-art performance on 13 datasets spanning the WILDS and DOMAINBED domain-generalization benchmarks.The authors describe the results as broadly applicable across different applications and domain-generalization subgenres.
  • Results: The experiments also show that Fish improves normalized inter-domain gradient inner products, whereas the inner product decreases during ERM training.This directly tests whether the proposed training behavior occurs in practice.

2 Related Work

Related domain-generalization methods include distributional robustness, domain-invariant representation learning, and invariant risk minimization, alongside ERM and meta-learning connections. The paper positions gradient matching as an efficient alternative to hard invariant constraints.

  • Domain Generalization: Domain generalization uses source-domain training data to predict in a new target domain without accessing target data during training.The paper focuses on settings with multiple known source domains.
  • Distributional Robustness: Distributional robustness minimizes worst-case loss over distributions constructed from training domains, addressing shifts such as covariate and subpopulation shift.These shifts can change input distributions or mixture weights while preserving aspects of the underlying domains.
  • Domain-invariant Representation Learning: Domain-invariant representation learning makes domains statistically indistinguishable and bases prediction on the resulting features.Examples include domain-adversarial penalties, feature-distribution matching, self-supervised tasks, and maximum mean discrepancy.
  • Invariant Risk Minimization: Invariant risk minimization learns representations whose optimal classifiers are shared across domains, aiming to exploit invariant causal effects and reduce spurious correlations.This approach differs from gradient matching in how invariance is specified.
  • ERM: ERM minimizes average training-domain loss and remains a strong practical baseline despite lacking theoretical guarantees when train and test distributions differ.The paper notes that recent benchmarks often find ERM competitive with domain-generalization methods.
  • Meta-learning: Meta-learning inspires several domain-generalization methods, while Fish draws specifically on Reptile’s first-order approximation to MAML.The paper distinguishes Fish’s domain-generalization goal from meta-learning’s focus on adapting to unseen tasks.
  • Gradient Matching: Inter-domain gradient matching offers an alternative to hard constraints on invariant features and combines standard gradient computations with parameter updates.The resulting algorithm is described as simple and efficient.

3 Methodology

The methodology frames domain generalization as learning from invariant cross-domain relationships rather than spurious correlations. It proposes IDGM to align domain gradients and Fish as a first-order approximation that retains similar behavior without second-order derivatives.

  • Goals: Domain generalization trains models for unseen domains, where invariant features should be learned while spuriously label-correlated features are ignored.The formal setup separates disjoint training and test domains, with performance evaluated on unseen test domains.
  • The pitfall of ERM: a linear example: 97% train accuracy from ERM falls to 57% test accuracy because the learned model places most weights on spurious features.In contrast, the invariant feature f1 maintains 93% accuracy across domains, while the spurious feature f2 reaches at most 57% on D2 and D3.
  • Inter-domain Gradient Matching (IDGM): IDGM maximizes the inner product between gradients from different domains while minimizing ERM loss, aligning gradient directions and encouraging shared optimization paths.When domain gradients point in similar directions, a gradient step improves performance on both domains; opposite directions do not provide this guarantee.
  • Inter-domain Gradient Matching (IDGM): 93% train and test accuracy show that IDGM focuses the linear model on invariant feature f1 despite lower training accuracy than ERM.The experiment uses normalized GIP during training to avoid exploding gradients.
  • Optimizing IDGM with Fish: Fish performs domain-specific inner-loop updates on a cloned model, then updates the original parameters using the weighted difference between clone and original.Each inner-loop step uses a minibatch from a permuted training domain, and the outer update is ε(eθ − θ).
  • Optimizing IDGM with Fish: Taylor expansion shows that Fish combines averaged inner-loop gradients with a minibatch GIP-gradient term, implicitly approximating IDGM while avoiding second-order derivatives.The approximation can lose direct control over the GIP scaling γ, although a smoothed version recovers that scaling term.

4 Experiments

Experiments evaluate Fish and the inter-domain gradient matching objective on a controlled shape-color task, WILDS, and DOMAINBED. Across these settings, Fish generalizes well, approximates direct IDGM optimization, and benefits from domain-aware grouping.

  • CDSPRITES-N: CDSPRITES-N uses domains with deterministic shape-color matches, making shape invariant and color spuriously correlated for testing domain generalization.The dataset varies the number of domains N from 5 to 50, with test colors reassigned independently to shapes.
  • CDSPRITES-N: Over 90% test accuracy is achieved by Fish and direct IDGM when N ≥10, reaching near 100% when N ≥20.All methods achieve 100% training accuracy, so the test split distinguishes their generalization performance.
  • CDSPRITES-N: Fish is about 10 times faster than directly optimizing IDGM while achieving similar performance on CDSPRITES-N.The experiment also reports that Fish increases the inter-domain gradient inner product whereas ERM does not.
  • WILDS: On WILDS, Fish is best overall, outperforming every baseline on 4 of 6 datasets and matching the best method on AMAZON and IWILDCAM.The evaluation covers six real-world distribution-shift datasets and follows WILDS configurations and model selection for fair comparison.
  • WILDS: Fish also works on CIVILCOMMENTS subpopulation shift without the special sampling strategies used by its baselines.The other WILDS datasets represent pure domain generalization with disjoint training and test domains.
  • DOMAINBED: On DOMAINBED, Fish ranks second across 7 datasets, is 0.1% behind Coral, and is one of three methods outperforming ERM.Results use training domains for validation and average five random trials.
  • Ablation: Random grouping performs worse than Fish on every dataset, with Fish reaching 100% test accuracy on CDSPRITES with N=10 while random grouping and ERM remain random.This indicates that Fish’s effectiveness depends on grouping data by domains rather than maximizing gradient inner products across random batches.

5 Conclusion

The paper presents IDGM for domain generalization and approximates its costly second-order optimization with the first-order Fish algorithm. Experiments show Fish learns invariant features and performs well across diverse domain-generalization tasks, while its advantage narrows when many domains are available.

  • Conclusion: IDGM matches gradient directions across domains, while Fish approximates this objective with a simple first-order update.The derivation links a component of Fish’s update to the gradient of the gradient inner product.
  • Conclusion: Fish is evaluated on simple invariant-feature tasks and on WILDS and DomainBed benchmarks spanning vision and language architectures.The reported architectures include DenseNet, ResNet-50, and BERT.
  • Conclusion: Fish performs well across different domain-generalization subgenres and surpasses baseline performance on a diverse range of tasks.The experiments are reported as replicable with 1500 GPU hours on NVIDIA V100.
  • Conclusion: When the number of domains is large, Fish struggles to outperform ERM.The authors identify scaling to datasets with orders of magnitude more domains as future work.
  • Conclusion: The derivation analyzes Fish’s inner-loop update using domain-sampled minibatches, Hessians, and Taylor approximations.The analysis assumes S inner-loop steps and distinguishes step indices from domain indices.

A.1 Fish and Reptile: Differences and Connections

Fish and Reptile share a meta-learning-style update but align gradients across different units: Fish uses different tasks, whereas Reptile uses batches from the same task. On CDSPRITES-N, their domain-generalization behavior differs substantially.

  • Differences and Connections: Domain generalization differs from meta-learning because the target domain provides no examples during training, corresponding to K = 0.Meta-learning instead exposes some test examples at test time, with K > 0.
  • Differences and Connections: Fish and Reptile use similar algorithmic structures but differ in the task sources used for their inner-loop minibatches.Algorithm 3 identifies the inner-loop sampling choice as their key algorithmic difference.
  • Differences and Connections: Fish aligns gradients between batches from different tasks, whereas Reptile aligns gradients between batches from the same task.This distinction explains why Fish maximizes across-task gradient inner products while Reptile maximizes within-task inner products.
  • Empirical Comparison: At N = 10, Fish’s CDSPRITES-N test accuracy reaches 100%, while Reptile’s remains 50% regardless of N.The passage reports markedly different train and test behavior despite the methods’ algorithmic similarity.

B.1 Derivation

The derivation introduces SmoothFish by scaling the gradient-inner-product component of Fish’s update. The parameter γ interpolates between ERM-like and Fish-like behavior, and the method is evaluated across several WILDS settings.

  • Derivation: A component of Fish’s update aligns with the gradient of the gradient inner product, enabling explicit control over its scaling.This connects Fish’s first-order update to the original IDGM objective.
  • Derivation: SmoothFish introduces γ to control the emphasis on the gradient inner product versus the average gradient.γ = 1 recovers Fish, while γ = 0 is equivalent to ERM with learning rate αS.
  • Derivation: Algorithm 4 uses permuted source domains and samples minibatches from each selected domain during the inner loop.The displayed procedure describes the smoothed Fish update implementation.
  • Experiments: SmoothFish experiments on six WILDS datasets evaluate γ values from 0.1 to 0.8, along with γ = 0 and γ = 1.Other hyperparameters are kept the same as in the main experiments.

C Discussions and Results on WILDS

The WILDS evaluation covers six datasets with real-world domain shifts, including a 23-country satellite-imagery task. On POVERTYMAP-WILDS, Fish achieves the highest test performance while matching the best baseline on validation, though the reported results use WILDS v1 splits.

  • Dataset Scope: The WILDS study reports results for six datasets and notes that its experiments use WILDS v1 rather than the later v2 splits.The authors state that they are working on updating the results to WILDS v2.
  • Results: The WILDS result materials include tables for dataset details, POVERTYMAP-WILDS, and CAMELYON17-WILDS.These table captions identify the organization of the reported benchmark results.
  • POVERTYMAP-WILDS: POVERTYMAP-WILDS predicts a real-valued asset wealth index from satellite imagery across 23 country domains.Because the number of domains is large, each iteration samples N << S domains; the selected value is N = 5.
  • POVERTYMAP-WILDS: POVERTYMAP-WILDS is evaluated using Pearson correlation (r) between predicted and ground-truth asset indices over three random seeds.The model uses a ResNet-18 backbone.
  • Results: Fish obtains the highest test performance while matching the best baseline’s validation performance on POVERTYMAP-WILDS.Its validation-to-test performance is more stable, with a smaller standard deviation than the baselines.

C.2 CAMELYON17-WILDS

Across several WILDS tasks, Fish improves performance on challenging domain-shift metrics, especially worst-domain or lower-tail accuracy, while showing task-dependent trade-offs on aggregate and class-balanced measures.

  • CAMELYON17-WILDS: 6%: Fish surpasses the best baseline in CAMELYON17 test accuracy, while its validation–test gap is only 3%.Other baselines show validation–test gaps of 11%–27%.
  • FMOW-WILDS: Fish achieves the highest FMOW worst-region accuracy on both test and validation sets, while ranking second in average accuracy behind ERM.Fish also has the smallest standard deviation across the reported metrics.
  • CIVILCOMMENTS-WILDS: 4% and 7%: Fish exceeds the best baseline on CIVILCOMMENTS test and validation worst-group accuracy, respectively, while matching ERM’s average accuracy within standard deviation.ERM has the highest average accuracy but the lowest worst-group accuracy.
  • CIVILCOMMENTS-WILDS: CIVILCOMMENTS uses subpopulation shift rather than the pure domain-generalization split used by the other WILDS datasets, with test domains forming a subpopulation of training domains.Its WILDS baselines are therefore tailored to avoiding systematic failures on minority subpopulations.
  • IWILDCAM-WILDS: Fish improves IWILDCAM test accuracy and weighted F1 by 1% over ERM, but trails ERM by 1%–3% on macro F1.The results suggest weaker performance on rarer species despite improved overall test accuracy.
  • AMAZON-WILDS: Fish matches ERM on AMAZON’s 10th-percentile reviewer accuracy and outperforms all other baselines, despite lower average accuracy than ERM.The 10th-percentile metric evaluates consistency across reviewers.

E Ablation Studies on Pre-trained Models

The pretrained-model ablation evaluates Fish from three ERM initialization points and finds that the preferred convergence level varies substantially across datasets.

  • Experimental setup: Three ERM initialization settings are compared: 10% of training data, 50% of training data, and a converged model.The comparison measures how pretrained-model convergence affects Fish’s training performance.
  • Results: CIVILCOMMENTS maintains good performance across different pretrained-model convergence levels.This dataset is comparatively insensitive to the initialization point in the reported ablation.
  • Results: FMOW and IWILDCAM prefer converged pretrained models, with performance worsening as the amount of pretraining data decreases.The reported trend contrasts with CAMELYON17’s preference for an earlier initialization.
  • Results: CAMELYON17 performs best when Fish starts from a model trained on only 10% of the data, with test accuracy declining at higher convergence levels.POVERTY is excluded because Fish can be trained from scratch on its smaller dataset.

F Ablation Studies on hyperparameters

The hyperparameter ablations vary Fish’s learning rates while holding their product constant and compare different numbers of sampled domains per meta-step.

  • Learning rates: Fish’s inner- and outer-loop learning rates α and ϵ are varied while keeping α · ϵ constant.The resulting comparisons are reported in Figure 8.
  • Meta steps: The meta-step ablation compares sampling N = 5, 10, and 20 domains from the available training domains.CAMELYON17 is excluded from this sampling comparison because it has T = 3 domains.
  • Results: Changing these hyperparameters generally has little effect on performance, although N = 20 slightly degrades results on some datasets, including POVERTY and CIVILCOMMENTS.This conclusion is based on the reported meta-step ablation.

G Tracking gradient inner product

The gradient-tracking experiment compares normalized inter-domain gradient inner products before and after matched Fish and ERM updates across WILDS datasets.

  • Measurement: The plots include CDSPRITES-N and five WILDS datasets, with dataset-specific trends reflecting differences in data, domain splits, and architectures.CDSPRITES-N and POVERTY show an initial dip because they are trained from scratch, unlike the fine-tuning settings used elsewhere.
  • Measurement: Normalized gradient inner products are recorded before and after updates while Fish and ERM train on the same minibatch sequence.The measurements use cosine similarity between minibatches from different domains sampled in each inner loop.
  • Results: ERM updates often decrease the gradient inner product, whereas Fish updates either increase it significantly or maintain it at roughly the same level.The comparison uses dotted pre-update and solid post-update curves.
  • Results: Fish’s post-update inter-domain gradient inner product is consistently higher than ERM’s across the plotted datasets.The authors interpret this pattern as evidence that Fish increases or maintains inter-domain gradient alignment.
Loading 2104.09937v3…