Gradient Matching for Domain Generalization

In plain words

A model trained on data from several sources often learns shortcuts. In one toy example it leans on a feature that looks reliable inside each source but means something different in the next one. The paper asks for a different signal: the weights should move in a direction that helps every source at the same time. If the gradients from the different sources point the same way, a step on one of them also helps the others, so the feature it learns is probably a real one.

The direct way to reward this needs second-order derivatives, which are slow to compute. The method, Fish, gets close to the same effect with a cheap procedure. It starts from the current weights, makes a copy, and trains the copy briefly on one batch from each source in turn. It then moves the real weights part of the way toward that copy. The copy's path is pulled toward the directions the sources share, so the update works like the direct objective without computing it.

The authors test this on classification and regression tasks with distribution shift, including a benchmark of real-world shifts and a benchmark with synthetic-to-real shifts. The method is competitive and wins on several of the tasks, and it does not win everywhere.

TL;DR

• Fish is a first-order approximation of inter-domain gradient matching: it maximises the inner product of gradients from different training domains without computing second derivatives §3.4 app. A.
• One meta-step runs S plain gradient updates on a copy of the model, one per training domain, then applies θ ← θ + ε(θ̃ − θ) alg. 1.
• Author-reported: Fish is the best of the compared methods on Poverty, Camelyon17, FMoW and CivilComments, and roughly tied on Amazon and iWildCam §4.2 tab. 3; on DomainBed it is second of ten by average accuracy, 0.1 points below Coral §4.3 tab. 4.
• The official code (MIT) implements Fish and ERM only; the direct IDGM objective, SmoothFish and the DomainBed runs are not in that repository main.py README.md.
• Two main-text results tables disagree with the per-dataset appendix tables for the same method (see Pitfalls), so check the numbers against the table you use.

Problem

Empirical risk minimisation (ERM) minimises the average loss over the training domains and ignores how the domains differ. It can therefore lock onto features that predict the label in the training domains but not in a new one §3.1 §2. The paper's toy example shows this: a feature that is only correlated with the label within the training domains gives ERM 97% training accuracy and 57% test accuracy, while the invariant feature alone gives 93% on both tab. 1 §3.2.

Prior invariance methods either enforce a hard constraint on the features (IRM, domain-adversarial training, CORAL, MMD) or minimise the variance of the gradients across domains (IGA). The authors argue that IGA behaves like ERM when ERM is already optimal on every domain, so it can fail in the same way §2. The paper's target is a method that prefers non-ERM solutions whenever the domain gradients agree, and that is cheap enough to run on large models.

Method

Setup and goal

There are S training domains, each a set of i.i.d. samples, and T test domains disjoint from them. The goal is a parameter vector θ with low expected loss on the unseen test domains §3.1, eq. 1:

\arg\min_{\theta} \; \mathbb{E}_{\mathcal{D} \sim \mathcal{D}_{te}} \mathbb{E}_{(x,y) \sim \mathcal{D}} \big[ \ell((x,y); \theta) \big] \qquad [\text{eq. 1}]

ERM minimises the average over the training domains §3.1, eq. 2:

\mathcal{L}_{erm}(\mathcal{D}_{tr}; \theta) = \mathbb{E}_{\mathcal{D} \sim \mathcal{D}_{tr}} \mathbb{E}_{(x,y) \sim \mathcal{D}} \big[ \ell((x,y); \theta) \big] \qquad [\text{eq. 2}]

Inter-domain gradient matching (IDGM)

Let G_s be the expected gradient of the loss on domain s §3.3, eq. 3. The gradient inner product (GIP) is the average over all pairs of distinct domains of G_i · G_j. The IDGM objective subtracts γ times the GIP from the ERM loss §3.3, eq. 4:

\mathcal{L}_{idgm} = \mathcal{L}_{erm}(\mathcal{D}_{tr}; \theta) - \gamma \, \frac{2}{S(S-1)} \sum_{i \ne j} G_i \cdot G_j \qquad [\text{eq. 4}]

Here G_i is the expected gradient on domain i §3.3, eq. 3. The paper notes that the sum of pairwise products can be computed in linear time from the sum of the gradients and their squared norms, up to a constant §3.3. For direct IDGM the paper uses a normalised GIP to avoid exploding gradients §3.2 §4.1; when it tracks the GIP it defines the normalised form as cosine similarity app. G.

