Task-Contribution Dynamics of Adam in Physics-Informed Neural Networks: Slow-Manifold Analysis and Gradient Modulation

원문 초록

Physics-informed neural networks (PINNs) often suffer from imbalanced optimization of competing loss components, particularly when approximating solutions with high-frequency or multiscale features. Although existing balancing strategies have proven effective in alleviating this difficulty, the mechanisms governing the evolution of individual task contributions to Adam's parameter updates remain insufficiently understood. We introduce task-wise indicators to characterize this evolution and identify three distinct training regimes: smooth variation, intermittent spikes, and persistent oscillations. Focusing on the early smooth regime, we combine discrete-time analysis with a continuous-time singular perturbation framework under the full-batch setting and the parameter regime $(1-\beta_2)/(1-\beta_1)\ll1$. A first-order approximation of the locally attracting slow manifold reveals two factors governing the step-to-step redistribution of task contributions: curvature-related alignment disparities and differences in the coordinatewise distributions of preconditioned task gradients. Both factors involve Adam's shared second-moment estimate. Inspired by AdaTask, we propose Gradient-Modulated Adam (GMAdam), which modulates incoming task gradients using task-specific second-moment estimates. Numerical experiments on high-frequency Poisson and wave equations, the Allen-Cahn equation, and two-dimensional lid-driven cavity flow ($Re=100,400$) demonstrate consistent improvements in predictive accuracy over AdaTask at comparable computational cost.