Adaptive Riemannian Graph Neural Networks

TL;DR

Adaptive Riemannian Graph Neural Networks (ARGNN) dynamically learn local geometric structures of graphs, enabling better modeling of complex heterogeneity with theoretical guarantees and improved empirical performance.

Contribution

Introduces a novel framework that learns node-specific Riemannian metrics, unifying fixed and mixed-curvature GNNs with a regularization for stable training.

Findings

Outperforms existing GNNs on diverse benchmarks

Learns interpretable local geometries

Provides theoretical convergence guarantees

Abstract

Graph data often exhibits complex geometric heterogeneity, where structures with varying local curvature, such as tree-like hierarchies and dense communities, coexist within a single network. Existing geometric GNNs, which embed graphs into single fixed-curvature manifolds or discrete product spaces, struggle to capture this diversity. We introduce Adaptive Riemannian Graph Neural Networks (ARGNN), a novel framework that learns a continuous and anisotropic Riemannian metric tensor field over the graph. It allows each node to determine its optimal local geometry, enabling the model to fluidly adapt to the graph's structural landscape. Our core innovation is an efficient parameterization of the node-wise metric tensor, specializing to a learnable diagonal form that captures directional geometric information while maintaining computational tractability. To ensure geometric regularity and stable training, we integrate a Ricci flow-inspired regularization that smooths the learned manifold. Theoretically, we establish the rigorous geometric evolution convergence guarantee for ARGNN and provide a continuous generalization that unifies prior fixed or mixed-curvature GNNs. Empirically, our method demonstrates superior performance on both homophilic and heterophilic benchmark datasets with the ability to capture diverse structures adaptively. Moreover, the learned geometries both offer interpretable insights into the underlying graph structure and empirically corroborate our theoretical analysis.