Handling Distribution Shifts on Graphs: An Invariance Perspective

In plain words

A graph neural network learns from a social network, a citation graph or a payment network, and then has to predict labels on a graph that looks a bit different. Models trained this way often latch onto shortcuts that held in the training graph but not in the new one. This is hard on graphs, because the nodes are linked to each other, so their data are not independent samples.

The paper's idea is to make the model prove it does not depend on the shortcut. It has only one graph, so it invents extra versions of it by adding and removing edges. Small learned "editors" choose the edits, and they are trained to make the model's errors as different as possible across the edited copies. The model is trained to do well on average and to be equally good on every copy. A feature that only works in one version of the graph gets punished, so the model learns to rely on features that work everywhere.

TL;DR

• Node-level out-of-distribution (OOD) learning on graphs, where training data is usually a single graph with no environment labels §2.1 §3.2.
• EERM trains K learned edge editors to produce K edited graphs, and minimises the mean plus variance of the losses on them §3.2 eq. 5.
• The authors report gains over standard ERM on all tested backbones, with the largest gains on the temporal and cross-domain splits §5 tab. 2 tab. 3.
• All numbers are author-reported. The method adds an editor for every virtual environment and a dense n × n parameter matrix per editor (see Pitfalls).

Problem

Standard empirical risk minimisation (ERM) minimises the average training loss. With spurious features, ERM can pick up correlations that depend on the environment. A toy linear example shows that its optimum keeps a reliance on the spurious feature, and the reliance grows as the environment's effect becomes less uncertain §3.1 [prop. 1].

For graphs, the OOD problem needs two changes. Nodes are not i.i.d., so the paper treats each node's L-hop ego-graph as the input and fragments the graph into ego-graphs §2.1. It also needs a notion of invariant features that uses the structure, so it extends the invariance assumption to a recursive readout over the BFS tree of each ego-graph §2.2 [asm. 1].

Methods such as Risk Extrapolation (REx) need risks from several environments. Node-level training data usually has one graph and no node-to-environment map, so this is not available §3.2.

Method

Problem formulation. For a node v with L-hop neighbourhood N_v, the input is the ego-graph G_v = (A_v, X_v) and the target is y_v. The OOD node-level problem is a min-max over environments of the expected risk §2.1 eq. 2:

\min_f \max_{e \in \mathcal E} \mathbb E_{G \sim p(\mathbf G \mid \mathbf e = e)} \left[ \frac{1}{|V|} \sum_{v \in V} \mathbb E_{y \sim p(\mathbf y \mid \mathbf G_{\mathbf v} = G_v, \mathbf e = e)} \big[ l(f(G_v), y) \big] \right] eq. 2

The paper notes that the outer sampling over G can be dropped in node-level tasks because there is usually a single input graph §2.1.

Motivating example. Two-dimensional node features, a one-layer mean aggregation, and an environment-dependent term ε with variance σ_e² §3.1:

y_v = \frac{1}{|N_v|} \sum_{u \in N_v} x_u^1 + n_v^1, \qquad x_v^2 = \frac{1}{|N_v|} \sum_{u \in N_v} y_u + n_v^2 + \epsilon §3.1

The stated results are that ERM's optimum is [(1+σ_e²)/(2+σ_e²), 1/(2+σ_e²)] [prop. 1], and that minimising the variance of risks across environments has its unique optimum at [1, 0] [prop. 2].

Step 1: Virtual environments. Instead of real environments, K context generators g_{w_k}(G) produce K edited graphs from the observed graph. Each generator has a Boolean edit matrix B^k, and the edited adjacency is §3.2:

A^k = A + B^k \circ (\overline{A} - A), \qquad \overline{A} = \mathbf 1 \mathbf 1^\top - I - A §3.2

Each row of the edit matrix is sampled with a multinomial distribution whose probabilities are a softmax over the parameter matrix P_k app. A. A sampled entry flips the pair: an existing edge is removed and a missing one is added §3.2.

Step 2: Objective. The GNN f = c ∘ h is trained on the risks of the K edited graphs. The paper's single-level objective, which is the form used by REx when environments are known, is §3.2 eq. 4:

\min_\theta \; \mathbb V_{\mathbf e}\big[L(G^e, Y^e;\theta)\big] + \beta \, \mathbb E_{\mathbf e}\big[L(G^e, Y^e;\theta)\big] eq. 4

