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

arXiv: https://arxiv.org/abs/2505.06331 · HTML: https://mlsanity.com/paper/2505.06331/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 22 of its 45 references that are in the map.
- First posted on arXiv 9 May 2025.

## Builds on (most cited first, 22 of 22)

- [Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://mlsanity.com/paper/1502.03167/) (2015)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [DeepSDF: Learning Continuous Signed Distance Functions for Shape Representation](https://mlsanity.com/paper/1901.05103/) (2019)
- [Neural Tangent Kernel: Convergence and Generalization in Neural Networks](https://mlsanity.com/paper/1806.07572/) (2018)
- [Implicit Neural Representations with Periodic Activation Functions](https://mlsanity.com/paper/2006.09661/) (2020)
- [When and why PINNs fail to train: A neural tangent kernel perspective](https://mlsanity.com/paper/2007.14527/) (2020)
- [Weight Normalization: A Simple Reparameterization to Accelerate Training of Deep Neural Networks](https://mlsanity.com/paper/1602.07868/) (2016)
- [Physics-informed neural networks (PINNs) for fluid mechanics: A review](https://mlsanity.com/paper/2105.09506/) (2021)
- [PhyGeoNet: Physics-Informed Geometry-Adaptive Convolutional Neural Networks for Solving Parameterized Steady-State PDEs on Irregular Domain](https://mlsanity.com/paper/2004.13145/) (2020)
- [Activation Functions in Deep Learning: A Comprehensive Survey and Benchmark](https://mlsanity.com/paper/2109.14545/) (2021)
- [Activation Functions: Comparison of trends in Practice and Research for Deep Learning](https://mlsanity.com/paper/1811.03378/) (2018)
- [A comprehensive and FAIR comparison between MLP and KAN representations for differential equations and operator networks](https://mlsanity.com/paper/2406.02917/) (2024)
- [Multi-Objective Loss Balancing for Physics-Informed Deep Learning](https://mlsanity.com/paper/2110.09813/) (2021)
- [Improved architectures and training algorithms for deep operator networks](https://mlsanity.com/paper/2110.01654/) (2021)
- [PirateNets: Physics-informed Deep Learning with Residual Adaptive Networks](https://mlsanity.com/paper/2402.00326/) (2024)
- [Gradient Alignment in Physics-informed Neural Networks: A Second-Order Optimization Perspective](https://mlsanity.com/paper/2502.00604/) (2025)
- [Physics Guided RNNs for Modeling Dynamical Systems: A Case Study in Simulating Lake Temperature Profiles](https://mlsanity.com/paper/1810.13075/) (2018)
- [ConFIG: Towards Conflict-free Training of Physics Informed Neural Networks](https://mlsanity.com/paper/2408.11104/) (2024)
- [Optimizing the Optimizer for Physics-Informed Neural Networks and Kolmogorov-Arnold Networks](https://mlsanity.com/paper/2501.16371/) (2025)
- [Learning Specialized Activation Functions for Physics-informed Neural Networks](https://mlsanity.com/paper/2308.04073/) (2023)
- [Robust Weight Initialization for Tanh Neural Networks with Fixed Point Analysis](https://mlsanity.com/paper/2410.02242/) (2024)
- [Element-wise Multiplication Based Deeper Physics-Informed Neural Networks](https://mlsanity.com/paper/2406.04170/) (2024)

## Abstract (arXiv)

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.