# Sampling-based Distributed Training with Message Passing Neural Network

arXiv: https://arxiv.org/abs/2402.15106 · HTML: https://mlsanity.com/paper/2402.15106/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 20 of its 25 references that are in the map.
- First posted on arXiv 23 Feb 2024.

## Builds on (most cited first, 20 of 20)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Semi-Supervised Classification with Graph Convolutional Networks](https://mlsanity.com/paper/1609.02907/) (2016)
- [Graph Attention Networks](https://mlsanity.com/paper/1710.10903/) (2017)
- [Inductive Representation Learning on Large Graphs](https://mlsanity.com/paper/1706.02216/) (2017)
- [Neural Message Passing for Quantum Chemistry](https://mlsanity.com/paper/1704.01212/) (2017)
- [Fourier Neural Operator for Parametric Partial Differential Equations](https://mlsanity.com/paper/2010.08895/) (2020)
- [Fast Graph Representation Learning with PyTorch Geometric](https://mlsanity.com/paper/1903.02428/) (2019)
- [DeepONet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators](https://mlsanity.com/paper/1910.03193/) (2019)
- [Learning to Simulate Complex Physics with Graph Networks](https://mlsanity.com/paper/2002.09405/) (2020)
- [Learning Mesh-Based Simulation with Graph Networks](https://mlsanity.com/paper/2010.03409/) (2020)
- [Convolutional Networks on Graphs for Learning Molecular Fingerprints](https://mlsanity.com/paper/1509.09292/) (2015)
- [Neural Operator: Graph Kernel Network for Partial Differential Equations](https://mlsanity.com/paper/2003.03485/) (2020)
- [Dynamic Edge-Conditioned Filters in Convolutional Neural Networks on Graphs](https://mlsanity.com/paper/1704.02901/) (2017)
- [GraphCast: Learning skillful medium-range global weather forecasting](https://mlsanity.com/paper/2212.12794/) (2022)
- [Physics-Constrained Deep Learning for High-dimensional Surrogate Modeling and Uncertainty Quantification without Labeled Data](https://mlsanity.com/paper/1901.06314/) (2019)
- [PhyCRNet: Physics-informed Convolutional-Recurrent Network for Solving Spatiotemporal PDEs](https://mlsanity.com/paper/2106.14103/) (2021)
- [AirfRANS: High Fidelity Computational Fluid Dynamics Dataset for Approximating Reynolds-Averaged Navier-Stokes Solutions](https://mlsanity.com/paper/2212.07564/) (2022)
- [A composable machine-learning approach for steady-state simulations on high-resolution grids](https://mlsanity.com/paper/2210.05837/) (2022)
- [An extensible Benchmarking Graph-Mesh dataset for studying Steady-State Incompressible Navier-Stokes Equations](https://mlsanity.com/paper/2206.14709/) (2022)
- [Multi-GPU Approach for Training of Graph ML Models on large CFD Meshes](https://mlsanity.com/paper/2307.13592/) (2023)

## Built on by (most cited first, 2 of 2)

- [DoMINO: A Decomposable Multi-scale Iterative Neural Operator for Modeling Large Scale Engineering Simulations](https://mlsanity.com/paper/2501.13350/) (2025)
- [Scalable and Consistent Graph Neural Networks for Distributed Mesh-based Data-driven Modeling](https://mlsanity.com/paper/2410.01657/) (2024)

## Abstract (arXiv)

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(105) 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). The DS-MPNN model demonstrates comparable accuracy to single-GPU implementation, can accommodate a significantly larger number of nodes compared to the single-GPU variant (S-MPNN), and significantly outperforms the node-based GCN.