where L(G^e, Y^e; θ) = (1/|V_e|) Σ_v l(f_θ(G_v^e), y_v^e) and β trades mean against variance §3.2.

Step 3: Bilevel EERM. Because the environments are not observed, the editors are optimised adversarially so that the variance is as large as possible, and the GNN then minimises the variance and mean on those graphs §3.2 eq. 5:

\min_\theta \; \operatorname{Var}\!\big(\{L(g_{w_k^*}(G), Y;\theta)\}_{k=1}^K\big) + \frac{\beta}{K} \sum_{k=1}^K L(g_{w_k^*}(G), Y;\theta)

\text{s.t.} \quad [w_1^*, \dots, w_K^*] = \arg\max_{w_1,\dots,w_K} \operatorname{Var}\!\big(\{L(g_{w_k}(G), Y;\theta)\}_{k=1}^K\big) eq. 5

The editors are updated with REINFORCE, treating the edit as an action and the variance as the reward §3.2 app. A. The gradient for editor k is ∇ log p_{w_k}(A^k) times the reward app. A.

Step 4: Training loop. The procedure is Algorithm 1 alg. 1.

# EERM training loop  [alg. 1] [§3.2]
# Input:  graph G = (A, X), labels Y, GNN params theta, editors w_1..w_K
# Output: trained GNN params theta

def train_eerm(G, Y, theta, editors, T, beta, alpha_f, alpha_g):
    # K editors, one per virtual environment  [§3.2]
    K = len(editors)
    # Repeat until convergence or the epoch limit is reached  [alg. 1]
    while not converged:
        for t in range(1, T + 1):
            # Edit the graph once per editor  [§3.2]
            graphs = [edit(G, w_k) for w_k in editors]
            # Risk of the GNN on each edited graph  [§3.2]
            risks = [risk(theta, g, Y) for g in graphs]
            # Editor objective J1: variance of the K risks  [alg. 1]
            j1 = variance(risks)
            # REINFORCE step for each editor, reward = j1  [app. A]
            for k in range(K):
                editors[k] += alpha_g * grad_log_p(editors[k], graphs[k]) * j1
            # On the last inner step, update the GNN on J2  [alg. 1]
            if t == T:
                # J2 = variance + beta/K * sum of risks  [alg. 1] [eq. 5]
                j2 = variance(risks) + (beta / K) * sum(risks)
                theta -= alpha_f * grad(j2, theta)
    return theta

# One editor's edit of the adjacency matrix  [§3.2] [app. A]
# Input:  adjacency A (N, N), editor params B_k (N, N), samples s
# Output: edited adjacency A_k (N, N)

def edit(A, B_k, s):
    # For each node n, p(a_nm) is a softmax over m  [app. A]
    probs = softmax(B_k)   # each row sums to 1 over m
    # Draw s target nodes m for each node n  [app. A]
    sampled = multinomial(probs, s)
    # Mark the sampled pairs (n, m)  [app. A]
    M = indicator(sampled)
    # Complement of A, without the diagonal  [§3.2]
    A_complement = ones - identity - A
    # Flip the marked pairs: add missing edges, delete existing ones  [§3.2]
    A_k = A + M * (A_complement - A)
    return A_k

Hyperparameters

Values are as stated in the paper. Learning rate for the GNN is α_f and for the editors α_g. The paper searches the grid α_f ∈ {1e-4, 2e-4, 1e-3, 5e-3, 1e-2}, α_g ∈ {1e-4, 1e-3, 5e-3, 1e-2}, β ∈ {0.2, 0.5, 1, 2, 3}, s ∈ {1, 5, 10} and T ∈ {1, 5} on the validation set app. F.2. K is not stated anywhere in the paper's text.

Other settings, all from the paper:

• Architecture: GCN, GAT, GraphSAGE use the PyTorch Geometric convolution with self-loops and batch normalisation after each layer; ReLU for GCN and GraphSAGE, ELU for GAT app. F.1. GCNII uses α = 0.1 and λ = 1.0 app. F.1. GPRGNN uses α = 0.1 with PPR initialisation app. F.1.
• Training: a fixed number of epochs, and the test result is taken at the epoch with the best validation score app. F.3. The number of epochs is not stated.
• Editor and GNN optimiser, and the number of sampled graphs K: not stated.
• Hardware and software: RTX 2080Ti (11 GB) and RTX 8000 (48 GB); Python 3.7, PyTorch 1.9.0, PyTorch Geometric 1.7.2 app. F.

