Accelerating Natural Gradient Descent for PINNs with Randomized Numerical Linear Algebra

Ivan Bioli; Carlo Marcati; Giancarlo Sangalli

math.NAMay 16, 2025

Accelerating Natural Gradient Descent for PINNs with Randomized Numerical Linear Algebra

Authors: Ivan Bioli, Carlo Marcati, Giancarlo Sangalli

Organizations: Department of Mathematics, University of Pavia, Via A. Ferrata 5, 27100 Pavia, Italy · Department of Civil Engineering and Architecture, University of Pavia, Via A. Ferrata 3, 27100 Pavia, Italy · Istituto di Matematica Applicata e Tecnologie Informatiche “E. Magenes”, CNR, Via A. Ferrata 1, 27100 Pavia, Italy

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 Numerical Linear Algebra (RandNLA) techniques for efficient preconditioning of the inner CG solver. The resulting algorithms demonstrate substantial performance improvements over existing NGD-based methods on a range of PDE problems discretized using neural networks, and offer competitive results compared to other state-of-the-art optimizers.

Figures & tables

Strategy	Setup Cost	MMP Cost	Memory Footprint
Precompute Gramian	$\mathcal{O}(q\texttt{[FP]}+p^{2}q)$	$\mathcal{O}(p^{2}\ell)$	$\mathcal{O}(p^{2}+pB)$
Precompute Jacobian only	$\mathcal{O}(q\texttt{[FP]})$	$\mathcal{O}(pq\ell)$	$\mathcal{O}(pq)$
Recompute batch Jacobian
on the fly	–	$\mathcal{O}(q\texttt{[FP]}+pq\ell)$	$\mathcal{O}(pB)$

Table 1 . Summary of computational and memory costs for the three matrix-vector product strategies. Here

p

is the number of parameters, [FP] is the cost of a forward pass,

\ell

is the number of simultaneous matvecs,

q

is the number of quadrature points, and

B

is the batch size for processing quadrature points. We assume that

p\leq q

, i.e., we focus on the regime where a highly accurate quadrature rule is employed and the number of quadrature points

q

exceeds the number of neural network parameters

p

Optimizer	Poisson 3D	Heat 3+1D	Kovasznay	Deep-Ritz Poisson	FEINNs 2D
NGD full	$1.56\times 10^{-6}$
$(1.37-1.77\,\times 10^{-6})$
317.7 s	$1.75\times 10^{-6}$
$(1.49-1.86\,\times 10^{-6})$
279.2 s	$1.08\times 10^{-6}$
$(0.965-1.25\,\times 10^{-6})$

Table 2 . Relative

\mathrm{H}^{1}(\Omega)

error at the plateau iterate for each optimizer and test problem. The first line reports the median error, the second line reports in smaller font and in parentheses the

25

th and

75

th percentiles in compact form as

(a-b\,\times 10^{k})

, and the third line reports the corresponding cumulative wall-clock time in seconds. The plateau criterion and the dot (

\bullet

) marking the corresponding iterate in the convergence curves are described in Section 5.2 ; see Figure 10 .

Optimizer	$\mathbf{3\times 64\times 64\times 64\times 1}$
$\mathbf{q=10^{4}}$	$\mathbf{3\times 64\times 64\times 64\times 1}$
$\mathbf{q=10^{5}}$	$\mathbf{3\times 128\times 128\times 128\times 1}$
$\mathbf{q=10^{5}}$
NGD full	$1.56\times 10^{-6}$
$(1.37-1.77\,\times 10^{-6})$
317.7 s

Table 3 . Scalability experiments on the 3D Poisson problem discretized via PINNs; see Section 5.3 . Relative

\mathrm{H}^{1}(\Omega)

error at the plateau iterate for each optimizer and architecture. The first line reports the median error, the second line reports in smaller font and in parentheses the

25

th and

75

th percentiles in compact form as

(a-b\,\times 10^{k})

, and the third line reports the corresponding cumulative wall-clock time in seconds, while the fourth line reports the per-iteration cost in seconds. Each column corresponds to a different network architecture and number of quadrature points

q

