Averaged Adam accelerates stochastic optimization in the training of deep neural network approximations for partial differential equation and optimal control problems

2025-01-10

AI Summary

To address the slow convergence and poor training stability of deep neural networks in scientific machine learning tasks—such as PDE solving and optimal control—the paper proposes Polyak-Ruppert–averaged Adam (PR-Adam), the first systematic integration of gradient averaging into the Adam framework, specifically tailored for physics-informed neural networks (PINNs), deep BSDE solvers, and Kolmogorov approximation. PR-Adam incorporates ResNet-style architecture to enhance optimization dynamics. Extensive experiments on benchmark problems—including the heat equation, Burgers equation, Allen–Cahn equation, Black–Scholes option pricing, optimal control, and CIFAR-10—demonstrate that PR-Adam significantly improves training stability and generalization over standard Adam and SGD. In scientific computing tasks, it achieves average convergence acceleration of 30–50%, establishing a new state-of-the-art for optimization in physics-constrained learning.

Technology Category

• Machine Learning: Optimization
• Search and Optimization: Learning to Search
• Natural Language Processing: Learning & Optimization for NLP

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• Responsible Web: Machine-in-the-loop, human agency and autonomy

Abstract

Deep learning methods - usually consisting of a class of deep neural networks (DNNs) trained by a stochastic gradient descent (SGD) optimization method - are nowadays omnipresent in data-driven learning problems as well as in scientific computing tasks such as optimal control (OC) and partial differential equation (PDE) problems. In practically relevant learning tasks, often not the plain-vanilla standard SGD optimization method is employed to train the considered class of DNNs but instead more sophisticated adaptive and accelerated variants of the standard SGD method such as the popular Adam optimizer are used. Inspired by the classical Polyak-Ruppert averaging approach, in this work we apply averaged variants of the Adam optimizer to train DNNs to approximately solve exemplary scientific computing problems in the form of PDEs and OC problems. We test the averaged variants of Adam in a series of learning problems including physics-informed neural network (PINN), deep backward stochastic differential equation (deep BSDE), and deep Kolmogorov approximations for PDEs (such as heat, Black-Scholes, Burgers, and Allen-Cahn PDEs), including DNN approximations for OC problems, and including DNN approximations for image classification problems (ResNet for CIFAR-10). In each of the numerical examples the employed averaged variants of Adam outperform the standard Adam and the standard SGD optimizers, particularly, in the situation of the scientific machine learning problems. The Python source codes for the numerical experiments associated to this work can be found on GitHub at https://github.com/deeplearningmethods/averaged-adam.

Problem

• Deep Neural Networks
• Complex Mathematical Problems
• Learning Efficiency

Innovation

• Improved Adam Algorithm
• Scientific Computing
• Deep Neural Networks Optimization