← ProjectsResearch · Software

Scalable natural gradient descent for neural PDE solvers

Making natural gradient training of physics-informed neural networks fast, using matrix-free solvers with randomized low-rank preconditioning.

Role: First author
With: C. Marcati, G. Sangalli
Tools: Python, JAX
Links: arXiv

The problem

Neural networks can be trained to solve partial differential equations by minimising the PDE residual — the idea behind Physics-Informed Neural Networks (PINNs). In practice, first-order optimisers such as Adam stall at modest accuracy. Natural Gradient Descent (NGD), which preconditions the gradient with the Gramian of the network’s Jacobian, converges far better, but each step requires solving a linear system

G(\theta_k)\, \delta_k = \nabla L(\theta_k), \qquad G = J^\top J,

whose size equals the number of network parameters. Forming and factorising $G$ is out of the question for realistic networks.

To see the difference, take the Poisson problem with exact solution $u(x,y) = \sin(2\pi x)\sin(2\pi y)$ and look at how each direction moves the function represented by the network: the plain (Euclidean) gradient, pushed forward to function space, is a smooth correction unrelated to the solution, while the natural gradient reproduces the oscillating structure of the target — it moves the function in the direction that actually matters.

Pushforward of the Euclidean gradient: a smooth, almost flat surface — Pushforward of the Euclidean gradient, Dρ(θ)[∇L(θ)].

Pushforward of the natural gradient: an oscillating surface matching the solution error — Pushforward of the natural gradient, Dρ(θ)[G(θ)^†∇L(θ)].

Approach

Matrix-free NGD solves the system with the conjugate gradient (CG) method using only Jacobian–vector products — but $G$ is severely ill-conditioned, so CG converges slowly. We

extend matrix-free NGD to broader classes of problems than previously considered, and
precondition the inner CG solver with Randomized Numerical Linear Algebra (randomized Nyström approximations), which capture the dominant spectrum of $G$ at low cost.

Outcome

The paper will appear in Computer Methods in Applied Mechanics and Engineering. The resulting algorithms show substantial speed-ups over existing NGD-based methods on a range of PDE problems discretised with neural networks, and are competitive with other state-of-the-art optimisers. Follow-up work on low-rank and geometry-aware preconditioning was presented at WCCM-ECCOMAS 2026 and will be presented at SIAM MDS26.