🔬 physics

An Optimisation Framework for the Well-Conditioned Training of Physics-Informed Neural Networks

The paper introduces DSGNAR, a scalable second-order optimization framework that overcomes the ill-conditioned loss landscapes of Physics-Informed Neural Networks (PINNs) to achieve unprecedented accuracy and speed across diverse partial differential equation problems, often surpassing state-of-the-art results by several orders of magnitude.

Original authors: Joseph Webb, Sadok Jerad, Coralia Cartis

Published 2026-07-03

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Joseph Webb, Sadok Jerad, Coralia Cartis

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

The Big Picture: Teaching a Robot to Solve Physics Puzzles

Imagine you have a robot (a Neural Network) and you want it to solve a complex physics puzzle, like predicting how water flows in a pipe or how heat spreads through a metal plate. These puzzles are described by mathematical rules called Partial Differential Equations (PDEs).

Usually, we teach robots by showing them thousands of examples (data). But with physics, we often don't have enough data. Instead, we tell the robot the rules of the game (the equations) and ask it to figure out the solution on its own. This is called a Physics-Informed Neural Network (PINN).

The Problem:
The paper argues that while this idea is great, the robots are currently terrible at solving these puzzles accurately. They get stuck or produce sloppy answers. The authors say the problem isn't the robot's brain (the network architecture); it's the training method (the optimizer).

Think of the training process like a hiker trying to find the lowest point in a vast, foggy, mountainous valley (the "loss landscape").

Old methods are like a hiker taking tiny, hesitant steps. They get stuck in small dips (local minima) or slide down the wrong side of a mountain because the ground is slippery and uneven (ill-conditioned).
The result: The hiker never reaches the true bottom, and the solution is inaccurate.

The Solution: DSGNAR

The authors introduce a new training framework called DSGNAR (Doubly-Sketched Gauss–Newton with Adaptive Ratio). It's like giving the hiker a high-tech drone, a perfect map, and a smart guide.

Here is how it works, broken down into three simple concepts:

1. The "Drone" View (Doubly-Sketched Gauss–Newton)

To find the bottom of the valley, the hiker needs to see the shape of the terrain. In math, this means calculating a massive, complex map (a matrix). For big problems, this map is so huge it would crash the hiker's computer (Out of Memory).

The Innovation: Instead of drawing the whole map, DSGNAR uses a "drone" to take a smart, compressed snapshot of the terrain from two angles at once.
- Angle 1 (Rows): It looks at the equations and compresses them.
- Angle 2 (Columns): It looks at the variables and compresses them.
The Result: It creates a tiny, manageable map that still holds all the important details. This allows the robot to "see" the shape of the problem clearly without needing a supercomputer. It's like looking at a high-resolution photo of a city from a drone instead of trying to walk every single street to understand the layout.

2. The "Smart Guide" (Adaptive Ratio)

Once the hiker has the map, they need to decide how big a step to take.

Old methods are rigid. They either take tiny steps (too slow) or huge steps (and fall off a cliff). They often get confused about how much "friction" (regularization) to apply to stop the hiker from slipping.
The Innovation: DSGNAR uses a "Smart Guide" that constantly checks a specific metric called the Ratio.
- The guide asks: "Is the step we just took actually helping us get closer to the goal, or did we just guess?"
- Phase 1 (The Setup): At first, the guide tells the hiker to take very cautious steps. The goal isn't to rush to the bottom, but to find a stable, flat area where the ground is solid. This fixes the "slippery" nature of the problem.
- Phase 2 (The Sprint): Once the ground is stable, the guide switches tactics. It says, "Okay, the ground is solid. Now, let's sprint to the bottom!"

This two-stage approach prevents the robot from rushing into a bad solution. It ensures the robot first finds a "well-conditioned" region (a safe spot) before trying to minimize the error aggressively.

3. The "Crowd Sourcing" (Dense Sampling)

Usually, PINNs only check a few random points to see if they are solving the puzzle correctly. This is like checking the weather in only three cities to predict the weather for the whole country.

The Innovation: Because DSGNAR is so efficient (thanks to the "Drone" view), it can check millions of points at once.
The Result: It doesn't need to guess where the tricky parts of the puzzle are. It checks everywhere. If there is a sharp shockwave in a fluid flow, the robot sees it immediately because it's looking at the whole picture, not just a few dots.

What Did They Achieve?

The paper tested this new method on a variety of difficult physics problems:

