Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Rishabh Bhattacharya, Naresh Manwani

TL;DR

EdgeMask-DG* tackles graph domain generalization under structural shifts by learning domain-invariant subgraphs through an adversarial edge masking framework applied to an enriched graph that fuses original topology with feature-derived edges. A two-network setup, comprising a TaskNet (GAT) and a MaskNet, performs a min-max game where the MaskNet sparsely masks edges to maximally degrade the TaskNet, while the TaskNet learns robust representations under these perturbations. The enrichment via spectral-clustered and kNN-based edges, combined with adversarial pruning, yields a practical approach that achieves new state-of-the-art results on diverse graphs, including citation networks, social networks, and temporal graphs. The authors provide a robust optimization interpretation and comprehensive experiments, demonstrating the value of adaptive topology search over static augmentations for cross-domain generalization in GNNs.

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\%, 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/

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

TL;DR

EdgeMask-DG* tackles graph domain generalization under structural shifts by learning domain-invariant subgraphs through an adversarial edge masking framework applied to an enriched graph that fuses original topology with feature-derived edges. A two-network setup, comprising a TaskNet (GAT) and a MaskNet, performs a min-max game where the MaskNet sparsely masks edges to maximally degrade the TaskNet, while the TaskNet learns robust representations under these perturbations. The enrichment via spectral-clustered and kNN-based edges, combined with adversarial pruning, yields a practical approach that achieves new state-of-the-art results on diverse graphs, including citation networks, social networks, and temporal graphs. The authors provide a robust optimization interpretation and comprehensive experiments, demonstrating the value of adaptive topology search over static augmentations for cross-domain generalization in GNNs.

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\%, 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/
Paper Structure (50 sections, 3 theorems, 48 equations, 3 figures, 12 tables, 1 algorithm)

This paper contains 50 sections, 3 theorems, 48 equations, 3 figures, 12 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Preliminaries
  4. Proposed Approach for Cross-Domain Generalization: EdgeMask-DG*
  5. Enriched Graph Representation
  6. Adversarial Edge Masking
  7. GAT Backbone with Edge Attributes
  8. Training Procedure
  9. Theoretical Grounding: Robust Optimization Perspective
  10. Computational Complexity
  11. Augment-then-prune rationale.
  12. Experiments
  13. Experimental Setup
  14. Main Results
  15. Performance on Citation Networks (ACM, DBLP, Citation)
  16. ...and 35 more sections

Key Result

Lemma 4.1

Let $\mathbf{s}^*$ be an optimal solution to the mask optimization problem equation eq:P_main for a fixed $\theta$. Assuming continuous differentiability of $\ell(f_\theta, \mathbf{s})$ w.r.t. $\mathbf{s}$ and constraint qualifications (Slater's condition holds), there exists an optimal dual variabl

Figures (3)

  • Figure 1: EdgeMask-DG*: (1) enrich the graph with kNN & spectral edges; (2) play a min–max game where MaskNet sparsifies edges and TaskNet (GAT) learns to stay accurate.
  • Figure 2: Average performance vs. kNN K (number of neighbors). Results averaged over 3 leave-one-out scenarios.
  • Figure 3: Average performance and training time vs. spectral $K$. F1 scores (left axis) and training time (right axis) averaged over three leave-one-out runs (note: $k=50$ omitted due to OOM).

Theorems & Definitions (5)

  • Lemma 4.1: Optimality Conditions for Adversarial Mask
  • Lemma C.1: Lagrangian Upper Bound
  • proof
  • Proposition C.2: Optimal Mask for Penalized Surrogate
  • proof