--- kind: concept slug: invariant-learning title: Invariant learning summary: >- Train a model to rely only on features whose link to the label holds across environments. Covered: EERM, which edits one graph into virtual environments when labels are missing; IRM, as a survey summarises it; and Fish, which matches gradients across labelled training domains. Trade-off: EERM's editors are dense n × n matrices, costly on large graphs. tags: [out-of-distribution, invariance, graphs, gradient-alignment] papers: ["2104.09937", "2109.00725", "2202.02466", "2206.15475"] related: [domain-generalization, causal-inference, graph-neural-network, policy-gradient, meta-learning] status: draft updated_at: 2026-11-01 --- ## 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 [arxiv: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 [arxiv: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 [arxiv:2202.02466 §3.2]. For example, invariant risk minimisation learns one feature map and one classifier that is optimal in every environment [arxiv: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 [arxiv:2109.00725 §4.1.2] [arxiv: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 [arxiv: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 [arxiv:2104.09937 §3.3] [arxiv:2104.09937 §3.4]. ## Variants ### Environment-editing risk minimisation (EERM) Card: [Handling Distribution Shifts on Graphs: An Invariance Perspective](/papers/2202.02466). - **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 [arxiv:2202.02466 §3.2] [app. A]. - **Objective.** The GNN minimises the variance of the risks across the K edited graphs, plus β times their mean [arxiv: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) [arxiv:2202.02466 §3.2]. See [Contradictions](#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 [arxiv:2202.02466 §3.2] [eq. 5]. The editors are updated with REINFORCE, with the variance as reward [arxiv:2202.02466 app. A] [alg. 1]. - **Reported gains.** All author-reported, against ERM with the same backbone. Twitch: up to 7.0% ROC-AUC [arxiv:2202.02466 §5.2] [fig. 4]. Facebook-100: up to 7.2% accuracy [arxiv:2202.02466 §5.2] [tab. 2]. Elliptic: average F1 gains of 9.6% (GraphSAGE) and 10.0% (GPRGNN) [arxiv: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 [arxiv: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 [arxiv: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](/papers/2206.15475). 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 [arxiv: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 [arxiv:2206.15475 §3.2]. - **REx.** The survey describes it as penalising the worst risk over affine combinations of environments [arxiv:2206.15475 §3.2]. See [Contradictions](#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 [arxiv:2206.15475 §3.2]. The survey does not state a value for λ [arxiv:2206.15475 §3.2]. ### Gradient matching across domains (Fish) Card: [Gradient Matching for Domain Generalization](/papers/2104.09937). 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 [arxiv:2104.09937 §3.3] [eq. 4]. Computing that objective needs second-order derivatives, which the paper's direct optimiser computes [arxiv: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 ε [arxiv: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 α [arxiv: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 [arxiv: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 [arxiv:2104.09937 §4.2] [arxiv:2104.09937 tab. 3]. On DomainBed it averages 67.1%, against 67.0% for ERM and 67.2% for CORAL [arxiv:2104.09937 §4.3] [arxiv: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% [arxiv:2104.09937 tab. 1]. - **Many domains.** The paper says Fish struggles to beat ERM when there are many domains [arxiv:2104.09937 §5]. On iWildCam (324 domains) ERM is ahead on the primary metric, and on Amazon (7,676 reviewers) the two are level [arxiv:2104.09937 §4.2] [arxiv: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](/papers/2104.09937) section; check them before you match reported numbers. ## When to use which | Situation | Use | Why | |---|---|---| | Environment labels exist for the training data | The mean-plus-variance objective (eq. 4) over those environments | The paper states this is the form used when environments are known [arxiv:2202.02466 §3.2]. No editors needed. | | Environment labels exist, relations are roughly linear, and the environments differ enough | IRM, per the survey's description | The survey says IRM works best in this setting [arxiv:2206.15475 §3.2]. Tune λ yourself, because the survey gives no value. | | One graph, node-level labels, no environment map | EERM | This is the paper's setting [arxiv:2202.02466 §2.1] [arxiv:2202.02466 §3.2]. | | Large graph (tens of thousands of nodes or more) | Plan an editing scheme that is not dense, or check memory first | Each editor is an n × n parameter matrix [code: synthetic/model.py#L86-L89]. The paper gives no sparse alternative. For OGB-Arxiv (n = 169,343) the matrix is about 115 GB in float32 per editor, by our arithmetic (n² × 4 bytes) [arxiv:2202.02466 app. F.2.3]. | | Text inputs (NLP) | Check first that you have labelled environments or a labelled nuisance factor | The survey says the distributional criteria need such labels and have been applied mostly to classification [arxiv:2109.00725 §4.1.2]. It says they are rarely applied to text [arxiv:2109.00725 §4.2]. | | Tight compute budget | Keep K and T small | Each inner step runs K full GNN passes [arxiv:2202.02466 §3.2]. The paper's grid uses T ∈ {1, 5} [arxiv:2202.02466 app. F.2]. | | Temporal shift on a graph | EERM, with the caveat that gains were reported on Elliptic and Arxiv only | The largest average gains in the paper's tables are on Elliptic [arxiv:2202.02466 §5.3]. Test on your own split. | | Labelled training domains, a moderate count, and the baseline can use the same pretrained checkpoint | Fish, with the paper's per-dataset rates and N | Fish is best in the main table on Poverty, Camelyon17, FMoW and CivilComments [arxiv:2104.09937 §4.2] [arxiv:2104.09937 tab. 3]. The paper says a fair baseline needs the same pretraining [arxiv:2104.09937 app. E]. | | Hundreds or thousands of domains | Run ERM first; Fish may not beat it | The paper reports ERM ahead or level on iWildCam (324 domains) and Amazon (7,676) [arxiv:2104.09937 §4.2] [arxiv:2104.09937 §5]. | | Many domains per epoch, limited compute | Sample N domains per inner loop, not all S | Algorithm 1 visits every domain [arxiv:2104.09937 alg. 1]. The paper's Wilds runs sample between 3 and 10 domains per inner loop [arxiv:2104.09937 tab. 14]. On the four tasks in its N sweep, N = 5 and N = 10 differ by at most about one point [arxiv:2104.09937 tab. 16]. | 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) [arxiv:2202.02466 tab. 2], and SAGE on OGB-Arxiv 2014–2016 (41.55 vs 42.09) [arxiv:2202.02466 tab. 3]. Many other gains are within one standard deviation [arxiv: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 [arxiv:2202.02466 alg. 1]. ```pseudocode # 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 [arxiv:2202.02466 app. A]. The code's softmax runs over dim 0 and is reused in the log-probability [code: synthetic/model.py#L99-L109]. Check the density you sample from against the one you differentiate. - **Reward.** The appendix says the reward "can be defined as the inverse loss" [arxiv:2202.02466 app. A]. Algorithm 1 and the code use the variance itself [arxiv:2202.02466 alg. 1] [code: synthetic/main.py#L113-L114]. - **Update schedule.** Algorithm 1 updates θ only on the last inner step [arxiv:2202.02466 alg. 1]. The code updates θ on the first inner step [code: synthetic/main.py#L110-L121]. - **Editor reset.** The code re-initialises the editors at the start of each epoch [code: synthetic/main.py#L108]. Algorithm 1 carries them across iterations [arxiv:2202.02466 alg. 1]. - **Diagonal.** The paper's complement Ā excludes the identity [arxiv:2202.02466 §3.2]. The code's complement includes the diagonal [code: synthetic/model.py#L98], 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 [arxiv:2202.02466 app. F.2] [arxiv: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 [arxiv:2104.09937 alg. 1]. ```pseudocode # 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 [arxiv:2202.02466 §3.2]. The survey describes REx as penalising the worst risk over affine combinations of environments [arxiv: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 [arxiv:2206.15475 §3.2]. EERM reports gains over ERM on graph benchmarks [arxiv:2202.02466 §5.2] [arxiv: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) [arxiv: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 [arxiv:2104.09937 tab. 3] [arxiv:2104.09937 app. C.2]. Amazon 10th-percentile accuracy is 53.3 and 56.0 [arxiv:2104.09937 tab. 3] [arxiv: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](/papers/2202.02466): 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](/papers/2104.09937): 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](/papers/2206.15475): 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](/papers/2109.00725): 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.