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

TL;DR

This work tackles the high computational cost of energy natural gradient descent (ENGD) for Physics-Informed Neural Networks (PINNs) by introducing three key enhancements. First, ENGD-W uses the Woodbury matrix identity to replace a $P\times P$ inverse with a kernel-based $N\times N$ inversion, achieving $\mathcal{O}(N^2P)$ per-iteration cost. Second, SPRING inserts momentum into the natural-gradient updates, accelerating convergence without requiring line searches. Third, a GPU-efficient Nyström sketch-and-solve approach further reduces costs for large batches in some settings. Empirically, the combination yields up to $75\times$ faster convergence to the same $L^2$ error on challenging high-dimensional PDEs, with SPRING delivering state-of-the-art performance in several cases; randomization helps mainly in early phases for low-dimensional problems and shows barriers in higher dimensions due to effective dimension constraints.

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.

Figures (19)

• Figure 1: Timeline of VMC, PINN and RLA methods. $*$: Stochastic reconfiguration (SR) and energy natural gradient descent (ENGD) both precondition their stochastic gradients with the inverse of an appropriate curvature matrix. $\dagger$: Inspired by MinSR, we apply Woodbury's matrix identity to use the kernel matrix instead of ENGD's Gauss-Newton matrix, thus reducing the cost to $O(N^2P)$ instead of $O(P^3)$, where $N$ denotes the batch size and $P$ the number of parameters; we further introduce SPRING for PINNs to transport curvature information across optimization steps; last, we use a GPU-efficient Nyström approximation to reduce the iteration cost for large batch sizes.
• Figure 2: 5D Poisson: Performance comparison of different optimization algorithms on a five‐dimensional Poisson problem discretized with 10065.0 parameters, trained with a tanh‐activated multilayer perceptron. Introducing the Woodbury matrix identity allows ENGD to take more than 30.0 times more steps and outperform the Hessian-free optimizer by a significant margin.
• Figure 3: Performance of ENGD and SPRING on four problems: 5d and 100d Poisson, 4+1d Heat, and 9+1d log-Fokker-Plancks. SPRING achieves $L^2$ errors not previously seen in the high dimensional settings, the 100d Poisson and 9+1d log-Fokker-Planck problems. In the latter, the error is an order of magnitude lower that the previous state-of-the-art KFAC.
• Figure 4: Effect of Nyström on ENGD for the 5D Poisson equation.
• Figure 5: Effect of Nyström on SPRING for the 100D Poisson equation.