Burgers' Equation: A classic test for fluid shocks. The new method was 100,000 times more accurate than previous best methods.
High-Dimensional Poisson: A problem with 10 dimensions (imagine a cube with 10 sides). Traditional computers can't solve this, but the new method solved it with incredible precision.
Navier–Stokes: The complex equations for fluid flow (like air over a wing). They solved it much faster and more accurately than before.

The "Magic" Numbers:

In Double Precision (high accuracy), they reached errors so small they are essentially zero for all practical purposes (e.g., $10^{-14}$ ).
In Single Precision (lower accuracy, faster speed), they solved a complex fluid problem in under 10 seconds with an accuracy that rivals much slower, older methods.

The Takeaway

The paper claims that the reason Physics-Informed Neural Networks haven't been perfect yet isn't because the AI is "dumb," but because the training method was clumsy.

By using a compressed map (sketching) to see the problem clearly and a smart guide (adaptive ratio) to know exactly when to be careful and when to be aggressive, they have created a system that solves physics puzzles with the precision of a supercomputer but the speed of a modern AI. They didn't just tweak the settings; they changed the strategy of how the robot learns.

Technical Summary: An Optimisation Framework for the Well-Conditioned Training of Physics-Informed Neural Networks

1. Problem Statement

Physics-Informed Neural Networks (PINNs) have emerged as a mesh-free approach for solving partial differential equations (PDEs), offering advantages in handling complex geometries and high-dimensional problems. However, despite their potential, PINNs have struggled to match the precision of classical numerical solvers. The primary bottleneck is identified not as a lack of data or model capacity, but as an optimisation problem arising from a severely ill-conditioned loss landscape.

Standard first-order optimisers (e.g., Adam) are often insufficient for the high-precision requirements of PINNs, where the goal is to drive residuals to the limits of machine precision. While second-order methods (e.g., BFGS, KFAC, Gauss–Newton) have been explored, they face significant challenges:

Memory and Cost: Storing or computing the full Hessian or Gauss–Newton matrix ( $J^T J$ ) is infeasible for problems with large numbers of collocation points ( $N$ ) and parameters ( $d_\theta$ ).
Ill-Conditioning: The loss landscape is non-convex and highly sensitive to regularisation. Traditional strategies for controlling step length (via trust-region radius $\Delta$ or regularisation parameter $\lambda$ ) often fail to navigate the landscape effectively, leading to premature convergence or instability.
Residual Selection: Existing methods often struggle with the "residual-selection problem," where the placement of collocation points significantly impacts performance.

2. Methodology: DSGNAR

The authors propose DSGNAR (Doubly-Sketched Gauss–Newton with Adaptive Ratio), a scalable second-order optimisation framework designed to confront ill-conditioning directly. The method is built on two core pillars: a doubly-sketched Gauss–Newton model and a novel step-selection strategy based on adaptive ratios.

2.1. Doubly-Sketched Gauss–Newton Model

To avoid the prohibitive cost of forming the full Gauss–Newton matrix, DSGNAR constructs a small, explicit approximation of the Jacobian using randomised linear algebra techniques.

Row Sketching (CountSketch): The residual dimension ( $N$ ) is compressed using CountSketch. This linear, non-invertible operation aggregates residuals from dense sampling into a smaller number of buckets ( $s$ ). This allows for cheap, dense residual sampling, where important regions of the PDE naturally dominate the sketch, effectively averaging out noise and removing the need for manual collocation point selection.
Column Sketching (SRCT): The parameter dimension ( $d_\theta$ ) is compressed using a Subsampled Randomised Cosine Transform (SRCT). Unlike row sketching, this must be a near-isometry to allow the computed step to be "lifted" back to the full parameter space faithfully. SRCT offers the computational efficiency of Fast Fourier Transforms ( $O(d_\theta \log d_\theta)$ ) while preserving singular values.
The Result: A small, square sketched Jacobian $\tilde{J} \in \mathbb{R}^{s \times s}$ (where $s \ll N, d_\theta$ ). This matrix is small enough to compute its Singular Value Decomposition (SVD) cheaply. The SVD enables the rapid calculation of Levenberg–Marquardt steps for any regularisation parameter $\lambda$ without re-solving linear systems.

2.2. Adaptive Ratio and Two-Stage Optimisation

Instead of manually tuning $\lambda$ or $\Delta$ , DSGNAR controls the optimisation trajectory using the decrease ratio $\varrho$ , defined as the ratio of actual objective reduction to predicted reduction.