The direct optimiser needs the gradient of the GIP with respect to θ, which is second order. Algorithm 2 shows that optimiser alg. 2.

# Direct IDGM step  [alg. 2]
# Input:  per-domain minibatches; learning rate eps; weight gamma
# Output: updated weights

def idgm_step(model, batches, eps, gamma):
    # Gradient of each domain's minibatch loss at the current weights  [alg. 2 l. 2-4]
    grads = [gradient(loss(model, b), model) for b in batches]
    # Average gradient, the ERM part  [alg. 2 l. 8]
    mean_grad = average(grads)
    # Gradient inner product over all pairs of different domains  [alg. 2 l. 8]
    gip = pairwise_inner_product_mean(grads)
    # Second-order term: gradient of the GIP with respect to the weights  [alg. 2 l. 8]
    gip_grad = gradient(gip, model)
    # Descend on ERM and ascend on the GIP  [alg. 2 l. 8]
    model = model - eps * (mean_grad - gamma * gip_grad)
    return model

Fish: a first-order approximation

Fish runs an inner loop over the domains on a copy of the weights, then moves the original weights toward the copy alg. 1. The inner loop visits every training domain once, in a random order alg. 1 l. 3.

# Fish  [alg. 1]
# Input:  domains D_1..D_S; inner rate alpha; outer rate eps; iterations
# Output: trained weights theta

def fish_train(model, domains, alpha, eps, iterations):
    # Repeat the outer loop  [alg. 1 l. 1]
    for step in range(iterations):
        # Clone the current weights to start the inner loop  [alg. 1 l. 2]
        inner = copy(model)
        # Visit every domain once, in random order  [alg. 1 l. 3]
        for domain in permute(domains):
            # Draw one minibatch from this domain  [alg. 1 l. 4]
            batch = sample_batch(domain)
            # One plain gradient step on the clone  [alg. 1 l. 5-6]
            grad = gradient(loss(inner, batch), inner)
            inner = inner - alpha * grad
        # Move the original weights toward the clone  [alg. 1 l. 8]
        model = model + eps * (inner - model)
    return model

Inside the inner loop, gradients are taken at the clone's weights rather than at θ. A Taylor expansion of the Fish update shows that, once the scaled ERM gradient is removed, what remains points in the same direction as the GIP gradient. This is the theorem in §3.4, with the proof in app. A. It is a limit as α → 0; the paper gives no bound at the α values it uses.

The difference from Reptile is which batches go into the inner loop. Reptile uses minibatches from one task, which encourages agreement within a task. Fish takes one minibatch from each training domain, which encourages agreement across domains §3.4 app. A.1.

SmoothFish (optional variant)

Fish gives no direct control over the strength of the GIP term. SmoothFish adds a coefficient γ on the difference between the Fish step and the scaled ERM step app. B. With γ = 1 it is Fish, and with γ = 0 it is ERM with learning rate αS app. B. The paper says γ made little difference in its experiments §3.4 app. B.

Hyperparameters

Paper: ERM baseline tab. 13

Source: app. D, tab. 13. The paper says the same ERM settings were used to make the Fish pretrained models. Poverty was not pretrained app. D, tab. 13.

Paper: Fish tab. 14

Fish uses the same model, batch size, validation metric and optimiser as ERM, and weight decay is always 0 app. D, tab. 14.

Source: app. D, tab. 14. The meta-steps column is the number of domains sampled per inner loop; the paper's Amazon row lists 7,676 domains while the text uses S = 5008 for the training split app. C.6.

Paper: CdSprites-N

The paper gives no separate Fish table for CdSprites-N. It tests N from 5 to 50 and uses the normalised GIP for direct IDGM §4.1. The code configuration for CdSprites is in the next subsection.

Code configuration (official repository)

Values from src/config.py; these are the defaults that main.py loads for each dataset. Effective outer step = meta_lr / meta_steps (see Paper vs code).

Source: config.py. CdSprites uses num_domains = 15 by default config.py and --num-domains on the command line main.py.

Results

All numbers are author-reported. Each row names its setup. Std values are as printed.

Toy linear example tab. 1

Two training domains, one test domain, four binary features; linear model.

Source: §3.2 tab. 1. The data layout is in fig. 2.

