🔢 mathematics

Accelerating Natural Gradient Descent for PINNs with Randomized Numerical Linear Algebra

This paper proposes a novel Natural Gradient Descent algorithm for Physics-Informed Neural Networks that leverages Randomized Numerical Linear Algebra techniques to precondition the inner Conjugate Gradient solver, thereby overcoming the ill-conditioning and high computational costs that typically hinder the practical application of NGD for solving partial differential equations.

Original authors: Ivan Bioli, Carlo Marcati, Giancarlo Sangalli

Published 2026-05-28

📖 4 min read🧠 Deep dive

CC BY 4.0

Original authors: Ivan Bioli, Carlo Marcati, Giancarlo Sangalli

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a very smart, but slightly confused, robot (a Neural Network) to solve a complex physics puzzle, like predicting how heat spreads through a metal plate or how water flows around a bridge. This is what scientists call training a Physics-Informed Neural Network (PINN).

The robot learns by making guesses, checking how wrong it is, and then adjusting its internal settings (parameters) to get closer to the right answer. The problem is that the "landscape" of possible answers is like a giant, bumpy mountain range with deep, narrow valleys. Standard learning methods (like the popular "Adam" optimizer) are like hikers who just look at the slope directly under their feet. They often get stuck in shallow dips or wander aimlessly, taking a very long time to find the true bottom of the valley.

The "Natural" Way to Walk
To fix this, researchers use a smarter strategy called Natural Gradient Descent (NGD). Instead of just looking at the slope, NGD looks at the shape of the entire valley. It knows exactly which direction leads straight to the bottom, ignoring the confusing bumps and twists. It's like having a GPS that knows the terrain perfectly.

However, there's a catch: Calculating this perfect direction is incredibly expensive. It's like trying to draw a detailed 3D map of the entire mountain range every single time the hiker takes a step. For a large robot, this map is so huge it would crash the computer's memory.

The "Matrix-Free" Shortcut
To solve the memory problem, scientists developed a "matrix-free" version. Instead of drawing the whole map, they just ask the computer, "If I take a step in this direction, what happens?" This avoids storing the giant map. But, because the mountain is so bumpy (mathematically "ill-conditioned"), the computer still takes tiny, slow steps, and the hiker gets tired before reaching the bottom.

The Paper's Big Idea: Randomized Shortcuts
This paper introduces a clever trick using Randomized Numerical Linear Algebra (RandNLA). Think of the mountain range as having a secret: most of the bumps are actually very small and unimportant. The mountain is mostly flat in most directions, with only a few steep slopes that really matter.

The authors propose using randomized sampling to quickly figure out which few directions are the important ones. They use two specific tools:

Nyström Approximation: Like taking a few random photos of the mountain to guess its overall shape.
RPCholesky: Like a smart explorer who picks the most interesting spots to check based on what they've already seen.

These tools build a preconditioner. In our analogy, a preconditioner is like giving the hiker a pair of special boots or a magic compass. These boots smooth out the tiny, annoying bumps and make the steep slopes easier to climb. Suddenly, the hiker can take giant, confident strides straight toward the solution.

The Results
The authors created two new algorithms based on this idea: NyströmNGD and RPCholNGD. They tested these on various physics puzzles (like heat flow and fluid dynamics) and found:

Speed: They reached the solution much faster than the standard "smart" methods, often in a fraction of the time.
Accuracy: They found more precise answers than the popular "Adam" optimizer and matched or beat other advanced methods.
Efficiency: They achieved this without needing a supercomputer, using much less memory than other high-end methods.

In Summary
The paper doesn't just say "let's walk faster." It says, "Let's stop trying to map the whole mountain and instead use a smart, random sampling trick to build a pair of magic boots that let us walk straight to the bottom." This makes training these physics-solving robots significantly faster, cheaper, and more reliable.

Technical Summary: Accelerating Natural Gradient Descent for PINNs with Randomized Numerical Linear Algebra

Problem Statement
Physics-Informed Neural Networks (PINNs) and related neural network-based solvers for Partial Differential Equations (PDEs) have shown promise in providing mesh-free approximations, particularly in high dimensions. However, their practical adoption is hindered by optimization challenges. Standard first-order optimizers like Adam often stagnate at low-accuracy solutions due to the ill-conditioning of the loss landscape. While Natural Gradient Descent (NGD) has emerged as a robust second-order alternative that accounts for the geometry of the function space, its application to large-scale neural networks is computationally prohibitive.

The primary bottleneck in NGD is the need to solve linear systems involving the Gramian matrix (or Fisher information matrix), $G(\theta)$ . Direct inversion or factorization of $G(\theta)$ incurs a cubic cost $O(p^3)$ with respect to the number of parameters $p$ , rendering it infeasible for large networks. Matrix-free approaches using the Conjugate Gradient (CG) method avoid explicit matrix assembly but suffer from slow convergence when $G(\theta)$ is ill-conditioned, a common occurrence in PDE discretizations.

