Sampling-based Distributed Training with Message Passing Neural Network

Understand

In this study, we introduce a domain-decomposition-based distributed training and inference approach for message-passing neural networks (MPNN).

• Our objective is to address the challenge of scaling edge-based graph neural networks as the number of nodes increases.
• Through our distributed training approach, coupled with Nystr\"om-approximation sampling techniques, we present a scalable graph neural network, referred to as DS-MPNN (D and S standing for distributed and sampled, respectively), capable of scaling up to $O(10^5)$ nodes.
• We validate our sampling and distributed training approach on two cases: (a) a Darcy flow dataset and (b) steady RANS simulations of 2-D airfoils, providing comparisons with both single-GPU implementation and node-based graph convolution networks (GCNs).