On the Global Convergence of Stochastic Adam for Physics-Informed Neural Networks

원문 초록

The Adam optimizer is widely used to train neural networks with highly nonlinear and nonconvex objective functions. Remarkably, Adam often drives the empirical training loss nearly to zero. This phenomenon has been partially explained by neural tangent kernel (NTK) analyses of gradient descent (GD) and stochastic gradient descent (SGD). However, analyzing Adam presents additional challenges. Its first-moment momentum couples stochastic gradients across iterations, while its adaptive second-moment preconditioner depends on the same stochastic gradient used in the update. In this paper, we establish linear convergence guarantees for Adam in training sufficiently wide two-layer neural networks with either full-batch or minibatch gradients. For full-batch Adam, we construct a Lyapunov function to control the effect of historical momentum. For minibatch Adam, our analysis builds on an exact virtual-iterate representation, a conditional analysis of the dependence between stochastic gradients and adaptive preconditioners, and uniform concentration bounds for the second-moment estimates. We prove geometric decay of the empirical loss for full-batch Adam and geometric decay in restricted expectation on a high-probability event for minibatch Adam. Beyond regression, we extend our convergence analysis to physics-informed neural networks (PINNs) for solving the Poisson equation, where differential operators and separate sampling of interior and boundary points introduce additional challenges. We establish analogous global convergence guarantees for both deterministic and stochastic Adam under explicit conditions on network width and Adam's parameters, characterizing the effects of minibatch size on convergence. These results help bridge the gap between Adam's empirical success and its theoretical understanding, providing sharper convergence guarantees for PINNs than existing analyses.