Mask-PINNs: Mitigating Internal Covariate Shift in Physics-Informed Neural Networks

TL;DR

Mask-PINNs introduce a learnable masking mechanism to mitigate internal covariate shift in physics-informed neural networks, enhancing training stability, accuracy, and robustness without compromising physical constraints.

Contribution

The paper proposes Mask-PINNs, a novel architecture with a learnable mask that controls feature distributions while preserving physical laws, addressing ICS in PINNs.

Findings

Improved prediction accuracy on PDE benchmarks.

Enhanced convergence stability and robustness.

Enabled effective training of wider PINN networks.

Abstract

Physics-Informed Neural Networks (PINNs) have emerged as a powerful framework for solving partial differential equations (PDEs) by embedding physical laws directly into the loss function. However, as a fundamental optimization issue, internal covariate shift (ICS) hinders the stable and effective training of PINNs by disrupting feature distributions and limiting model expressiveness. Unlike standard deep learning tasks, conventional remedies for ICS -- such as Batch Normalization and Layer Normalization -- are not directly applicable to PINNs, as they distort the physical consistency required for reliable PDE solutions. To address this issue, we propose Mask-PINNs, a novel architecture that introduces a learnable mask function to regulate feature distributions while preserving the underlying physical constraints of PINNs. We provide a theoretical analysis showing that the mask suppresses the expansion of feature representations through a carefully designed modulation mechanism. Empirically, we validate the method on multiple PDE benchmarks -- including convection, wave propagation, and Helmholtz equations -- across diverse activation functions. Our results show consistent improvements in prediction accuracy, convergence stability, and robustness. Furthermore, we demonstrate that Mask-PINNs enable the effective use of wider networks, overcoming a key limitation in existing PINN frameworks.