Accelerating Natural Gradient Descent for PINNs with Randomized Nyström Preconditioning

TL;DR

This work tackles the optimization bottleneck in neural-network PDE solvers by addressing the ill-conditioning of the NGD Gramian $\mathbf{G}(\boldsymbol{\theta})$. It introduces NyströmNGD, which combines matrix-free NGD with randomized Nyström preconditioning to accelerate the inner Conjugate Gradient solves for $(\mathbf{G}(\boldsymbol{\theta})+\mu\mathbf{I})\mathbf{d}=\nabla L(\boldsymbol{\theta})$; this yields substantial improvements in both iteration count and wall-clock time across PINNs and FEINNs on diverse PDEs. The method leverages fast Gramian matvecs via automatic differentiation and exploits strong Gramian spectral decay to obtain an effective low-rank preconditioner, with adaptive updates of $\mu$ and the Nyström rank $\ell$. Empirically, NyströmNGD achieves higher final accuracy and significantly faster convergence than NGD variants and BFGS across Poisson, heat, and Navier–Stokes-like problems, while providing explicit memory control through $\ell_{\max}$. These results indicate NyströmNGD's practical impact for scalable, mesh-free PDE solvers with variational formulations.

Abstract

Natural Gradient Descent (NGD) has emerged as a promising optimization algorithm for training neural network-based solvers for partial differential equations (PDEs), such as Physics-Informed Neural Networks (PINNs). However, its practical use is often limited by the high computational cost of solving linear systems involving the Gramian matrix. While matrix-free NGD methods based on the conjugate gradient (CG) method avoid explicit matrix inversion, the ill-conditioning of the Gramian significantly slows the convergence of CG. In this work, we extend matrix-free NGD to broader classes of problems than previously considered and propose the use of Randomized Nyström preconditioning to accelerate convergence of the inner CG solver. The resulting algorithm demonstrates substantial performance improvements over existing NGD-based methods and other state-of-the-art optimizers on a range of PDE problems discretized using neural networks.

Figures (11)

• Figure 1: Spectral decay of the Gramian matrix when training a PINN for the 3D Poisson problem with exact solution $u(x,y,z) = \sin(\pi x)\sin(\pi y)\sin(\pi z)$, using Energy Natural Gradient Descent muller_achieving_2023. Left: spectral decay of the Gramian $\mathbf{G}(\boldsymbol{\theta})$ at different iterations, with each color corresponding to a different iteration. Right: $\mathrm{H}^1(\Omega)$ relative error during training.
• Figure 2: Convergence for the 3D Poisson equation discretized via PINNs.
• Figure 3: 3D Poisson equation discretized via PINNs and trained with NGD using an SVD-based pseudoinverse. Left: decay of the normalized singular values $\{\sigma_i/\sigma_1\}_{i=1}^p$ of the Gramian across iterations, with each color representing a different iteration. Right: relative $\mathrm{H}^1(\Omega)$ error vs. iterations.
• Figure 4: Convergence for the Heat equation in 3+1D discretized via PINNs.
• Figure 5: Heat equation in 3+1D discretized via PINNs and trained with NGD using the SVD-based pseudoinverse. Left: decay of the normalized singular values $\{\sigma_i / \sigma_1\}_{i=1}^p$ of the Gramian across iterations, with each color representing a different iteration. Right: relative $\mathrm{H}^1(\Omega)$ error vs. iterations.

Theorems & Definitions (12)

• Theorem 3.0
• proof
• Remark 3.1
• Remark 3.2
• Proposition 4.1
• Remark 4.2
• Proposition 4.3
• Remark 4.4
• Remark 4.5
• Remark 5.1