Methodology
The authors propose a modular framework that accelerates matrix-free NGD by integrating techniques from Randomized Numerical Linear Algebra (RandNLA) to construct effective preconditioners for the inner CG solver.

Generalized Matrix-Free NGD: The work extends matrix-free NGD to a broader class of metrics relevant to PDEs. By leveraging automatic differentiation and a specific assumption on the metric structure (integral-based forms), the authors demonstrate that matrix-vector products (matvecs) with the Gramian can be computed efficiently without assembling the full matrix. This applies to both strong formulations (PINNs) and weak formulations (Finite Element Interpolated Neural Networks, or FEINNs).
RandNLA Preconditioning: To address the ill-conditioning of $G(\theta)$ , the paper introduces two novel preconditioning strategies based on the observation that the Gramian exhibits strong spectral decay (rapid decay of singular values):
- NyströmNGD: Utilizes the randomized Nyström approximation to construct a low-rank preconditioner. This method approximates the Gramian using a random test matrix and computes a preconditioner that leverages the dominant eigenspace.
- RPCholNGD: Employs the Randomly Pivoted Partial Cholesky (RPCholesky) decomposition. This adaptive method selects pivots based on the residual diagonal, constructing a low-rank approximation $FF^\top$ that serves as the basis for the preconditioner.
Algorithmic Framework: The authors present a meta-algorithm that decouples the NGD implementation from the preconditioner construction. This allows for flexible combinations of direct solvers, matrix-free iterative solvers, and various preconditioning strategies. The framework includes adaptive heuristics for selecting the regularization parameter $\mu$ and the rank $\ell$ of the preconditioner, ensuring numerical stability and balancing computational cost with preconditioner quality.

Key Contributions

Generalization of Matrix-Free NGD: The paper provides a methodology to compute Gramian matvecs via automatic differentiation for a wide range of PDE-related metrics, generalizing beyond previous specific implementations.
Novel Preconditioned Algorithms: The introduction of NyströmNGD and RPCholNGD, which combine matrix-free NGD with randomized Nyström and RPCholesky preconditioning, respectively.
Scalable Framework: A modular architecture that reduces the memory footprint from $O(p^2)$ (required by full-memory quasi-Newton methods) to $O(p\ell)$ , where $\ell$ is a tunable low-rank parameter.
Efficiency: The proposed methods transform the computational challenge of ill-conditioning into an asset by exploiting the low-rank structure of the Gramian, enabling efficient preconditioning without explicit matrix assembly.

Results
Extensive numerical experiments were conducted on five PDE problems, including the 3D Poisson equation, the 3+1D Heat equation, Kovasznay flow, and 2D diffusion-reaction-transport problems, using both PINNs and FEINNs.

Performance vs. Baselines: The preconditioned matrix-free NGD variants (NyströmNGD and RPCholNGD) achieved final accuracies comparable to NGD with a direct solver and the state-of-the-art SSBroyden optimizer.
Computational Efficiency: The proposed methods significantly outperformed unpreconditioned matrix-free NGD, which failed to reach high accuracy due to poor CG convergence. Compared to SSBroyden and other optimizers (Adam, BFGS, L-BFGS), NyströmNGD and RPCholNGD demonstrated superior or comparable accuracy with significantly reduced wall-clock time on most benchmarks.
Memory and Scalability: The methods maintained a memory footprint of $O(p\ell)$ , allowing them to scale to larger networks where full-matrix methods or full-memory quasi-Newton methods would be infeasible.
Rank Adaptivity: The experiments showed that the effective rank required for preconditioning ( $\ell$ ) was often much smaller than the total number of parameters, confirming the strong spectral decay of the Gramian in these settings.

Significance and Claims
The paper claims that by bridging Natural Gradient Descent with Randomized Numerical Linear Algebra, it is possible to overcome the computational bottlenecks that have limited the practical use of NGD in PDE solvers. The authors assert that their preconditioned matrix-free approach offers a superior trade-off between accuracy and computational time compared to existing state-of-the-art optimizers for PINNs.

Specifically, the work highlights that:

Preconditioning is critical for matrix-free NGD; without it, the method is ineffective.
The low-rank structure of the Gramian, often observed empirically, can be systematically exploited via RandNLA to create efficient preconditioners.
The proposed framework is particularly well-suited for neural network-based solvers targeting weak formulations of PDEs (like FEINNs), where the variational structure naturally guides metric selection.
While SSBroyden remains competitive (and slightly faster in specific cases like Kovasznay flow), the proposed methods generally achieve comparable or better accuracy in less time and with lower memory requirements, making them a promising direction for large-scale PDE-constrained optimization.

Technical Summary: Accelerating Natural Gradient Descent for PINNs with Randomized Numerical Linear Algebra

More like this