PaperlayerView exact revision on arXiv

Source-linked AI summary

Parallel Physics-Informed Neural Networks via Domain Decomposition

Khemraj Shukla, Ameya D. Jagtap, George Em Karniadakis

arXiv:2104.10013v3cs.DC

TL;DR

PINN training is computationally costly, motivating distributed domain decomposition for cPINNs and XPINNs. The paper develops MPI + CPU/GPU algorithms and evaluates their scaling and applications. cPINNs reduce communication cost for spatial decomposition, while XPINNs offer broader decomposition and geometry flexibility, including a ten-region inverse conductivity problem.

  • Problem

    PINNs have high neural-network training costs, particularly for forward multi-scale problems, motivating distributed domain-decomposition approaches.

  • Method

    The paper develops unified distributed cPINN and XPINN algorithms using MPI + CPUs or GPUs, with physics-based interface conditions and subdomain-specific hyperparameters.

  • Results

    For spatial decomposition, cPINNs are more communication-efficient, while XPINNs handle space-time decomposition, arbitrary complex subdomains, and an inverse conductivity problem on ten U.S. regions.

  • Takeaways & Limitations

    Domain decomposition provides parallelization and flexibility for PINNs across forward and inverse scientific-computing problems.

  • Takeaways & Limitations

    The manually partitioned irregular-domain test can suffer load imbalance, and optimized point-cloud or K-way partitioning could improve communication.

Abstract

from arXiv · show

We develop a distributed framework for the physics-informed neural networks (PINNs) based on two recent extensions, namely conservative PINNs (cPINNs) and extended PINNs (XPINNs), which employ domain decomposition in space and in time-space, respectively. This domain decomposition endows cPINNs and XPINNs with several advantages over the vanilla PINNs, such as parallelization capacity, large representation capacity, efficient hyperparameter tuning, and is particularly effective for multi-scale and multi-physics problems. Here, we present a parallel algorithm for cPINNs and XPINNs constructed with a hybrid programming model described by MPI $+$ X, where X $\in \{\text{CPUs},~\text{GPUs}\}$. The main advantage of cPINN and XPINN over the more classical data and model parallel approaches is the flexibility of optimizing all hyperparameters of each neural network separately in each subdomain. We compare the performance of distributed cPINNs and XPINNs for various forward problems, using both weak and strong scalings. Our results indicate that for space domain decomposition, cPINNs are more efficient in terms of communication cost but XPINNs provide greater flexibility as they can also handle time-domain decomposition for any differential equations, and can deal with any arbitrarily shaped complex subdomains. To this end, we also present an application of the parallel XPINN method for solving an inverse diffusion problem with variable conductivity on the United States map, using ten regions as subdomains.

1. Introduction

PINNs address scientific forward and inverse problems, but their training cost motivates domain-decomposed cPINNs and XPINNs. This paper develops distributed implementations and compares their scaling, flexibility, and performance with conventional parallel approaches.

  • PINNs infer solutions to governing physical laws and identify unknown parameters from multimodal measurements in forward and inverse problems.
  • Large training costs, especially for forward multi-scale problems, are a major limitation of vanilla PINNs.
  • cPINNs decompose space and enforce state and flux continuity across interfaces, whereas XPINNs support space-time decomposition for general PDEs and irregular geometries.
  • Data parallelism replicates one model across processors, while model parallelism distributes layers and pipelines batches through them.
  • Unlike data parallelism, cPINNs and XPINNs allow neural-network hyperparameters to be optimized separately in each subdomain.
  • The paper presents unified CPU- and GPU-based distributed algorithms, including weak scaling, strong scaling, and comparisons across forward and inverse problems.

2. Problem setup

The paper formulates parametrized PDEs with differential operators, boundary conditions, forcing terms, and model parameters, then converts their solution into a neural-network optimization problem.

  • A parametrized PDE combines a differential operator, solution, model parameters, boundary conditions, and forcing term over a spatial or spatio-temporal domain.
  • For transient problems, time is treated as a component of the input and initial conditions as boundary conditions on the spatio-temporal domain.
  • The paper treats forward problems with fixed parameters and inverse problems that learn unknown parameters from observed data.
  • PINNs convert PDE solution into an optimization problem whose loss-function minima correspond to PDE solutions.

3. Fully connected feed-forward neural networks

The paper represents solutions with fully connected feed-forward neural networks defined by layered weights, biases, activations, and trainable activation slopes. These adaptive slopes are intended to improve training speed and learning capacity.

  • A fully connected network maps R^Di to R^Do through L layers, with N_k neurons in layer k and identity activation in the last layer.
  • Weights W_k and biases b_k parameterize each layer, while the input is z and the layer output is N_k(z).
  • Layer-wise trainable activation slopes change activation-function steepness, increase training speed, and enhance learning capacity especially early in training.
  • The parameter collection ˜Θ includes all weights, biases, and activation slopes, defining the neural-network parameter space.

