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

arXiv: https://arxiv.org/abs/2602.05571 · HTML: https://mlsanity.com/paper/2602.05571/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 33 of its 45 references that are in the map.
- First posted on arXiv 5 Feb 2026.

## Builds on (most cited first, 30 of 33)

- [Semi-Supervised Classification with Graph Convolutional Networks](https://mlsanity.com/paper/1609.02907/) (2016)
- [Graph Attention Networks](https://mlsanity.com/paper/1710.10903/) (2017)
- [Inductive Representation Learning on Large Graphs](https://mlsanity.com/paper/1706.02216/) (2017)
- [How Powerful are Graph Neural Networks?](https://mlsanity.com/paper/1810.00826/) (2018)
- [A Comprehensive Survey on Graph Neural Networks](https://mlsanity.com/paper/1901.00596/) (2019)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [Fast Graph Representation Learning with PyTorch Geometric](https://mlsanity.com/paper/1903.02428/) (2019)
- [Open Graph Benchmark: Datasets for Machine Learning on Graphs](https://mlsanity.com/paper/2005.00687/) (2020)
- [Invariant Risk Minimization](https://mlsanity.com/paper/1907.02893/) (2019)
- [Simplifying Graph Convolutional Networks](https://mlsanity.com/paper/1902.07153/) (2019)
- [Distributionally Robust Neural Networks for Group Shifts: On the Importance of Regularization for Worst-Case Generalization](https://mlsanity.com/paper/1911.08731/) (2019)
- [Graph Contrastive Learning with Augmentations](https://mlsanity.com/paper/2010.13902/) (2020)
- [The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables](https://mlsanity.com/paper/1611.00712/) (2016)
- [Pitfalls of Graph Neural Network Evaluation](https://mlsanity.com/paper/1811.05868/) (2018)
- [Revisiting Semi-Supervised Learning with Graph Embeddings](https://mlsanity.com/paper/1603.08861/) (2016)
- [Graph Contrastive Learning with Adaptive Augmentation](https://mlsanity.com/paper/2010.14945/) (2020)
- [Generalizing to Unseen Domains via Adversarial Data Augmentation](https://mlsanity.com/paper/1805.12018/) (2018)
- [EvolveGCN: Evolving Graph Convolutional Networks for Dynamic Graphs](https://mlsanity.com/paper/1902.10191/) (2019)
- [Adversarial Attacks on Neural Networks for Graph Data](https://mlsanity.com/paper/1805.07984/) (2018)
- [Adversarial Attack on Graph Structured Data](https://mlsanity.com/paper/1806.02371/) (2018)
- [Graph Random Neural Network for Semi-Supervised Learning on Graphs](https://mlsanity.com/paper/2005.11079/) (2020)
- [Data Augmentation for Graph Neural Networks](https://mlsanity.com/paper/2006.06830/) (2020)
- [Adaptive Risk Minimization: Learning to Adapt to Domain Shift](https://mlsanity.com/paper/2007.02931/) (2020)
- [Handling Distribution Shifts on Graphs: An Invariance Perspective](https://mlsanity.com/paper/2202.02466/) (2022)
- [Out-Of-Distribution Generalization on Graphs: A Survey](https://mlsanity.com/paper/2202.07987/) (2022)
- [Shift-Robust GNNs: Overcoming the Limitations of Localized Graph Training Data](https://mlsanity.com/paper/2108.01099/) (2021)
- [Twitch Gamers: a Dataset for Evaluating Proximity Preserving and Structural Role-based Node Embeddings](https://mlsanity.com/paper/2101.03091/) (2021)
- [Mind the Label Shift of Augmentation-based Graph OOD Generalization](https://mlsanity.com/paper/2303.14859/) (2023)
- [Do We Need Anisotropic Graph Neural Networks?](https://mlsanity.com/paper/2104.01481/) (2021)
- [Improving Out-of-Distribution Generalization by Adversarial Training with Structured Priors](https://mlsanity.com/paper/2210.06807/) (2022)

## Abstract (arXiv)

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%, a +3.8 pp improvement over the prior state of the art (74.2%). The source code for our experiments can be found here: https://anonymous.4open.science/r/TMLR-EAEF/