TMLR 2026

EdgeMask-DG*: Learning Domain-Invariant Graph Structures via Adversarial Edge Masking

Rishabh Bhattacharya, Naresh Manwani

Main figure for EdgeMask-DG*: Learning Domain-Invariant Graph Structures via Adversarial Edge Masking

Abstract

Structural shifts pose a significant challenge for graph neural networks, as graph topology acts as a covariate that can vary across domains. Existing domain generalization methods rely on fixed structural augmentations or training on globally perturbed graphs, mechanisms that do not pinpoint which specific edges encode domain-invariant information. We argue that domain-invariant structural information is not rigidly tied to a single topology but resides in the consensus across multiple graph structures derived from topology and feature similarity. To capture this, we first propose EdgeMask-DG, a novel min-max algorithm where an edge masker learns to find worst-case continuous masks subject to a sparsity constraint, compelling a task GNN to perform effectively under these adversarial structural perturbations. Building upon this, we introduce EdgeMask-DG*, an extension that applies this adversarial masking principle to an enriched graph. This enriched graph combines the original topology with feature-derived edges, allowing the model to discover invariances even when the original topology is noisy or domain-specific. EdgeMask-DG* is the first to systematically combine adaptive adversarial topology search with feature-enriched graphs. We provide a formal justification for our approach from a robust optimization perspective. We demonstrate that EdgeMask-DG* achieves new state-of-the-art performance on diverse graph domain generalization benchmarks, including citation networks, social networks, and temporal graphs. Notably, on the Cora OOD benchmark, EdgeMask-DG\* lifts the worst-case domain accuracy to 78.0\

Proposed method: Graph-DG

Benchmark Results

Tested on 1 dataset

Tables show presented results in global context (best 3 papers + GCN + MLP).

Classic

#ModelAccuracy ↑Paper
1
GraphTARIF
0.9703 ±0.0019
2
GAT
0.9669 ±0.0014
3
GAT+T
0.9660 ±0.0007
7
SAGE
0.9643 ±0.0027
69
EdgeMask-DG*
0.9480 ±0.0060
This paper
258
GRM
0.9270 ±0.0160
This paper
357
EERM
actual #357 vs paper #357
0.9180 ±0.0090
actual #357 vs paper #357
paper 0.9180 ±0.0090

Improvement within reported uncertainty; not highlighted.

paper 0.9180 ±0.0090 Δ at pub 0.0101
375
LiSA
0.9150 ±0.0150
This paper
390
MLP
0.9128 ±0.0063
441
IS-GIB
0.9020 ±0.0090
This paper
466
MARIO
0.8940 ±0.0140
This paper
483
ERM
actual #483 vs paper #483
0.8860 ±0.0130
actual #483 vs paper #483
paper 0.8860 ±0.0130

Improvement within reported uncertainty; not highlighted.

paper 0.8860 ±0.0130 Δ at pub 0.0054
510
MMD
0.8480 ±0.0060
This paper
545
DRNN
0.7710 ±0.0120
This paper
549
ARM
0.7460 ±0.0070
This paper
SOTA performance (not significantly different from the best) Baseline / comparison (this paper) Global context (other papers)