4. Brief overview of the PINNs, cPINNs and XPINNs

PINNs learn PDE solutions by minimizing losses that combine data mismatch with PDE residuals. cPINNs and XPINNs extend this framework through subdomain-specific networks coupled by interface conditions.

  • PINNs: PINNs learn a surrogate PDE solution by minimizing a loss built from training data and governing-equation residuals.The data term enforces initial and boundary conditions, while the residual term evaluates the PDE operator.
  • PINNs: Automatic differentiation computes PDE residual derivatives without truncation or round-off errors from numerical differentiation.This enables a grid-free formulation that avoids mesh generation.
  • cPINNs: cPINNs divide the domain into non-overlapping subdomains with separate neural networks and subdomain-wise loss functions.Their interface losses enforce average solution continuity and normal flux continuity.
  • cPINNs: Interface conditions propagate information between neighboring subdomains and support convergence where training data are unavailable.The cPINN formulation includes average solution and flux continuity terms at common interfaces.
  • XPINNs: XPINNs add residual continuity to average solution continuity and can impose further interface constraints suited to the differential equation and interface orientation.Examples include flux, higher-order derivative, and invariant continuity.

5. Parallel implementation of cPINN and XPINN

The parallel implementation decomposes the computational domain across processors or accelerators, assigns a neural network to each subdomain, and alternates concurrent computation with neighbor communication.

  • Parallel implementation: The implementation uses preprocessing and parallel solution stages for distributed heterogeneous CPU and GPU execution.The domain is decomposed into subdomains corresponding to available processors or accelerators.
  • Pre-processing: The domain is divided into non-overlapping subdomains, with training, residual, and interface points prepared for each subdomain.The number of subdomains equals the number of processors or accelerators, and irregular geometries are supported.
  • Pre-processing: Each subdomain uses an individual neural network and loss function, with interface conditions coupling neighboring networks.The network parameters are updated using gradients computed separately for each subdomain.
  • Parallel solution: Subdomains concurrently compute network outputs, fluxes, and PDE residuals because these terms require no neighboring data.These quantities are computed during the computation stage of each epoch.
  • Parallel solution: Processors communicate with neighboring subdomains through a logical two-dimensional layout and non-blocking point-to-point send-receive operations.Neighbor ranks are determined from local processor coordinates and directional edges.
  • Parallel solution: After communication and synchronization, each subdomain computes its loss and independently optimizes its neural-network parameters.GPU mappings use CUDA-aware MPI.

6. Optimization on distributed systems

The distributed cPINN and XPINN formulation optimizes separate subdomain losses without requiring global convergence guarantees. Experiments indicate accurate solutions under well-posedness and sufficient network expressivity.

  • Optimization: Each subdomain minimizes its own loss function, allowing distributed optimization of separate neural-network parameter sets.The optimization process has no theoretical guarantee of reaching a global minimum.
  • Optimization: Accurate solutions are observed when the PDE is well-posed with a unique solution and the networks and residual-point sets are sufficiently expressive.This statement is based on the authors’ computational experiments rather than a convergence theorem.
  • Optimization: Distributed optimization procedures for single-loss-function methods can be incorporated without significant code changes.The paper situates its approach alongside existing distributed optimization procedures.

7. Computational experiments

The experiments define weak and strong scaling metrics, describe the parallel algorithm’s stages, and profile PINN computation costs across residual-point and network-size changes.

  • Scaling: Weak scaling keeps work per processor constant while increasing problem size and processing elements.It is used for memory-bound problems that do not fit in one node; its efficiency is We = T1/TNPW.
  • Scaling: Strong scaling fixes problem size while increasing processing elements, thereby reducing workload per processor.These measurements target compute-bound problems.
  • Algorithm: The parallel algorithm initializes MPI, assigns one processor per subdomain, prepares points, communicates with neighbors, and minimizes each subdomain loss.The algorithm includes preprocessing, computation, communication, and local optimization operations.
  • Cost profile: Figure 4 profiles data-loss, residual-loss, and backward-pass computation times for Burgers-equation PINNs.The profiling varies residual-point count, hidden-layer count, and neurons per layer.

7.3. A pedagogical example on the computational complexity of PINNs

The Burgers-equation example profiles how PINN computational cost changes with network and sampling choices. Residual-loss evaluation is identified as the dominant computational expense because it requires PDE derivatives.

  • The example uses the one-dimensional time-dependent Burgers equation with specified initial and boundary conditions.The setup uses hyperbolic-tangent activation and learning rate 1e-4.
  • Computational profiling varies network depth, width, and the number of residual points to assess their effects on performance.
  • Residual-loss evaluation consumes the most time among data loss, residual loss, and backward-pass computations.Residual loss requires computing partial derivatives in the governing PDEs.

7.4. Two-dimensional steady-state incompressible Navier–Stokes equations

