# Topology-Aware Dynamic Reweighting for Distribution Shifts on Graph

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

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 32 of its 44 references that are in the map.
- First posted on arXiv 3 Jun 2024.

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

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Semi-Supervised Classification with Graph Convolutional Networks](https://mlsanity.com/paper/1609.02907/) (2016)
- [Graph Attention Networks](https://mlsanity.com/paper/1710.10903/) (2017)
- [Domain-Adversarial Training of Neural Networks](https://mlsanity.com/paper/1505.07818/) (2015)
- [Fast Graph Representation Learning with PyTorch Geometric](https://mlsanity.com/paper/1903.02428/) (2019)
- [Invariant Risk Minimization](https://mlsanity.com/paper/1907.02893/) (2019)
- [Distributionally Robust Neural Networks for Group Shifts: On the Importance of Regularization for Worst-Case Generalization](https://mlsanity.com/paper/1911.08731/) (2019)
- [Class-Balanced Loss Based on Effective Number of Samples](https://mlsanity.com/paper/1901.05555/) (2019)
- [Deep CORAL: Correlation Alignment for Deep Domain Adaptation](https://mlsanity.com/paper/1607.01719/) (2016)
- [Geom-GCN: Geometric Graph Convolutional Networks](https://mlsanity.com/paper/2002.05287/) (2020)
- [Out-of-Distribution Generalization via Risk Extrapolation (REx)](https://mlsanity.com/paper/2003.00688/) (2020)
- [GNNExplainer: Generating Explanations for Graph Neural Networks](https://mlsanity.com/paper/1903.03894/) (2019)
- [Towards Out-Of-Distribution Generalization: A Survey](https://mlsanity.com/paper/2108.13624/) (2021)
- [Certifying Some Distributional Robustness with Principled Adversarial Training](https://mlsanity.com/paper/1710.10571/) (2017)
- [Deep Gaussian Embedding of Graphs: Unsupervised Inductive Learning via Ranking](https://mlsanity.com/paper/1707.03815/) (2017)
- [Learning Models with Uniform Performance via Distributionally Robust Optimization](https://mlsanity.com/paper/1810.08750/) (2018)
- [Discovering Invariant Rationales for Graph Neural Networks](https://mlsanity.com/paper/2201.12872/) (2022)
- [Rethinking Graph Neural Networks for Anomaly Detection](https://mlsanity.com/paper/2205.15508/) (2022)
- [Handling Distribution Shifts on Graphs: An Invariance Perspective](https://mlsanity.com/paper/2202.02466/) (2022)
- [GOOD: A Graph Out-of-Distribution Benchmark](https://mlsanity.com/paper/2206.08452/) (2022)
- [Characteristic Functions on Graphs: Birds of a Feather, from Statistical Descriptors to Parametric Models](https://mlsanity.com/paper/2005.07959/) (2020)
- [Heterogeneous Risk Minimization](https://mlsanity.com/paper/2105.03818/) (2021)
- [Shift-Robust GNNs: Overcoming the Limitations of Localized Graph Training Data](https://mlsanity.com/paper/2108.01099/) (2021)
- [TAM: Topology-Aware Margin Loss for Class-Imbalanced Node Classification](https://mlsanity.com/paper/2206.12917/) (2022)
- [OOD Link Prediction Generalization Capabilities of Message-Passing GNNs in Larger Test Graphs](https://mlsanity.com/paper/2205.15117/) (2022)
- [Unleashing the Power of Graph Data Augmentation on Covariate Distribution Shift](https://mlsanity.com/paper/2211.02843/) (2022)
- [Mind the Label Shift of Augmentation-based Graph OOD Generalization](https://mlsanity.com/paper/2303.14859/) (2023)
- [Adaptive Graph Contrastive Learning for Recommendation](https://mlsanity.com/paper/2305.10837/) (2023)
- [OpenGSL: A Comprehensive Benchmark for Graph Structure Learning](https://mlsanity.com/paper/2306.10280/) (2023)
- [Debiased Graph Neural Networks with Agnostic Label Selection Bias](https://mlsanity.com/paper/2201.07708/) (2022)

## Built on by (most cited first, 3 of 3)

- [Data Heterogeneity Modeling for Trustworthy Machine Learning](https://mlsanity.com/paper/2506.00969/) (2025)
- [When Molecular Similarity Works: Property Cliffs Reveal Hidden Errors](https://mlsanity.com/paper/2605.17265/) (2026)
- [A Unified Uncertainty Representation for Graph Neural Networks via Doubly-Spectral Stochastic Expansion](https://mlsanity.com/paper/2609.35703/) (2026)

## Abstract (arXiv)

Graph Neural Networks (GNNs) are widely used for node classification tasks but often fail to generalize when training and test nodes come from different distributions, limiting their practicality. To overcome this, recent approaches adopt invariant learning techniques from the out-of-distribution (OOD) generalization field, which seek to establish stable prediction methods across environments. However, the applicability of these invariant assumptions to graph data remains unverified, and such methods often lack solid theoretical support. In this work, we introduce the Topology-Aware Dynamic Reweighting (TAR) framework, which dynamically adjusts sample weights through gradient flow in the geometric Wasserstein space during training. Instead of relying on strict invariance assumptions, we prove that our method is able to provide distributional robustness, thereby enhancing the out-of-distribution generalization performance on graph data. By leveraging the inherent graph structure, TAR effectively addresses distribution shifts. Our framework's superiority is demonstrated through standard testing on four graph OOD datasets and three class-imbalanced node classification datasets, exhibiting marked improvements over existing methods.
