Fast and Distributed Equivariant Graph Neural Networks by Virtual Node Learning

Area of Science:

• Geometric deep learning
• Scientific machine learning
• Graph neural networks

Background:

• Equivariant Graph Neural Networks (GNNs) excel in scientific applications but struggle with large, sparse, or distributed geometric graphs.
• Existing GNNs face efficiency limitations and performance degradation on computationally tractable, sparsified graphs.

Purpose of the Study:

• To introduce novel enhancements, FastEGNN and DistEGNN, for efficient and scalable equivariant GNNs on large geometric graphs.
• To address the computational challenges of processing large-scale and distributed geometric graph data.

Main Methods:

• FastEGNN utilizes a small set of virtual nodes to approximate large graphs, employing distinct message passing and aggregation for virtual nodes.
• Maximum Mean Discrepancy (MMD) is minimized between virtual and real coordinates for global distributedness in FastEGNN.
• DistEGNN extends FastEGNN for extremely large graphs by using virtual nodes as distributed bridges between subgraphs, reducing overhead.

Main Results:

• FastEGNN demonstrates high accuracy and efficiency on large-scale sparse graphs.
• DistEGNN effectively reduces memory and computational costs for extremely large-scale distributed geometric graphs.
• Evaluations across N-body systems, protein dynamics, Water-3D, and the Fluid113K benchmark show superior efficiency and performance.

Conclusions:

• FastEGNN and DistEGNN establish new capabilities for large-scale equivariant graph learning.
• The proposed methods significantly improve efficiency and performance for geometric GNNs on large and distributed datasets.
• Open-source code is provided for reproducibility and further research.