Results

Author-reported. Setup: each task is the one listed in §5; baselines are standard ERM with the same GNN backbone.

Artificial spurious features (Cora, Amazon-Photo) §5.1. The data are generated from 10 environments; one for training, one for validation and eight for testing §5.1 app. E.1. Results are over 20 trials fig. 2 fig. 3. The authors describe EERM as consistently ahead of ERM on Cora and Photo across the eight test graphs T1–T8 §5.1 fig. 2 fig. 3. When spurious features are removed at inference, the training accuracy of ERM drops more than that of EERM, which the authors read as ERM relying more on spurious features §5.1 fig. 2 fig. 3. Figure values are not transcribed here.

Unseen domains (Twitch, Facebook-100) §5.2. Twitch: train on DE, validate on ENGB, test on five other networks, using ROC-AUC. Up to 7.0% improvement on ROC-AUC §5.2 fig. 4. Facebook-100: three graphs train, two validate and three test, using accuracy. Up to 7.2% improvement §5.2 tab. 2.

Dynamic data (Elliptic, OGB-Arxiv) §5.3. Elliptic: F1 on nine test folds, with averaged improvements of 9.6% (GraphSAGE) and 10.0% (GPRGNN) §5.3 fig. 5. Both methods drop after the seventh fold, which the authors link to a market shutdown §5.3 fig. 5. OGB-Arxiv, inductive setting: EERM is better in five of six cases, and up to 8.1% relative improvement with GPRGNN §5.3 tab. 3.

Table 1: Dataset summary tab. 1

Table 2: Facebook-100 test accuracy (mean ± std) tab. 2

Table 3: OGB-Arxiv test accuracy by time interval (mean ± std, inductive setting) tab. 3

Table 4: Twitch-Explicit graph statistics tab. 4 in app. E.1

Paper vs code

Code: raw/code/2202.02466 at commit ffdc4a11161976fac7dd71e2aa1dcd72db6e44e9, from the first author's repository linked in the paper's abstract footnote. The repository has no LICENSE file, so no code is quoted. Paths below are relative to the repository root.

• Editor sampling and its log-probability use different normalisations. The paper defines p(a_nm) as a softmax over m for each row n app. A. The code takes the softmax over dim 0, which normalises over rows for each column, then samples with torch.multinomial on the rows of that matrix, so each row is treated as a distribution over columns model.py. The log-probability then uses the column-wise normalisation again model.py. So the sampling density and the density used in the REINFORCE gradient appear not to match.
• Normaliser subtracted once per column, not once per sample. The paper's log-probability is a product over the s samples per row, so the normaliser is counted s times app. A. The code sums the s sampled logits and subtracts the column normaliser once model.py. This differs for s > 1.
• Reward. Appendix A says the reward "can be defined as the inverse loss" app. A. Algorithm 1 and the code use the variance itself as reward (J1 = Var) alg. 1 main.py.
• Timing of the GNN update. Algorithm 1 updates θ only when t == T alg. 1. The code updates the GNN on the first inner step (m == 0) and updates the editors on every inner step main.py.
• Editors reset every epoch. The code calls model.gl.reset_parameters() at the start of each epoch main.py, which re-initialises the editor matrix uniformly model.py. The paper's Algorithm 1 carries the editors across iterations and says nothing about resets alg. 1.
• Diagonal in the complement. The paper defines Ā without the identity term §3.2. The code's complement is ones - A, which includes the diagonal, so the editor can add or delete self-pairs model.py.
• Variance estimator. The code uses torch.var_mean, which is the unbiased (N−1) estimator model.py. The paper's Var is not further specified.
• β schedule. The code computes β as 1 * beta * epoch / epochs + beta * (1 - epoch / epochs), which equals β at every epoch main.py. The effective β is constant, and the paper uses a single fixed β §3.2.
• Cora settings in run.sh differ from §5.1. The paper gives α_f = 0.01, α_g = 0.005, β = 2.0, s = 5 for both Cora and Photo app. F.2.1. The Cora run uses --lr 0.005 (α_f), --lr_a 0.001 (α_g), --beta 1.0, --num_sample 1 and K = 10 run.sh. The Photo run matches α_f and α_g but uses --beta 1.0, --num_sample 1 and K = 5 run.sh. The Cora GAT and SGC variants use K = 10 and 5 run.sh run.sh.
• Other experiments match the appendix. The Twitch, Facebook-100, Elliptic and OGB-Arxiv settings in run.sh agree with app. F.2.2 and app. F.2.3 for the learning rates, β, s and T run.sh run.sh run.sh run.sh run.sh. The code supplies K (3 or 5), which the paper does not state.
• Weight decay. The Elliptic runs set --weight_decay 0. run.sh. The paper does not state weight decay for §5.3.
• Synthetic-data generators. The paper says "randomly initialized GNN" app. E.1. The code uses a 2-layer GCN with hidden width 10 and 10 outputs, plus a linear context term, to generate 10 environments create_synthetic.py. The script's __main__ only generates the GAT variant create_synthetic.py.
• Epochs. The paper says "a fixed number of epochs" app. F.3 but gives no number. The code default is 200 parse.py.
• Dropout. The code default is 0.0 parse.py. The paper does not state dropout.
• Dense editor parameters. Each editor holds an n × n matrix, so the parameters are K × n × n model.py model.py. The OGB-Arxiv runs use the full graph with n = 169,343 main.py. That is about 115 GB per editor in float32, before the K multiplier. The paper does not discuss this.

