# Breaking Boundaries: Distributed Domain Decomposition with Scalable Physics-Informed Neural PDE Solvers

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

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 25 of its 63 references that are in the map.
- First posted on arXiv 28 Aug 2023.

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

- [Gaussian Error Linear Units (GELUs)](https://mlsanity.com/paper/1606.08415/) (2016)
- [Fourier Neural Operator for Parametric Partial Differential Equations](https://mlsanity.com/paper/2010.08895/) (2020)
- [DeepONet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators](https://mlsanity.com/paper/1910.03193/) (2019)
- [Large Batch Optimization for Deep Learning: Training BERT in 76 minutes](https://mlsanity.com/paper/1904.00962/) (2019)
- [Characterizing possible failure modes in physics-informed neural networks](https://mlsanity.com/paper/2109.01050/) (2021)
- [When and why PINNs fail to train: A neural tangent kernel perspective](https://mlsanity.com/paper/2007.14527/) (2020)
- [DeepXDE: A deep learning library for solving differential equations](https://mlsanity.com/paper/1907.04502/) (2019)
- [Learning the solution operator of parametric partial differential equations with physics-informed DeepOnets](https://mlsanity.com/paper/2103.10974/) (2021)
- [Neural Operator: Learning Maps Between Function Spaces](https://mlsanity.com/paper/2108.08481/) (2021)
- [PDEBENCH: An Extensive Benchmark for Scientific Machine Learning](https://mlsanity.com/paper/2210.07182/) (2022)
- [Machine Learning for Fluid Mechanics](https://mlsanity.com/paper/1905.11075/) (2019)
- [hp-VPINNs: Variational Physics-Informed Neural Networks With Domain Decomposition](https://mlsanity.com/paper/2003.05385/) (2020)
- [Finite Basis Physics-Informed Neural Networks (FBPINNs): a scalable domain decomposition approach for solving differential equations](https://mlsanity.com/paper/2107.07871/) (2021)
- [Parallel Physics-Informed Neural Networks via Domain Decomposition](https://mlsanity.com/paper/2104.10013/) (2021)
- [PyTorch Distributed: Experiences on Accelerating Data Parallel Training](https://mlsanity.com/paper/2006.15704/) (2020)
- [The Old and the New: Can Physics-Informed Deep-Learning Replace Traditional Linear Solvers?](https://mlsanity.com/paper/2103.09655/) (2021)
- [NVIDIA SimNet^{TM}: an AI-accelerated multi-physics simulation framework](https://mlsanity.com/paper/2012.07938/) (2020)
- [D3M: A deep domain decomposition method for partial differential equations](https://mlsanity.com/paper/1909.12236/) (2019)
- [CFDNet: a deep learning-based accelerator for fluid simulations](https://mlsanity.com/paper/2005.04485/) (2020)
- [Mosaic Flows: A Transferable Deep Learning Framework for Solving PDEs on Unseen Domains](https://mlsanity.com/paper/2104.10873/) (2021)
- [Distributed physics informed neural network for data-efficient solution to partial differential equations](https://mlsanity.com/paper/1907.08967/) (2019)
- [BubbleML: A Multi-Physics Dataset and Benchmarks for Machine Learning](https://mlsanity.com/paper/2307.14623/) (2023)
- [MLPerf HPC: A Holistic Benchmark Suite for Scientific Machine Learning on HPC Systems](https://mlsanity.com/paper/2110.11466/) (2021)
- [SURFNet: Super-resolution of Turbulent Flows with Transfer Learning using Small Datasets](https://mlsanity.com/paper/2108.07667/) (2021)
- [Scalable communication for high-order stencil computations using CUDA-aware MPI](https://mlsanity.com/paper/2103.01597/) (2021)

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

- [Machine learning and domain decomposition methods -- a survey](https://mlsanity.com/paper/2312.14050/) (2023)
- [A Flow-rate-conserving CNN-based Domain Decomposition Method for Blood Flow Simulations](https://mlsanity.com/paper/2509.15900/) (2025)
- [Localized PCA-Net Neural Operators for Scalable Solution Reconstruction of Elliptic PDEs](https://mlsanity.com/paper/2509.18110/) (2025)

## Abstract (arXiv)

Mosaic Flow is a novel domain decomposition method designed to scale physics-informed neural PDE solvers to large domains. Its unique approach leverages pre-trained networks on small domains to solve partial differential equations on large domains purely through inference, resulting in high reusability. This paper presents an end-to-end parallelization of Mosaic Flow, combining data parallel training and domain parallelism for inference on large-scale problems. By optimizing the network architecture and data parallel training, we significantly reduce the training time for learning the Laplacian operator to minutes on 32 GPUs. Moreover, our distributed domain decomposition algorithm enables scalable inferences for solving the Laplace equation on domains 4096 times larger than the training domain, demonstrating strong scaling while maintaining accuracy on 32 GPUs. The reusability of Mosaic Flow, combined with the improved performance achieved through the distributed-memory algorithms, makes it a promising tool for modeling complex physical phenomena and accelerating scientific discovery.