Improving Energy Natural Gradient Descent through Woodbury, Momentum, and Randomization

TL;DR

This paper enhances energy natural gradient descent (ENGD) for Physics-Informed Neural Networks by integrating Woodbury, momentum, and randomization techniques, significantly boosting efficiency and maintaining accuracy.

Contribution

The paper introduces novel methods combining Woodbury formula, subsampled algorithms, and randomization to improve ENGD's computational efficiency for PINNs.

Findings

Achieved up to 75x faster training while maintaining the same $L^2$ error.

Demonstrated effectiveness of randomization in early training stages for low-dimensional problems.

Identified barriers to acceleration in high-dimensional or complex scenarios.

Abstract

Natural gradient methods significantly accelerate the training of Physics-Informed Neural Networks (PINNs), but are often prohibitively costly. We introduce a suite of techniques to improve the accuracy and efficiency of energy natural gradient descent (ENGD) for PINNs. First, we leverage the Woodbury formula to dramatically reduce the computational complexity of ENGD. Second, we adapt the Subsampled Projected-Increment Natural Gradient Descent algorithm from the variational Monte Carlo literature to accelerate the convergence. Third, we explore the use of randomized algorithms to further reduce the computational cost in the case of large batch sizes. We find that randomization accelerates progress in the early stages of training for low-dimensional problems, and we identify key barriers to attaining acceleration in other scenarios. Our numerical experiments demonstrate that our methods outperform previous approaches, achieving the same $L^2$ error as the original ENGD up to $75\times$ faster.