Pitfalls

• Editor normalisation. Implement the editor as the paper defines it: softmax over the column index for each row, with a log-probability that subtracts the normaliser once per sample. Do not copy the code's dim-0 softmax without checking the density model.py.
• Reward sign and definition. Use the variance as reward, as in Algorithm 1 and the code. The appendix's "inverse loss" wording is not what either implements app. A alg. 1.
• Dense memory. A full n × n matrix per editor is infeasible for graphs of OGB-Arxiv or Elliptic size. The paper gives no sparse or subsampled alternative. Any large-graph use needs an explicit design choice (e.g. sparse editing of a candidate-edge set) that the paper does not describe.
• Cost. Each inner step runs K full GNN forward passes, so the cost per epoch grows with K and T. K is not stated in the paper, and the code uses 3 to 10 run.sh run.sh.
• REINFORCE variance. The editors are updated from s sampled edits per row, with a scalar reward (the variance) shared by all K editors. The paper does not discuss the resulting variance or any baseline subtraction.
• Hyperparameter mismatches. For Cora, the paper's §5.1 values and the run.sh values differ in α_f, β and s (see Paper vs code). Start from the paper's values and the grid in app. F.2, and report which one was used.
• Unstated values. K, the number of epochs, the number of runs, the optimiser for the editors and the GNN, and weight decay for §5.3 are not stated in the paper.
• Elliptic split overlap. §5.3 says 5/5/33 snapshots for train/valid/test, but the appendix gives training 7–11, validation 12–17 and test 17–49, so snapshot 17 appears in both validation and test, and validation has six snapshots, not five §5.3 app. E.3.
• Table 2 exception. For the WashU + Brandeis + Carnegie row, Texas EERM (56.12) is below ERM (56.25) tab. 2. Many EERM gains in Table 2 are within one standard deviation, e.g. Bingham + Duke + Princeton, Brown: 52.76 ± 3.40 vs 50.43 ± 4.58 tab. 2.
• Table 3 exception. For SAGE on 2014–2016, EERM (41.55) is below ERM (42.09) tab. 3.
• Synthetic results are plots only. The synthetic-shift results are shown only as figures, so their numbers are not transcribed here fig. 2 fig. 3.
• Theory. The optimality results need Assumption 1 (invariance and sufficiency) and Assumption 2 (environment heterogeneity) [asm. 1] [asm. 2]. Theorem 3's bound needs I_{e'}(G_v; y | z) = I_e(G_v; y | z) [thm. 3], which cannot be checked from the observed training data. The proofs are in app. C and app. D and were not transcribed here.
• Test selection. Test results are taken at the best-validation epoch app. F.3. Compare with other methods using the same rule, not the final epoch.

Lineage

• Builds on: the invariance principle as in Invariant Risk Minimization and the causal invariance works cited in §2.2; the risk-variance objective of Risk Extrapolation (REx), which the paper names as the multi-environment version of Eq. 4 §3.2.
• Compared against: standard ERM with the same GNN backbone in all experiments §5.
• Uses: the policy-gradient (REINFORCE) estimator for the editors app. A; graph-structure adversarial editing following the graph-robustness works cited in §3.2.