This paper presents a theoretical framework to determine the optimal depth of Graph Neural Networks (GNNs) by analyzing Optimal Transport (OT) geodesics in the space of probability measures over graph node representations. It establishes a connection between GNN message-passing and Wasserstein gradient flows, characterizing the evolution of node representations as geodesic curves and deriving a closed-form expression for optimal depth based on graph spectral properties and curvature.
Key findings
Optimal GNN depth corresponds to the geodesic length in Wasserstein space that balances representational expressivity against over-smoothing degradation.
The theoretical results yield a closed-form characterization of optimal depth as a function of graph spectral properties, curvature, and task-specific information requirements.
Depth selection guided by OT geodesics outperforms heuristic approaches in extensive experiments on benchmark datasets.
Limitations & open questions
The framework's applicability may be limited to specific types of GNN architectures and graph structures.
Empirical validation is based on benchmark datasets, and further testing on diverse real-world graphs is needed.