, as indicated in the header.

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Experiment	NGD (all)	(L-)BFGS	SSBroyden	Adam	Network Architecture
Poisson 3D	1 000	50 000	50 000	150 000	[3, 64, 64, 64, 1]
Heat 3+1D	1 000	50 000	50 000	150 000	[4, 64, 64, 64, 1]
Kovasznay	2 000	25 000	25 000	75 000	$u$ : [2, 45, 45, 45, 2]
$p$ : [2, 45, 45, 45, 1]
Deep-Ritz Poisson	1 000	25 000	25 000	75 000	[2, 64, 64, 64, 1]
FEINNs 2D	1 000	25 000	25 000	75 000	[2, 64, 64, 64, 1]

Appendix

Table 4 . Iteration budgets for the evaluated optimizers and neural network architectures for each PDE experiment. Network shapes are denoted as [input_dim, hidden_1, ..., hidden_N, output_dim] .

Experiment	Interior / Mesh configuration	Boundary Points	Evaluation set
Poisson 3D	10 000 random	1 000 random	100 000 random
Heat 3+1D	10 000 random	1 000 random	100 000 random
Kovasznay	10 000 random	1 000 random	100 000 random
Deep-Ritz Poisson	Quad mesh, $h=1/30$ , $d=4$ (14.5k DoF)	-	Gauss–Jacobi, order $16$
FEINNs 2D	Quad mesh, $h=1/30$ , $d=4$ (14.5k DoF)	-	Gauss–Jacobi, order $16$

Appendix

Table 5 . Domain discretization, boundary conditions, and evaluation setups for each numerical experiment.

$\boldsymbol{\gamma}$	Poisson 3D	Heat 3+1D	Deep-Ritz Poisson	FEINNs 2D
$\mathbf{0.1}$	$1.87\times 10^{-6}$
$(0.886-3.27\,\times 10^{-6})$
$1.004\times$	$7.42\times 10^{-6}$
$(6.05-8.60\,\times 10^{-6})$
$1.023\times$	$7.49\times 10^{-1}$
$(3.99-10.4\,\times 10^{-1})$

Appendix

Table 6 . Sensitivity analysis of

\gamma

for NyströmNGD with Gaussian test matrices. The first line reports the median relative

\mathrm{H}^{1}(\Omega)

error, the second line reports in smaller font and in parentheses the

25

th and

75

th percentiles in compact form as

(a-b\,\times 10^{k})

, and the third line reports the wall-clock time relative to the default

\gamma=10

$\boldsymbol{\gamma}$	Poisson 3D	Heat 3+1D	Deep-Ritz Poisson	FEINNs 2D
$\mathbf{0.1}$	$8.42\times 10^{-3}$
$(7.79-8.49\,\times 10^{-3})$
$1.062\times$	$7.48\times 10^{-3}$
$(5.24-9.63\,\times 10^{-3})$
$1.055\times$	$2.12$
$(1.08-2.53)$

Appendix

Table 7 . Sensitivity analysis of

\gamma

for RPCholNGD. The first line reports the median relative

\mathrm{H}^{1}(\Omega)

error, the second line reports in smaller font and in parentheses the

25

th and

75

th percentiles in compact form as

(a-b\,\times 10^{k})

, and the third line reports the wall-clock time relative to the default

\gamma=10

$\boldsymbol{r}$	Poisson 3D	Heat 3+1D	Kovasznay	Deep-Ritz Poisson	FEINNs 2D
$\mathbf{0.1}$	$1.68\times 10^{-6}$
$(1.55-1.75\,\times 10^{-6})$
$1.00\times$	$2.03\times 10^{-6}$
$(2.03-2.12\,\times 10^{-6})$
$1.00\times$	$6.95\times 10^{-7}$
$(6.80-7.48\,\times 10^{-7})$

Appendix

Table 8 . Sensitivity analysis of the parameter

r

for NyströmNGD with Gaussian test matrices. The first line reports the median relative

\mathrm{H}^{1}(\Omega)

error, the second line reports in smaller font and in parentheses the

25

th and

75

th percentiles in compact form as

(a-b\,\times 10^{k})

, and the third line reports the wall-clock time relative to the default

r=10

