# Distributed Multigrid Neural Solvers on Megavoxel Domains

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

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Builds on 24 of its 54 references that are in the map.
- First posted on arXiv 29 Apr 2021.

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

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [3D U-Net: Learning Dense Volumetric Segmentation from Sparse Annotation](https://mlsanity.com/paper/1606.06650/) (2016)
- [DGM: A deep learning algorithm for solving partial differential equations](https://mlsanity.com/paper/1708.07469/) (2017)
- [Solving high-dimensional partial differential equations using deep learning](https://mlsanity.com/paper/1707.02568/) (2017)
- [Physics-informed neural networks with hard constraints for inverse design](https://mlsanity.com/paper/2102.04626/) (2021)
- [Horovod: fast and easy distributed deep learning in TensorFlow](https://mlsanity.com/paper/1802.05799/) (2018)
- [Physics-Constrained Deep Learning for High-dimensional Surrogate Modeling and Uncertainty Quantification without Labeled Data](https://mlsanity.com/paper/1901.06314/) (2019)
- [hp-VPINNs: Variational Physics-Informed Neural Networks With Domain Decomposition](https://mlsanity.com/paper/2003.05385/) (2020)
- [Hidden Physics Models: Machine Learning of Nonlinear Partial Differential Equations](https://mlsanity.com/paper/1708.00588/) (2017)
- [Bayesian Deep Convolutional Encoder-Decoder Networks for Surrogate Modeling and Uncertainty Quantification](https://mlsanity.com/paper/1801.06879/) (2018)
- [Demystifying Parallel and Distributed Deep Learning: An In-Depth Concurrency Analysis](https://mlsanity.com/paper/1802.09941/) (2018)
- [Physics-Informed Generative Adversarial Networks for Stochastic Differential Equations](https://mlsanity.com/paper/1811.02033/) (2018)
- [Accelerating Eulerian Fluid Simulation With Convolutional Networks](https://mlsanity.com/paper/1607.03597/) (2016)
- [An Energy Approach to the Solution of Partial Differential Equations in Computational Mechanics via Machine Learning: Concepts, Implementation and Applications](https://mlsanity.com/paper/1908.10407/) (2019)
- [NVIDIA SimNet^{TM}: an AI-accelerated multi-physics simulation framework](https://mlsanity.com/paper/2012.07938/) (2020)
- [Learning Algebraic Multigrid Using Graph Neural Networks](https://mlsanity.com/paper/2003.05744/) (2020)
- [A Multigrid Method for Efficiently Training Video Models](https://mlsanity.com/paper/1912.00998/) (2019)
- [DiscretizationNet: A Machine-Learning based solver for Navier-Stokes Equations using Finite Volume Discretization](https://mlsanity.com/paper/2005.08357/) (2020)
- [Learning optimal multigrid smoothers via neural networks](https://mlsanity.com/paper/2102.12071/) (2021)
- [Highly-scalable, physics-informed GANs for learning solutions of stochastic PDEs](https://mlsanity.com/paper/1910.13444/) (2019)
- [Solving Irregular and Data-enriched Differential Equations using Deep Neural Networks](https://mlsanity.com/paper/1905.04351/) (2019)
- [Deep Generative Models that Solve PDEs: Distributed Computing for Training Large Data-Free Models](https://mlsanity.com/paper/2007.12792/) (2020)
- [Translating Numerical Concepts for PDEs into Neural Architectures](https://mlsanity.com/paper/2103.15419/) (2021)

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

- [Scientific Machine Learning through Physics-Informed Neural Networks: Where we are and What's next](https://mlsanity.com/paper/2201.05624/) (2022)
- [Geometry-Informed Neural Operator for Large-Scale 3D PDEs](https://mlsanity.com/paper/2309.00583/) (2023)
- [NeuFENet: Neural Finite Element Solutions with Theoretical Bounds for Parametric PDEs](https://mlsanity.com/paper/2110.01601/) (2021)
- [Neural PDE Solvers for Irregular Domains](https://mlsanity.com/paper/2211.03241/) (2022)
- [SciAI4Industry -- Solving PDEs for industry-scale problems with deep learning](https://mlsanity.com/paper/2211.12709/) (2022)

## Abstract (arXiv)

We consider the distributed training of large-scale neural networks that serve as PDE solvers producing full field outputs. We specifically consider neural solvers for the generalized 3D Poisson equation over megavoxel domains. A scalable framework is presented that integrates two distinct advances. First, we accelerate training a large model via a method analogous to the multigrid technique used in numerical linear algebra. Here, the network is trained using a hierarchy of increasing resolution inputs in sequence, analogous to the 'V', 'W', 'F', and 'Half-V' cycles used in multigrid approaches. In conjunction with the multi-grid approach, we implement a distributed deep learning framework which significantly reduces the time to solve. We show the scalability of this approach on both GPU (Azure VMs on Cloud) and CPU clusters (PSC Bridges2). This approach is deployed to train a generalized 3D Poisson solver that scales well to predict output full-field solutions up to the resolution of 512x512x512 for a high dimensional family of inputs.