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

Abstract: Physics-informed neural networks (PINNs) have emerged as a promising route to solve partial differential equations, yet they have struggled to reach the precision of classical solvers. The obstacle is increasingly understood to be one of optimisation, owing to the severely ill-conditioned loss landscape. We present $\textbf{DSGNAR}$: Doubly-Sketched Gauss-Newton with Adaptive Ratio, a scalable second-order optimisation framework that confronts this ill-conditioning and, in doing so, obtains unprecedented accuracy and speed. $\textbf{DSGNAR}$ couples a doubly-sketched Gauss-Newton model with a novel strategy that carefully controls both regularisation and step length. Across a suite of problems spanning nonlinear, chaotic, multi-scale, high-dimensional, and Navier-Stokes, the framework greatly improves on the state of the art: able to attain relative $\ell_2$ errors as low as $3\times10<sup>{-16}$ in double precision, improve contemporary results by five orders of magnitude on the canonical Burgers' equation, and as much as eight orders on a high-dimensional Poisson problem, while remaining markedly faster. We further show that, in single precision, solutions at the limit of round-off error can be obtained very quickly: Burgers' equation to $$\ell_2^{\text{rel}} = 4.75 \times 10<sup>{-7}$$ in under ten seconds. The framework is also robust to the choice of architecture, arithmetic precision, and initial hyperparameters. The code is available at https://www.github.com/wephy/physics-informed-neural-networks

• The paper introduces DSGNAR, a novel second-order optimization framework that overcomes PINN ill-conditioning using a conditioning-first strategy.
• It employs doubly-sketched Gauss–Newton models with adaptive regularisation to efficiently solve complex PDE benchmarks.
• The study demonstrates orders-of-magnitude improvements in accuracy and efficiency, significantly reducing the need for manual hyperparameter tuning.

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

Introduction

This paper presents DSGNAR ("Doubly-Sketched Gauss–Newton with Adaptive Ratio"), a highly scalable second-order optimisation framework aimed at resolving persistent optimisation pathologies in Physics-Informed Neural Networks (PINNs) (2607.02194). Despite promising attributes (mesh-free operation, ease of incorporating physical priors), PINNs have historically lagged behind classical solvers in attainable accuracy, largely due to severe ill-conditioning in their loss landscapes and the inadequacy of first-order optimisers. This work identifies the bottleneck in current PINN training as the challenge of overcoming this ill-conditioning, and introduces a suite of algorithmic innovations centering on sketch-based second-order optimisation, adaptive regularisation, and step selection by explicit control of the objective decrease ratio.

Methodological Innovations

Conditioning-First Optimisation

A key conceptual departure is the conditioning-first step selection principle: rather than greedily minimising the objective at each iteration, DSGNAR prioritises finding regions of parameter space where the local model (as captured by the Gauss–Newton approximation) is maximally well-conditioned. This is accomplished by dynamically adjusting both the regularisation parameter $\lambda$ and the trust-region radius $\Delta$, so that each Gauss–Newton step achieves a prescribed ratio $\varrho$ between actual and predicted loss decrease. Critically, the optimisation proceeds in two phases:

1. Conditioning phase (Stage 1): Early iterations focus on lowering $\lambda$ by maintaining a small decrease ratio ($\varrho \leq 0.2$), driving the system into an operational regime where large steps are justified by the model’s local accuracy.
2. Descent phase (Stage 2): Once the algorithm detects that the minimum attainable regularisation is reached, it transitions to larger ratios ($\varrho \geq 0.5$), thereby accelerating actual descent of the objective.

Figure 2: The optimisation step is selected using geometrically spaced probes for the trust-region radius and the corresponding regularisation; interpolation at the target ratio $\varrho$ determines the next step and regularisation.

Doubly-Sketched Gauss–Newton Model

The intractability of forming and factorising the full Jacobian or Gauss–Newton matrix in PINNs (with millions of residuals) motivates the introduction of doubly-sketched models. DSGNAR compresses the residual (row) dimension with CountSketch and the parameter (column) dimension with a subsampled randomised cosine transform (SRCT), explicitly constructing a small $s \times s$ Jacobian sketch:

• Row Sketch (CountSketch): Allows aggressive, dense residual sampling and efficient aggregation, enabling the algorithm to capture ill-posed regions and high-variance areas in the PDE domain.
• Column Sketch (SRCT): Preserves parameter geometry via a near-isometric random projection, ensuring that the right singular vectors and step directions remain faithful to the original problem.

This combination results in a square Jacobian sketch permitting computationally efficient SVD-based solution of the Levenberg–Marquardt subproblem and fast interpolation across many candidate steps.

Figure 4: (a) Full Jacobian for viscous Burgers' equation; (b) corresponding doubly-sketched Jacobian, showcasing the preservation of main structural features with significant compression.

Probing for Targeted Decrease Ratio

A central element of DSGNAR is step selection by explicit probing. The SVD of the sketched Jacobian allows repeated, cheap evaluation of steps and corresponding regularisations for a geometrically spaced array of radii. The algorithm interpolates the decrease ratio achieved by these steps to select one closest to the current target, rather than tuning regularisation with heuristic or history-dependent schedules.

