Learning Divergence Fields for Shift-Robust Graph Representations

TL;DR

This paper introduces a geometric diffusion model with learnable divergence fields designed to improve the generalization of graph representations across different data distributions, especially under distribution shifts.

Contribution

It proposes a novel diffusion framework with stochastic diffusivity and a causal inference-based learning objective for shift-robust graph models, along with three instantiations inspired by GCN, GAT, and Transformers.

Findings

Enhanced out-of-distribution generalization on real-world datasets

Robustness of the proposed models against distribution shifts

Effective modeling of interdependent data patterns

Abstract

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.