PINN Balls: Scaling Second-Order Methods for PINNs with Domain Decomposition and Adaptive Sampling

TL;DR

PINN Balls addresses the scalability bottleneck of second-order training for physics-informed neural networks by marrying a locally sparse mixture-of-experts with a fully learned domain decomposition guided by adversarial adaptive sampling. The method trains submodels with a Levenberg–Marquardt update while aligning the decomposition to focus training where the PDE residuals are hardest, via a minimax sampling process that can converge to uniform coverage as residuals vanish. Theoretical contributions extend AAS convergence to practical Dirichlet/entropy regularizations, and empirical results across Helmholtz, Burgers, and Navier–Stokes demonstrate state-of-the-art accuracy under memory constraints and favorable scaling with the number of experts. While the approach offers strong memory efficiency and accuracy, it relies on efficient sparse solvers and careful DD initialization, indicating directions for solver development and DD optimization. Overall, the work advances SciML by enabling efficient, accurate second-order PINN training on large domains, with clear implications for industrial PDE solving.

Abstract

Recent advances in Scientific Machine Learning have shown that second-order methods can enhance the training of Physics-Informed Neural Networks (PINNs), making them a suitable alternative to traditional numerical methods for Partial Differential Equations (PDEs). However, second-order methods induce large memory requirements, making them scale poorly with the model size. In this paper, we define a local Mixture of Experts (MoE) combining the parameter-efficiency of ensemble models and sparse coding to enable the use of second-order training. Our model -- \textsc{PINN Balls} -- also features a fully learnable domain decomposition structure, achieved through the use of Adversarial Adaptive Sampling (AAS), which adapts the DD to the PDE and its domain. \textsc{PINN Balls} achieves better accuracy than the state-of-the-art in scientific machine learning, while maintaining invaluable scalability properties and drawing from a sound theoretical background.

Figures (6)

• Figure 1: High-level diagram of the minimax training routine devised for the PINN Balls model.
• Figure 2: Mean value with shaded variance, for different ensemble sizes, of various statistics of the LM update step on PINN Balls. From left to right, the runtime in seconds, the percentage of non-zero elements of $D_\theta L$, and the CPU memory consumed to invert and store the matrix $D_\theta L^T D_\theta L$
• Figure 3: Median Loss value and minimax range over 10 independent trainings of PINN Balls with an increasing number of submodels on Helmholtz (left), Burgers (center) and Navier Stokes' (right) equations. Each submodel has 3 hidden layers and 10 neurons per layer.
• Figure 4: Prediction with learnt center and radii of the PINN Balls (left); correct solution (center); and $L^2$ Error achieved by our best performing model on Helmholtz equation.
• Figure 5: Prediction with learnt center and radii of the PINN Balls (left); correct solution (center); and $L^2$ Error achieved by our best performing model on Burgers' equation.

Theorems & Definitions (14)

• Theorem 3.1
• proof : Idea of the proof
• Remark 3.2
• Remark 3.3
• Definition 4.1: PINN Balls
• Remark 4.2
• Remark A.1
• Theorem A.2
• Remark A.3
• proof : Proof of Theorem \ref{['thm:aas-grad']}