Invariant learning

Overview

A model can score well on its training data by relying on a feature that is predictive only in that data. This is a spurious feature. Its reliance grows when the environment's effect on it is less uncertain, and the reliance of a plain risk minimiser does not go away 2202.02466 §3.1 [prop. 1].

Invariant learning asks the model to use only the features whose relation to the label holds in every environment. The assumption is that such features exist and are sufficient for prediction 2202.02466 §2.2 [asm. 1]. The usual route is to train on several environments and penalise disagreement between them. This needs environment labels, which node-level data on a single graph usually does not have 2202.02466 §3.2. For example, invariant risk minimisation learns one feature map and one classifier that is optimal in every environment 2206.15475 §3.2. A causal-inference survey states the same idea for text: the relation to the label should hold across environments. It also notes that these criteria need labelled environments or a labelled nuisance factor, have been applied mostly to classification, and are rarely applied to text 2109.00725 §4.1.2 2109.00725 §4.2.

The method covered first creates its own environments. It edits the observed graph into several versions and trains on those, so the model is pushed to behave the same way on all of them 2202.02466 §3.2.

When the training data already has labelled domains, the variants below need no editing. Inter-domain gradient matching asks that the gradients of the loss from different domains point the same way. Fish approximates this cheaply 2104.09937 §3.3 2104.09937 §3.4.

Variants

Environment-editing risk minimisation (EERM)

Card: Handling Distribution Shifts on Graphs: An Invariance Perspective.

• Virtual environments. K learned editors each produce an edited copy of the graph. An edit samples node pairs and flips them: an existing edge is removed, and a missing one is added 2202.02466 §3.2 app. A.
• Objective. The GNN minimises the variance of the risks across the K edited graphs, plus β times their mean 2202.02466 §3.2 eq. 4. The paper notes that eq. 4 is the form used when the environments are known, and that it corresponds to Risk Extrapolation (REx) 2202.02466 §3.2. See Contradictions for how the survey describes REx.
• Adversarial editors. The editors are chosen to maximise the same variance, and the GNN then minimises it. This is a bilevel problem 2202.02466 §3.2 eq. 5. The editors are updated with REINFORCE, with the variance as reward 2202.02466 app. A alg. 1.
• Reported gains. All author-reported, against ERM with the same backbone. Twitch: up to 7.0% ROC-AUC 2202.02466 §5.2 fig. 4. Facebook-100: up to 7.2% accuracy 2202.02466 §5.2 tab. 2. Elliptic: average F1 gains of 9.6% (GraphSAGE) and 10.0% (GPRGNN) 2202.02466 §5.3 fig. 5. OGB-Arxiv, inductive: EERM is better in five of six cases, and up to 8.1% relative with GPRGNN 2202.02466 §5.3 tab. 3.
• Synthetic shifts. On Cora and Amazon-Photo with artificial spurious features, the paper describes EERM as ahead of ERM on all eight test graphs. The numbers are only in figures 2202.02466 §5.1 fig. 2 fig. 3.

Invariant risk minimisation (IRM) and REx

Card for the survey that summarises these methods: Causal Machine Learning: A Survey and Open Problems. The primary IRM and REx papers have no card, so this section relies on the survey's description only.

• Invariant feature learning. Find content features whose relation to the label is the same in all environments. With multiple environments, the survey lists feature maps with invariant classifiers and matching of content features across environments 2206.15475 §3.1.
• IRM. Learn a feature map φ and one classifier that is optimal in every environment. The survey gives the bilevel form and a Lagrangian relaxation that penalises the squared gradient norm of each environment's risk, with weight λ. The constraint is recovered as λ goes to infinity 2206.15475 §3.2.
• REx. The survey describes it as penalising the worst risk over affine combinations of environments 2206.15475 §3.2. See Contradictions.
• Reported behaviour. The survey reports that IRM often performs no better than ERM, and works best with linear causal relations and enough environment heterogeneity 2206.15475 §3.2. The survey does not state a value for λ 2206.15475 §3.2.

Gradient matching across domains (Fish)

Card: Gradient Matching for Domain Generalization. The paper's direct objective is the ERM loss minus γ times the average inner product of domain gradients, over all pairs of different training domains 2104.09937 §3.3 eq. 4. Computing that objective needs second-order derivatives, which the paper's direct optimiser computes 2104.09937 alg. 2.

• Fish. A first-order approximation. It copies the weights, takes one plain gradient step on a minibatch from each training domain in random order, then moves the real weights toward the copy by a rate ε 2104.09937 alg. 1. The paper's theorem says this points along the gradient-matching direction only in the limit of a small inner rate α 2104.09937 §3.4 app. A. It gives no bound at the α it uses.
• Difference from Reptile. Reptile's inner loop uses minibatches from one task. Fish uses one minibatch per domain, which is what encourages agreement across domains 2104.09937 §3.4 app. A.1.
• Needs labelled domains. The method is defined over a set of training domains S. It does not need the editing step of EERM.
• Reported results. All author-reported. On Wilds, Fish is best on four of six datasets and close on the other two 2104.09937 §4.2 2104.09937 tab. 3. On DomainBed it averages 67.1%, against 67.0% for ERM and 67.2% for CORAL 2104.09937 §4.3 2104.09937 tab. 4, so the average gain there is within noise. On a toy linear task, plain ERM reaches 57% test accuracy and Fish 93% 2104.09937 tab. 1.
• Many domains. The paper says Fish struggles to beat ERM when there are many domains 2104.09937 §5. On iWildCam (324 domains) ERM is ahead on the primary metric, and on Amazon (7,676 reviewers) the two are level 2104.09937 §4.2 2104.09937 app. C.5.
• Released code. The official code is not a literal copy of Algorithm 1. It divides the outer rate by the number of sampled domains and keeps the inner optimiser's state across inner loops. Details are in the card's Paper vs code section; check them before you match reported numbers.

