Adaptive Riemannian Graph Neural Networks
Xudong Wang, Chris Ding, Tongxin Li, Jicong Fan
TL;DR
ARGNN introduces a node-adaptive Riemannian framework for graphs by learning a diagonal metric tensor $\mathbf{G}_i=\text{diag}(\mathbf{g}_i)$ at each node, enabling anisotropic, continuous geometry. A local metric estimator computes $\mathbf{g}_i$ from $(\mathbf{h}_i, \mathbf{a}_i)$, while geometry-aware message passing employs geodesic distances $d_{\mathbf{G}_i}$, geometric modulation $\tau_{ij}$, and geometry-aware attention $\alpha_{ij}$ to update node representations. Stability and expressiveness are ensured via Ricci-flow-inspired regularization and a universal approximation analysis showing the framework subsumes fixed-curvature and product-manifold GNNs as special cases. Theoretical results yield convergence guarantees and efficiency, and experiments across homophilic and heterophilic graphs show ARGNN consistently outperforms state-of-the-art baselines while offering interpretable learned geometries.
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.