Mask-PINNs: Mitigating internal covariate shift in physics-informed neural networks

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. Conventional remedies for ICS-such as Batch Normalization and Layer Normalization-aim to stabilize feature distributions through statistical regularization. However, PINNs require deterministic coordinate-to-solution mappings for enforcing physical constraints, making such strategies fundamentally misaligned with their formulation. To address this issue, we propose Mask-PINNs, which introduce a smooth, learnable mask to adaptively regulate internal features without altering the pointwise physics-based formulation. 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 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. The codes of the experiments can be found on https://github.com/flongjiang/Mask-PINNs.