The Navier–Stokes experiments assess solution accuracy and distributed runtime for cPINN and XPINN. cPINN generally reduces communication and achieves higher efficiency for spatial decomposition, while XPINN supports broader decomposition choices.

  • Accuracy of Algorithm 1: For Re = 100 lid-driven cavity flow, cPINN and XPINN use four subdomains and are compared with reference velocity results from Ghia et al.Each subdomain has 12000 residual points, 80 boundary points, and 250 interface points.
  • Computation versus Communication times: Figure 6 compares CPU computation and communication times for cPINN and XPINN with 100 residual points and 20 interface points.
  • Computation versus Communication times: XPINN communication time exceeds cPINN for spatial decomposition because XPINN computes residuals along interfaces before exchanging neighboring-domain values.cPINN can reuse flux-related quantities, reducing computational-graph passes and communication.
  • Weak scaling and Strong scaling: Both distributed methods show good weak scaling on CPUs and GPUs, but XPINN has greater communication time for spatial decomposition.The stated advantage of XPINN appears when spatial and temporal decompositions are combined.
  • Weak scaling and Strong scaling: cPINN reaches average strong-scaling efficiency of 80%, whereas XPINN reaches 70%.The paper attributes the difference to greater communication in XPINN and addresses improving XPINN efficiency later.
  • Comparison between the runtime obtained from data-parallel vanilla PINN, cPINN and XPINN: Runtime comparisons include data-parallel vanilla PINN, cPINN, and XPINN under uniform eight-layer, 64-neuron network architectures.The comparison reports walltimes averaged over 10 runs for the first 200 iterations.

7.5. Viscous Burgers equation with space-time domain decomposition

The viscous Burgers experiment tests XPINN with decomposition across both space and time. It demonstrates the method’s efficacy for time-direction decomposition and its suitability for transient problems.

  • Problem setup: The experiment uses the one-dimensional viscous Burgers equation with specified initial and boundary conditions and ν = 0.01/π.
  • Space-time decomposition: XPINN decomposes the Burgers domain in space-time and distributes residual and interface points across the resulting subdomains.The whole domain contains 80k residual points, with 20 interface points along both time and space directions.
  • Space-time decomposition: Each subdomain uses a five-hidden-layer neural network with 20 neurons per layer, hyperbolic-tangent activation, and learning rate 8e-4.
  • Results: The example shows XPINN efficacy when decomposition is performed along the time direction.The paper characterizes XPINN as promising for transient problems.

7.6. Inverse problem: Steady-state heat conduction with variable conductivity

The inverse heat-conduction problem infers unknown spatially varying thermal conductivity from limited conductivity-boundary data and interior temperature values on a ten-region US-map domain. XPINN accurately handles the irregular decomposition, while parallel performance depends on precision, hardware, communication, and load balance.

  • Problem: The governing problem uses steady-state heat conduction with temperature and conductivity subject to Dirichlet boundary conditions derived from an exact solution.The variable conductivity is K(x, y) = 20 + exp(0.1y) sin(0.5x).
  • Domain decomposition: The US material domain is divided into 10 subdomains with arbitrarily shaped interfaces, boundary data, and training data available in every subdomain.XPINN is used because it can handle complex, non-convex interfaces.
  • Method: The inverse task infers unknown thermal conductivity from a few boundary conductivity measurements and interior temperature values, using one PINN per subdomain.Each network has 3 hidden layers with 80 neurons per layer and a fixed learning rate of 6e-3.
  • Accuracy: XPINN accurately inferred temperature and thermal conductivity, with corresponding absolute point-wise errors reported for the ten-region problem.The predictions and errors are shown in Figure 12.
  • Parallel performance: 7X and 9X scaling were obtained on 10 GPUs for single- and double-precision arithmetic, respectively.The implementation uses one GPU per subdomain; insufficient residual points increase data-fetching and inter-GPU communication costs.
  • Limitations: The manually selected partition can cause load imbalance, so point-cloud or K-way partitioning could further optimize communication.Subdomain 7 has only 800 residual points and must wait for the other devices.

8. Summary

The paper develops and evaluates a hybrid MPI + X framework for distributed cPINNs and XPINNs. Spatial decomposition gives XPINN higher communication overhead than cPINN, while XPINN supports time decomposition and broader interface geometries.

  • Summary: The framework parallelizes domain-decomposition-based cPINNs and XPINNs using MPI + X, where X is CPUs or GPUs.It evaluates forward and inverse problems with weak and strong scaling, speedup, and efficiency measurements.
  • Summary: cPINN has lower communication overhead for spatial decomposition, whereas XPINN can decompose along time and apply to physical problems beyond conservation laws.The two methods therefore offer distinct practical trade-offs.
  • Summary: XPINN supports arbitrarily shaped complex subdomains and allows neural-network hyperparameters to adapt separately in each subdomain.This flexibility distinguishes it from classical data- and model-parallel approaches.
Loading 2104.10013v3…