# Learning Divergence Fields for Shift-Robust Graph Representations

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

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 28 of its 48 references that are in the map.
- First posted on arXiv 7 Jun 2024.

## Builds on (most cited first, 28 of 28)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [A Simple Framework for Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2002.05709/) (2020)
- [Semi-Supervised Classification with Graph Convolutional Networks](https://mlsanity.com/paper/1609.02907/) (2016)
- [Graph Attention Networks](https://mlsanity.com/paper/1710.10903/) (2017)
- [mixup: Beyond Empirical Risk Minimization](https://mlsanity.com/paper/1710.09412/) (2017)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [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)
- [ImageNet-trained CNNs are biased towards texture; increasing shape bias improves accuracy and robustness](https://mlsanity.com/paper/1811.12231/) (2018)
- [WILDS: A Benchmark of in-the-Wild Distribution Shifts](https://mlsanity.com/paper/2012.07421/) (2020)
- [Distributionally Robust Neural Networks for Group Shifts: On the Importance of Regularization for Worst-Case Generalization](https://mlsanity.com/paper/1911.08731/) (2019)
- [Recognition in Terra Incognita](https://mlsanity.com/paper/1807.04975/) (2018)
- [Domain Generalization via Invariant Feature Representation](https://mlsanity.com/paper/1301.2115/) (2013)
- [Importance Weighted Autoencoders](https://mlsanity.com/paper/1509.00519/) (2015)
- [Towards Deeper Graph Neural Networks](https://mlsanity.com/paper/2007.09296/) (2020)
- [VAE with a VampPrior](https://mlsanity.com/paper/1705.07120/) (2017)
- [Handling Distribution Shifts on Graphs: An Invariance Perspective](https://mlsanity.com/paper/2202.02466/) (2022)
- [Uncertainty Aware Semi-Supervised Learning on Graph Data](https://mlsanity.com/paper/2010.12783/) (2020)
- [DIFFormer: Scalable (Graph) Transformers Induced by Energy Constrained Diffusion](https://mlsanity.com/paper/2301.09474/) (2023)
- [Beltrami Flow and Neural Diffusion on Graphs](https://mlsanity.com/paper/2110.09443/) (2021)
- [Graph Posterior Network: Bayesian Predictive Uncertainty for Node Classification](https://mlsanity.com/paper/2110.14012/) (2021)
- [Subgroup Generalization and Fairness of Graph Neural Networks](https://mlsanity.com/paper/2106.15535/) (2021)
- [Energy-based Out-of-Distribution Detection for Graph Neural Networks](https://mlsanity.com/paper/2302.02914/) (2023)
- [Twitch Gamers: a Dataset for Evaluating Proximity Preserving and Structural Role-based Node Embeddings](https://mlsanity.com/paper/2101.03091/) (2021)
- [Geometric Knowledge Distillation: Topology Compression for Graph Neural Networks](https://mlsanity.com/paper/2210.13014/) (2022)
- [DPPIN: A Biological Repository of Dynamic Protein-Protein Interaction Network Data](https://mlsanity.com/paper/2107.02168/) (2021)
- [Towards Out-of-Distribution Sequential Event Prediction: A Causal Treatment](https://mlsanity.com/paper/2210.13005/) (2022)
- [The nonlinear heat equation on dense graphs and graph limits](https://mlsanity.com/paper/1302.5804/) (2013)

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

- [SCL-GNN: Towards Generalizable Graph Neural Networks via Spurious Correlation Learning](https://mlsanity.com/paper/2603.08270/) (2026)
- [CGRL: Causal-Guided Representation Learning for Node-Level Out-of-Distribution Generalization](https://mlsanity.com/paper/2603.24304/) (2026)

## Abstract (arXiv)

Real-world data generation often involves certain geometries (e.g., graphs) that induce instance-level interdependence. This characteristic makes the generalization of learning models more difficult due to the intricate interdependent patterns that impact data-generative distributions and can vary from training to testing. In this work, we propose a geometric diffusion model with learnable divergence fields for the challenging generalization problem with interdependent data. We generalize the diffusion equation with stochastic diffusivity at each time step, which aims to capture the multi-faceted information flows among interdependent data. Furthermore, we derive a new learning objective through causal inference, which can guide the model to learn generalizable patterns of interdependence that are insensitive across domains. Regarding practical implementation, we introduce three model instantiations that can be considered as the generalized versions of GCN, GAT, and Transformers, respectively, which possess advanced robustness against distribution shifts. We demonstrate their promising efficacy for out-of-distribution generalization on diverse real-world datasets.