CdSprites-N §4.1 fig. 4

Shape classification with N domains, each with a fixed shape-colour pairing that changes at test time. Five runs. The paper reports:

• Train accuracy is 100% for all methods at every N.
• Fish and direct IDGM exceed 90% test accuracy for N ≥ 10 and reach nearly 100% for N ≥ 20.
• ERM stays near chance until N = 20 and reaches 95% only for N ≥ 40.
• The text says Fish is about 10 times faster than direct IDGM §4.1. No table or timing figure backs this; it is a claim in the text.

Wilds benchmark §4.2 tab. 2 tab. 3

Six Wilds datasets (Wilds v1, per app. C), same learning rates, batch sizes, architectures and model selection as the Wilds baselines §4.2. Metrics follow Wilds. All runs use 3 seeds except Camelyon17 with 10 §4.2.

Table 2, dataset summary tab. 2:

Table 3, main results tab. 3 (test set; bold marks the paper's best per column):

Source: §4.2 tab. 3. The paper says Fish is best on four datasets and close to the best on Amazon and iWildCam §4.2.

Per-dataset appendix results app. C.1–C.6. These are the numbers the appendix gives for the same test sets. They differ from tab. 3 (see Pitfalls).

Poverty tab. 7, validation / test Pearson r: Fish 0.81 ± 6e-3 / 0.81 ± 9e-3; IRM 0.81 ± 4e-2 / 0.78 ± 3e-2; ERM 0.80 ± 3e-2 / 0.78 ± 3e-2; ERM (ours) 0.80 ± 3e-2 / 0.77 ± 5e-2; Coral 0.80 ± 4e-2 / 0.77 ± 5e-2 app. C.1.

Camelyon17 tab. 8, validation / test avg. acc. %: Fish 82.5 ± 1.2 / 79.5 ± 6.0; ERM 84.3 ± 2.1 / 73.3 ± 9.9; ERM (ours) 84.1 ± 2.4 / 70.5 ± 12.1; IRM 86.2 ± 2.1 / 60.9 ± 15.3; Coral 86.3 ± 2.2 / 59.2 ± 15.1 app. C.2.

FMoW tab. 9, validation / test, avg. and worst %:

Source: app. C.3.

CivilComments tab. 10, validation / test, avg. and worst %:

Source: app. C.4.

iWildCam tab. 11, validation / test (accuracy, weighted F1, macro F1):

Source: app. C.5.

Amazon tab. 12, validation / test, average and 10th-percentile %: Fish 73.8 ± 0.1 / 57.3 ± 0.0 (val) and 72.9 ± 0.2 / 56.0 ± 0.0 (test); ERM 74.3 ± 0.0 / 57.3 ± 0.0 and 73.5 ± 0.1 / 56.0 ± 0.0; ERM (ours) 74.0 ± 0.0 / 57.3 ± 0.0 and 73.3 ± 0.1 / 56.0 ± 0.0; IRM 73.6 ± 0.2 / 56.4 ± 0.8 and 73.3 ± 0.2 / 55.1 ± 0.8; Reweighted 69.6 ± 0.0 / 53.3 ± 0.0 and 69.2 ± 0.1 / 52.4 ± 0.8 app. C.6.

The paper's reading: on iWildCam the ERM baseline is best on the primary metric, and on Amazon Fish and ERM are level or ERM is slightly ahead. The paper attributes the weak domain-generalisation results on these two datasets to the large number of domains (324 and 7,676) §4.2.

DomainBed §4.3 tab. 4

Seven DomainBed datasets, 5 random trials, training-domain validation for model selection §4.3. Test accuracy in %.

The paper ranks Fish second of ten by average, 0.1 points below Coral and one of three methods above ERM §4.3. The DomainBed experiments are not in the official repository (see Paper vs code).

Analysis §4.4

• Gradient inner product during training (CdSprites-N, N = 15): Fish's normalised GIP between domain minibatches rises during training, while ERM's stays flat §4.4 fig. 5. The appendix shows the same comparison on five Wilds datasets, with the GIP measured before and after each update, and says ERM updates often lower the GIP while Fish keeps it or raises it app. G.
• Random grouping §4.4 tab. 5 (test accuracy %):

The paper's reading: grouping minibatches by domain is what makes Fish work; random grouping, which still maximises the inner product between minibatches, does not reach the same generalisation.

Ablations app. E app. F

Pretrained ERM models at different convergence levels, test numbers tab. 15:

The paper's reading: CivilComments tolerates any convergence level; FMoW and iWildCam prefer the converged model; Camelyon17 does best with the least pretraining app. E.

Meta steps N (domains sampled per inner loop) tab. 16:

The paper reports that α and ε were varied with their product held fixed, and that N = 20 slightly hurts Poverty and CivilComments app. F fig. 7.

Compute

The paper estimates about 1,500 GPU-hours on NVIDIA V100 to replicate all its experiments §5.

Dataset details tab. 6

Source: app. C tab. 6.

Paper vs code

Repository: https://github.com/YugeTen/fish at commit 333efa24572d99da0a4107ab9cc4af93a915d2a9, MIT licence, copyright Facebook LICENSE. The GitHub user is YugeTen, the abstract links the repository [abstract], and the README directs contact to the first author's address README.md.

What the code contains: Fish and ERM training, for six Wilds datasets and CdSprites-N main.py README.md. Not in the repository: direct IDGM (Algorithm 2), SmoothFish, the GIP tracking experiments and the DomainBed runs; the README points to a separate DomainBed fork for those README.md.

Discrepancies found, each with both citations:

1. Outer step size is divided by S. The paper's update applies ε directly alg. 1 l. 8. The code passes meta_lr / meta_steps to the update main.py utils.py. The table values match the undivided meta_lr tab. 14 config.py config.py. The effective outer rate is therefore ε/N: 0.0033 for Camelyon17, 0.002 for FMoW and Amazon, 0.001 for iWildCam, 0.01 for CivilComments and CdSprites, and 0.02 for Poverty.
2. Inner optimiser. Algorithm 1 uses a plain gradient step alg. 1 l. 6. The code uses the dataset's ERM optimiser (Adam or SGD with momentum), and for every dataset except CdSprites it carries the optimiser state across inner loops main.py config.py. The paper says Fish uses the same optimiser as ERM app. D but does not say whether the state persists.
3. Camelyon17 inner learning rate. The paper gives α = 1e-3 tab. 14; the code's Camelyon17 learning rate is 1e-4 config.py. The paper's ERM learning rate for Camelyon17 is also 1e-3 tab. 13, while the code uses 1e-4 with momentum 0.9 config.py. The paper does not state a momentum value.
4. Weight decay for Civil and Amazon ERM. The paper's ERM table gives weight decay 0.01 for both tab. 13. The code uses 0 for both config.py config.py.
5. FMoW model selection metric. The paper's ERM table gives validation average accuracy tab. 13. The code selects on worst-region accuracy config.py. The code also saves the last FMoW checkpoint, not the best one utils.py.
6. Poverty selection metric. The paper's text names Pearson r app. C.1. The code selects on r_wg config.py. That metric is not defined in this repository; it presumably comes from the Wilds package, which is not in the repository.
7. Model backbones. The paper says BERT for CivilComments and Amazon §4.2 app. C.4 app. C.6. The code uses DistilBERT civil.py amazon.py. The paper names DenseNet-121 for Camelyon17 but does not state an ImageNet initialisation. The code uses DenseNet-121 without pretrained weights camelyon.py. FMoW is initialised from ImageNet fmow.py.
8. Camelyon17 ERM pretraining length. The paper's cut-off is "iter 500" tab. 13. The code's Camelyon17 pretraining runs 10,000 iterations and evaluates only at the end, because print_iters is -1 config.py config.py main.py.
9. iWildCam pretraining length. The paper says the pretrained model is trained "till convergence" app. C.5. The code runs a fixed 24,000 iterations config.py.
10. CivilComments domain count. The paper lists 16 domains (8 identities × toxicity) tab. 14. The code builds 18 groups: a "no identity" group plus 8 identities, each split by toxicity datasets.py.
11. Domain sampling per inner loop. Algorithm 1 loops over all S domains alg. 1 l. 3. The code samples meta_steps domains per inner loop and stops an epoch once fewer than meta_steps domains have batches left, so some batches go unused each epoch main.py utils.py.
12. CdSprites-N batches differ between ERM and Fish. Fish batches come from one domain at a time with its own colour palette datasets.py datasets.py. The ERM path reads items through __getitem__, which picks a random palette per image datasets.py. The paper describes both methods as training on the same N domains §4.1. The comparison may therefore not be like-for-like.
13. Poverty loss. The paper does not state a loss. The code uses mean squared error for Poverty and cross-entropy for the others utils.py. The Poverty backbone is a ResNet-18 with 8 input channels poverty.py. No discrepancy found for: the Fish inner-loop structure (S steps, one minibatch per domain), the update form θ ← θ + ε(θ̃ − θ) up to the S factor in item 1, the meta steps for Camelyon17, CivilComments, FMoW, iWildCam, Poverty and Amazon in tab. 14, and the batch sizes in tab. 13.

Pitfalls

• Main-text and appendix numbers disagree. For the same method and test sets, tab. 3 gives Camelyon17 74.7 but app. C.2 gives 79.5; Amazon 53.3 in tab. 3 but 56.0 in app. C.6; CivilComments worst-group 72.8 in tab. 3 but 74.2 in app. C.4; Poverty 0.80 in tab. 3 but 0.81 in app. C.1; FMoW worst-case 34.6 in tab. 3 but 34.3 in app. C.3. The paper does not explain the gap. Pick one source and say which you used.
• Paper's Fish step is not the one in the code. An implementation that follows Algorithm 1 literally (plain step, undivided ε) will take steps about S times larger than the released runs (see Paper vs code 1 and 2). Match the code's behaviour if you want to reproduce the reported numbers.
• Inner optimiser state. Carrying the Adam state across inner loops main.py is an implementation choice the paper does not describe. Results may depend on it.
• Second-order term is absent by design. Fish drops the GIP gradient on purpose; the claimed alignment holds only as α → 0 §3.4 app. A. The paper tests α = 1e-3 to 2e-6 tab. 14 with no bound on the error.
• Cost per meta-step. Each meta-step does S forward/backward passes on a model copy main.py, so cost grows with the number of domains sampled. The paper's Wilds runs sample N ≪ S for large S app. C.1 app. C.3 app. C.4 app. C.6. The 10× speed claim §4.1 has no timing table to check.
• Large domain counts. The paper says Fish struggles to beat ERM when there are many domains §5. iWildCam (324) and Amazon (7,676) show this §4.2. Both results in tab. 3 and app. C.5 show ERM ahead on the primary metric for iWildCam.
• Pretraining dependence. Camelyon17, FMoW, iWildCam, Amazon and CivilComments start from an ERM checkpoint, and the results depend on how far it converged app. E tab. 15. Poverty and CdSprites start from scratch tab. 13 [code: src/config.py#L46-L47] config.py. A fair baseline needs the same pretraining.
• Model selection. The paper says it uses the Wilds selection strategy §4.2. The code evaluates at a fixed interval and keeps the checkpoint with the best validation metric, except for FMoW, where it keeps the last one utils.py. Reported test numbers depend on this choice.
• Seed count. Most Wilds datasets use three seeds; Camelyon17 uses ten §4.2 app. C.2. Small seed counts make the ± values unreliable for close comparisons, such as Fish vs. ERM on Amazon.
• Stochastic domain order. The inner loop visits domains in random order alg. 1 l. 3 utils.py. Results may vary with the seed.
• Wilds version. The paper reports Wilds v1 and says v2 results were not ready at submission app. C. Do not compare these numbers with Wilds v2 leaderboards without rerunning.
• Epoch and checkpoint defaults are code-only. The number of epochs and the evaluation interval are not in the paper config.py. Using them reproduces the code, not the paper's text.

Lineage

• Builds on Reptile, a first-order meta-learning method, whose inner loop uses minibatches from one task instead of different domains §2 app. A.1. Reptile is cited as Nichol, Achiam and Schulman, "On First-Order Meta-Learning Algorithms" (no arXiv id in the source).
• Compared against IGA, which minimises the variance of gradients across domains instead of maximising their inner product §2.
• Compared against IRM (arXiv:1907.02893), CORAL, MMD, DANN, Group DRO, Mixup and MLDG on DomainBed (arXiv:2007.01434) §4.3 tab. 4.
• Evaluated on the Wilds benchmark (arXiv:2012.07421) §4.2.
• Superseded by: not stated in the source.