When to use which

Gains are not uniform. Some cells in the paper's tables show EERM below ERM: for the WashU + Brandeis + Carnegie row in Facebook-100, Texas (56.12 vs 56.25) 2202.02466 tab. 2, and SAGE on OGB-Arxiv 2014–2016 (41.55 vs 42.09) 2202.02466 tab. 3. Many other gains are within one standard deviation 2202.02466 tab. 2. Report the variance over runs, not only the mean.

Implementation sketch

Common core of the EERM variant. The sketch merges the inner loop of Algorithm 1 into one step, so it does not show the T inner steps 2202.02466 alg. 1.

# One EERM training step  [arxiv:2202.02466 §3.2] [eq. 5]
# Input: one graph G, labels Y, GNN f_theta, K editors w_1..w_K

def eerm_step(G, Y, f_theta, editors, beta):
    # Build K virtual environments by editing the observed graph  [§3.2]
    envs = [edit(G, w_k) for w_k in editors]
    # Risk of the current model on each edited graph  [§3.2]
    risks = [risk(f_theta, env, Y) for env in envs]
    # Editors take a REINFORCE step that increases the variance  [app. A]
    update_editors(editors, envs, reward=variance(risks))
    # GNN step on variance plus beta times mean of the risks  [eq. 4]
    loss = variance(risks) + beta * mean(risks)
    f_theta.step(loss)

Notes on the paper's definitions and the released code. Follow the paper unless you have checked the code:

• Editor normalisation. The paper takes a softmax over the column index for each row 2202.02466 app. A. The code's softmax runs over dim 0 and is reused in the log-probability model.py. Check the density you sample from against the one you differentiate.
• Reward. The appendix says the reward "can be defined as the inverse loss" 2202.02466 app. A. Algorithm 1 and the code use the variance itself 2202.02466 alg. 1 main.py.
• Update schedule. Algorithm 1 updates θ only on the last inner step 2202.02466 alg. 1. The code updates θ on the first inner step main.py.
• Editor reset. The code re-initialises the editors at the start of each epoch main.py. Algorithm 1 carries them across iterations 2202.02466 alg. 1.
• Diagonal. The paper's complement Ā excludes the identity 2202.02466 §3.2. The code's complement includes the diagonal model.py, so self-pairs can be edited.
• Unstated values. The paper does not give K, the number of epochs, the optimisers for the editors and the GNN, or weight decay for the §5.3 experiments 2202.02466 app. F.2 2202.02466 app. F.3.

Common core of the Fish variant, for labelled training domains. It is the paper's Algorithm 1 with the inner loop written out 2104.09937 alg. 1.

# One Fish meta-step  [arxiv:2104.09937 alg. 1]
# Input: domains D_1..D_S, inner rate alpha, outer rate eps, weights theta

def fish_meta_step(theta, domains, alpha, eps):
    # Start the inner loop from a copy of the current weights  [alg. 1 l. 2]
    inner = copy(theta)
    # Visit every training domain once, in random order  [alg. 1 l. 3]
    for d in permute(domains):
        # One plain gradient step on a minibatch of this domain  [alg. 1 l. 4-6]
        inner = inner - alpha * gradient(loss(inner, sample_batch(d)), inner)
    # Move the real weights part of the way toward the copy  [alg. 1 l. 8]
    return theta + eps * (inner - theta)

Contradictions

• What REx is. The EERM card reports that eq. 4 of that paper, the variance plus β times the mean of the risks, corresponds to Risk Extrapolation 2202.02466 §3.2. The survey describes REx as penalising the worst risk over affine combinations of environments 2206.15475 §3.2. These are two different descriptions of REx, and this page does not decide which is correct. Check the REx primary source before relying on either.
• Invariance methods against ERM. The survey reports that IRM often performs no better than ERM 2206.15475 §3.2. EERM reports gains over ERM on graph benchmarks 2202.02466 §5.2 2202.02466 §5.3. The two claims concern different methods and settings, so they are not compared here. Fish's own Wilds table is consistent with the survey's remark: IRM is below ERM on Camelyon17 (64.2 vs 70.3 average accuracy) and on iWildCam (15.1 vs 31.0 macro F1) 2104.09937 tab. 3. That is one paper's numbers, not a general result.
• Fish's main-text and appendix numbers. The paper's main results table and its per-dataset appendix give different test numbers for the same method and test sets. Camelyon17 average accuracy is 74.7 in the main table and 79.5 in the appendix 2104.09937 tab. 3 2104.09937 app. C.2. Amazon 10th-percentile accuracy is 53.3 and 56.0 2104.09937 tab. 3 2104.09937 app. C.6. The paper does not explain the gap, so this page does not pick one. Say which table a reported number comes from.

Papers

• Handling Distribution Shifts on Graphs: An Invariance Perspective: proposes EERM, which uses learned graph editors and a mean-plus-variance objective for node-level out-of-distribution learning on graphs.
• Gradient Matching for Domain Generalization: proposes Fish, a first-order approximation of inter-domain gradient matching that takes one step per training domain on a copy of the model, evaluated on Wilds and DomainBed.
• Causal Machine Learning: A Survey and Open Problems: surveys causal ML and summarises invariant feature learning, IRM and REx; it runs no experiments of its own.
• Causal Inference in Natural Language Processing: Estimation, Prediction, Interpretation and Beyond: a survey of causal inference for text; its section on robust prediction covers invariance across environments and its limits for text, and it runs no experiments.