Adaptive Hyperparameter and Weight Control

Two supporting mechanisms further improve robustness:

• Residual-Based Adaptive Reweighting: Condition weights for PINN objectives (dynamics, initial, and boundary terms) are adjusted online to maintain commensurate residual scales, compensating for extreme disparities induced by physics operators.
• Automatic Stage Transition: The algorithm tracks regularisation trends (using statistical tests on log-linear windows) to switch from the conditioning phase to aggressive descent at optimal timing, removing reliance on manual tuning for schedule changes.

Computational Efficiency

Through explicit sketching and SVD-based step computation, DSGNAR sidesteps the memory and operation count bottlenecks of matrix-free methods, delivering remarkable scalability. The sketch size $s$ serves as a direct accuracy-cost control, where moderate values ($\sim$1/3 to 1/2 the parameter count) suffice for high-fidelity solves.

Numerical Experiments

A comprehensive suite of PDEs, both canonical and highly challenging, demonstrate the efficacy of DSGNAR.

• Viscous Burgers' Equation (Double Precision): Achieves $\Delta$0 in 346.1 seconds—over five orders of magnitude more accurate than prior state-of-the-art PINN optimisers.
• Kuramoto–Sivashinsky Equation: Accurately reconstructs chaotic dynamics with relative errors $\Delta$1, substantially outperforming literature baselines.
• High-Dimensional Poisson (10D): Obtains $\Delta$2 in $\Delta$3 seconds, several orders of magnitude below previous method thresholds.
• Lid-Driven Cavity Navier–Stokes: Resolves both primary and secondary vortical structures at $\Delta$4, achieving errors at the $\Delta$5 scale, highlighting efficacy for coupled, constrained systems.

Figure 5: DSGNAR solution for viscous Burgers' equation, capturing the sharp internal transition region and driving the training loss below $\Delta$6 in only 331 iterations.

Figure 7: Quantitative metrics for Burgers’ equation solution, highlighting low error, rapid training, and convergence to numerical precision limits.

Figure 9: Ten-dimensional Poisson problem solution, showing high-accuracy recovery across a subdomain slice, with globally low error.

Figure 3: Lid-driven cavity Navier–Stokes flow, accurately resolving complex multi-vortex structures with efficient training and low residuals.

The empirical studies confirm that:

• The two-phase optimisation (conditioning-first, then descent) produces characteristic sharp drops in error upon stage transition.
• The trust-region radius remains stable over many orders of magnitude changes in regularisation, validating the search and step selection architecture.
• Single-precision runs achieve errors near the machine epsilon (e.g., $\Delta$7 for Burgers' in under 10 seconds), providing a new operational regime for PINN solvers where ultra-low precision is acceptable.
• The framework is robust to architectural choices (MLPs, SIRENs, GaborNets, SPINNs) and works uniformly across a wide variety of PDE benchmarks.

Figure 12: Burgers' equation in single precision; near machine-level accuracy achieved in roughly 10 seconds.

Practical and Theoretical Implications

Practical

• Orders-of-magnitude accuracy improvement: DSGNAR consistently surpasses—or nullifies—benchmark limits set by classical solvers for PDEs where existing PINN optimisers stagnate.
• Resource efficiency: By exploiting sketching, solutions below $\Delta$8 relative error are attained in a fraction of the time or energy cost previously required, and can be pushed to single-precision limits where appropriate.
• Hyperparameter minimisation: Adaptive weighting, step probing, and automated schedule switching relieve the user of brittle manual tuning, delivering robust out-of-the-box performance.
• Extensibility: The framework admits integration with advanced architectures (e.g., FBPINNs, feature-embedded networks) to potentially raise accuracy ceilings further.

Theoretical

• Algorithmic universality: The approach generalises to diverse PDE structures (e.g., multi-scale, chaotic, high-dimensional, coupled systems) without modification, underlying the broad applicability of its conditioning-first philosophy.
• Linear convergence regime: Experimental results indicate that the method achieves and sustains linear convergence deep into the optimisation trajectory, even when other methods exhibit critical slowdowns.
• Open theoretical questions: The composition of CountSketch and SRCT lacks a complete subspace-embedding analysis in the literature; developing theoretical bounds for doubly-sketched models and direct ratio-targeted step selection remains an area for further research.

Figure 14: Ablation showing that smaller target ratios yield smaller regularisation and higher solution quality, at a (slightly) increased iteration cost; the two-stage approach balances these effects efficiently.

Conclusion

DSGNAR reframes PINN optimisation as a problem of seeking well-conditioned parameter regimes before aggressive model descent, leveraging doubly-sketched second-order models for computational tractability and adaptive control strategies for robust, architecture-independent convergence. Empirical results substantiate strong claims: consistently obtaining accuracy up to ten orders of magnitude lower than contemporaneous PINN approaches, with reduced wall time, and minimal hyperparameter sensitivity. Future directions include the analysis of doubly-sketched embeddings, formalisation of ratio-targeted step theory, and integration with emerging scientific ML frameworks. The framework sets a new baseline in practical and high-precision PINN training (2607.02194).