$\boldsymbol{r}$	Poisson 3D	Heat 3+1D	Kovasznay	Deep-Ritz Poisson	FEINNs 2D
$\mathbf{0.01}$	$1.83\times 10^{-6}$
$(1.78-1.87\,\times 10^{-6})$
$1.00\times$	$2.11\times 10^{-6}$
$(1.81-2.53\,\times 10^{-6})$
$1.00\times$	$9.79\times 10^{-7}$
$(9.14-10.3\,\times 10^{-7})$

Appendix

Table 9 . Sensitivity analysis of the RPCholesky stopping tolerance scaling factor

r

(

\varepsilon=r\cdot p\cdot\epsilon_{\mathrm{mach}}

) for RPCholNGD. The first line reports the median relative

\mathrm{H}^{1}(\Omega)

error, the second line reports in smaller font and in parentheses the

25

th and

75

th percentiles in compact form as

(a-b\,\times 10^{k})

, and the third line reports the wall-clock time relative to the default

r=1

Explore similar work

Energy Manifold Natural Gradient Descent: Riemannian Optimization for Neural PDE Solvers
Jul 24, 2026Zhangyong Liang, Huanhuan GaoNeural Partial Differential Equation Solvers Riemannian Optimization
National Center for Applied Mathematics Tianjin University Tianjin, 300072, China · School of Mechanical and Aerospace Engineering Jilin University Changchun, 130025, China
Global Convergence of DGM and PINN Algorithms for Solving Nonlinear PDEs
Jul 27, 2026Justin Sirignano, Konstantinos Spiliopoulos, Samuel CohenPartial Differential Equations Variational Formulation
Mathematical Institute, University of Oxford · Department of Mathematics and Statistics, Boston University
Mitigating Gradient Pathology in PINNs through Aligned Constraint
May 24, 2026Yichen Luo, Peiyu Zhu, Dongxiao Hu +5Parametric Physics-Informed Neural Network Partial Differential Equations
Department of Information Science and Engineering, KTH Royal Institute of Technology, Stockholm, Sweden · School of Advanced Manufacturing and Robotics, Peking University, Beijing, China · School of Advanced Technology, Xi’an Jiaotong-Liverpool University, Suzhou, China +2

Jul 24, 2026, cs.LGJ/K move · Enter open · S save

Energy Manifold Natural Gradient Descent: Riemannian Optimization for Neural PDE Solvers

Zhangyong Liang, Huanhuan Gao

National Center for Applied Mathematics Tianjin University Tianjin, 300072, China · School of Mechanical and Aerospace Engineering Jilin University Changchun, 130025, China

Neural Partial Differential Equation Solvers Riemannian Optimization Riemannian Manifolds Gradient Descent Manifolds

Open paper AskPDF arXiv

Energy natural gradient descent (ENGD) aligns parameter updates with the curvature of an underlying function-space energy, but existing formulations assume an unconstrained Euclidean parameter domain. We introduce \EMNGDfull{}, a manifold optimization framework for physics-informed and variational neural PDE solvers whose parameters lie on a Riemannian manifold. EMNGD restricts the energy-induced quadratic model to feasible tangent directions and uses retractions to preserve parameter constraints throughout optimization. Under coercivity, we prove that the push-forward of the undamped EMNGD direction is the best feasible approximation to the function-space Newton vector in the energy metric. We establish coordinate invariance, exact reduction to ENGD in Euclidean space, global first-order convergence with Armijo backtracking, and robustness to inexact tangent solves. For quadratic residual energies and generalized Gauss--Newton pullbacks, the Woodbury identity transfers the tangent system to sample space without changing the direction. Nyström approximation provides scalable sample-space solves with controlled direction error and recovers the exact direction after iterative convergence. On the evaluated neural PDE benchmarks, EMNGD achieves higher accuracy and faster convergence than the compared state-of-the-art baselines. Woodbury preserves the EMNGD direction, while scalable-solver diagnostics quantify the accuracy--cost trade-off of preconditioning and residual subsampling.

Abstract

Explore similar work

Energy Manifold Natural Gradient Descent: Riemannian Optimization for Neural PDE Solvers

Global Convergence of DGM and PINN Algorithms for Solving Nonlinear PDEs

Mitigating Gradient Pathology in PINNs through Aligned Constraint