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

TL;DR

PINN Balls introduces a scalable second-order training method for Physics-Informed Neural Networks by combining domain decomposition, adaptive sampling, and a Mixture of Experts approach, improving accuracy and efficiency in solving PDEs.

Contribution

The paper proposes PINN Balls, a novel framework that enables scalable second-order training of PINNs through domain decomposition and adaptive sampling, with a learnable structure and theoretical grounding.

Findings

Achieves better accuracy than state-of-the-art methods.

Maintains scalability with large models.

Utilizes a learnable domain decomposition structure.

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.