Conditioning-First Principle: The authors argue that a PINN's accuracy ceiling is determined by the effective regularisation required to achieve a decrease. A "successful" step that reduces the objective early on may lead to a poor solution if the problem remains ill-conditioned.
Two-Stage Strategy:
1. Stage 1 (Exploration): A conservative target ratio ( $\varrho_{\text{Stage 1}} \le 0.2$ ) is used. The optimiser prioritises finding a region of parameter space where the required regularisation $\lambda$ is minimal. This phase drives $\lambda$ down over many orders of magnitude, effectively "conditioning" the problem.
2. Stage 2 (Descent): Once $\lambda$ bottoms out (detected via a log-linear fit), the target ratio is raised ( $\varrho_{\text{Stage 2}} \ge 0.5$ ). The optimiser then aggressively descends towards the minimum, leveraging the well-conditioned region found in Stage 1.
Step Selection: The algorithm probes a geometrically spaced set of trust-region radii, computes the corresponding steps and ratios using the pre-computed SVD, and interpolates to find the step that achieves the current target ratio.

2.3. Supporting Mechanisms

Automatic Reweighting: An adaptive weighting scheme ( $\text{UpdateWeights}$ ) balances the contributions of different PDE conditions (e.g., dynamics vs. boundary conditions) at each iteration, ensuring residuals remain on comparable scales for the sketching process.
Hybrid Precision: To prevent early divergence in double precision due to poor initialisation, the framework runs the first ten iterations in single precision before switching to double precision.

3. Key Contributions

A Conditioning-First Optimisation Philosophy: The paper introduces a rule where step selection is driven by the target decrease ratio rather than direct control of $\lambda$ or $\Delta$ . This prioritises finding well-conditioned regions before aggressive descent.
Doubly-Sketched Gauss–Newton: A novel method that explicitly builds a sketched Jacobian using CountSketch (rows) and SRCT (columns). This allows for dense sampling of residuals and efficient SVD-based step computation, bypassing the memory bottlenecks of matrix-free methods.
Robustness and Scalability: The framework is shown to be robust to architecture choices (MLP, SIREN, SPINN, GaborNet), arithmetic precision (single/double), and initial hyperparameters.
State-of-the-Art Performance: The method achieves unprecedented accuracy and speed across a diverse suite of PDEs.

4. Experimental Results

The framework was evaluated on a comprehensive suite of problems, including nonlinear, chaotic, multi-scale, high-dimensional, and Navier–Stokes equations. Key results include:

Burgers' Equation: Achieved a relative $\ell_2$ error of $7.97 \times 10^{-14}$ in double precision (346.1s), improving upon state-of-the-art results by five orders of magnitude. In single precision, it reached $4.75 \times 10^{-7}$ in under 10 seconds.
Kuramoto–Sivashinsky Equation: Solved the chaotic dynamics with an error of $5.12 \times 10^{-7}$ , significantly outperforming previous benchmarks.
High-Dimensional Poisson (10D): Achieved an error of $5.96 \times 10^{-12}$ , an improvement of eight orders of magnitude over existing methods.
Navier–Stokes (Lid-Driven Cavity): Solved the coupled system with a velocity error of $1.13 \times 10^{-4}$ in 402.9 seconds.
Precision Efficiency: The results demonstrate that single precision can often suffice for high-accuracy PINN solutions, offering a substantially cheaper operating point without sacrificing the ability to reach near-machine-precision limits.

5. Significance and Claims

The paper claims that DSGNAR represents a significant leap forward in the practical application of PINNs. By reframing the training process as a conditioning problem rather than a simple minimisation task, and by leveraging randomised linear algebra to make second-order methods scalable, the framework:

Overcomes the Ill-Conditioning Barrier: It directly addresses the primary obstacle preventing PINNs from rivaling classical solvers.
Delivers Classical Solver Precision: It demonstrates that PINNs can achieve relative errors as low as $3 \times 10^{-16}$ in double precision, rivaling traditional numerical methods.
Provides a Practical Pathway: The ability to trade sketch size for computational cost allows for "aggressive" solves that are orders of magnitude faster than current state-of-the-art while maintaining competitive accuracy.

The authors note that while the method is robust, the bottleneck for certain coupled systems (like Navier–Stokes) may shift from the optimiser to the underlying PINN formulation itself, suggesting future work could involve alternative PINN objectives. Theoretically, while the components (CountSketch, SRCT) are well-understood, a dedicated analysis of the combined doubly-sketched construction remains an open question.