Augmented Physics-Informed Neural Networks (APINNs): A gating network-based soft domain decomposition methodology

Physics-Informed Neural Networks (PINNs) and extended PINNs (XPINNs) have emerged as a promising approach in computational science and engineering for solving partial differential equations (PDEs) by combining the power of artificial intelligence (AI) with the underlying physics to accurately model and predict the solutions to complex problems in science and engineering. In this work, we propose the augmented physics-informed neural network (APINN), which adopts soft and trainable domain decomposition and flexible parameter sharing to further improve the extended PINN (XPINN) as well as the vanilla PINN methods. Concretely, a trainable gate network is employed to mimic the hard decomposition of XPINN, which can be flexibly fine-tuned for discovering a potentially better partition. The gate network satisfying the partition-of unity property, weighted averages several sub-networks as the final output. APINN does not require complex interface conditions, whose sub-nets can utilize all training samples rather than just part of the training data in their subdomains. Lastly, each sub-net shares part of the common parameters to capture the similar components in each decomposed function. Furthermore, following the PINN generalization theory (Hu et al., 2022), APINN is shown to improve generalization by proper gate network initialization and general domain & function decomposition. Extensive experiments on different partial differential equations (PDEs) demonstrate how APINN improves PINN and XPINN. Specifically, we present examples where XPINN performs similarly to or worse than PINN, so that APINN can significantly improve both. We also show cases where XPINN is already better than PINN, so APINN can still slightly improve XPINN. Furthermore, we visualize the optimized gating networks and their optimization trajectories, and connect them with their performance, which helps discover the possibly optimal decomposition. Interestingly, if initialized by different decomposition, the performances of corresponding APINNs can differ drastically. This, in turn, shows the potential to design an optimal domain decomposition for the PDE under consideration.