Curvature-Aware Optimization for High-Accuracy Physics-Informed Neural Networks

Anas Jnini; Elham Kiyani; Khemraj Shukla; Jorge F. Urban; Nazanin Ahmadi Daryakenari; Johannes Muller; Marius Zeinhofer; George Em Karniadakis

Figure 1: Schematic illustration of PINN training and optimization methods considered in this work. Our benchmarks include PDEs of elliptic (Helmholtz and Stokes), parabolic (2D viscous Burgers), hyperbolic (inviscid Burgers and 1D Euler) type, and a stiff PK–PD ODE system, characterized by oscillatory solutions, nonlinear diffusion, shock-induced discontinuities, and stiffness. The methods studied here are BFGS, SSBFGS, SSBroyden, NG, SPRING, and SOAP. SOAP ∗ is included for comparison as a curvature-aware preconditioned method, although it is not a classical Newton or quasi-Newton method.

Family	Preconditioner	Properties	Variants
AdaGrad	$\left(\sum_{i=1}^{k}\mathbf{g}_{i}\mathbf{g}_{i}^{\top}\right)^{1/2}$	Scale invariant Rotation invariant	Adam Shampoo SOAP
Newton’s method	$\nabla^{2}\mathcal{L}({\bm{\theta}}_{k})$	Affinely invariant	SSBFGS SSBroyden
Gauss-Newton Natural Gradient	$J_{k}^{\top}J_{k}$	Reparameterization invariant Function-space interpretation	NG / GN

Table 1: Overview of the optimizers considered in this work. The listed properties refer to the idealized exact methods, rather than to the practical variants, which typically rely on additional approximations.

Figure 2: Loss landscapes for two representative PDEs are shown along random directions and projections onto pairs of singular vectors. (a) Top row: 1D Euler equation 59 (Sod problem), projected along orthogonal random directions, leading singular vectors, and next dominant pair (left to right). Bottom row: Stokes equation 42 , following the same projection scheme. Axes correspond to

(\alpha,\beta)

.

Figure 3: Visualization of the NG (top), SSBroyden (middle), and SOAP (bottom) training dynamics for the 2D Helmholtz problem with

a_{1}=1

,

a_{2}=4

, and

k=1

; shown are side-by-side snapshot comparisons of the updates of the optimizers (left) and the current error of the prediction (right).

Problem	Challenges
4.1 2D Helmholtz	Oscillatory solution
4.4 (1+1)D Inviscid Burgers	Shock waves
4.3 (2+1)D viscous Burgers	Steep gradients / parabolic
4.5 (1+1)D Euler with Roe Flux	shock and contact waves but dissipative
4.5.2 (1+1)D Euler with HLLC Flux	Entropy solution, stability, and non-dissipative
4.2 2D Stokes	Multiscale solution

Table 2: Overview over the computational experiments used in this study and the main sources of training difficulty for PINNs in each case.

Figure 4: 2D Helmholtz problem with

a_{1}=1

,

a_{2}=4

, and

k=1

. SSBroyden prediction, reference solution given by Equation ( 39 ), and the pointwise absolute error

|u_{\mathrm{pred}}-u_{\mathrm{ref}}|

.

Figure 7

Case	Optimizer	# collocation points	Relative $L_{\infty}$	Training time (s)	# parameters
1	SSBFGS with TrustRegion	10000	$1.5\times 10^{-8}$	287	2971
2	SSBFGS with TrustRegion	20000	$7.1\times 10^{-8}$	489	2971
3	SSBFGS with TrustRegion	25000	$7.4\times 10^{-9}$	564	2971

Table 4: 2D Helmholtz problem with

a_{1}=1

,

a_{2}=4

, and

k=1

. Relative

l_{\infty}

error and training time for SSBFGS with different numbers of collocation points. All cases use the same PINN architecture with four hidden layers and 30 neurons per layer.

Figure 6: 2D Helmholtz problem with

a_{1}=6

and

a_{2}=6

: (a)

k=1

and (b)

k=100

. During training for

k=100

, collocation points are sampled progressively from

[-0.2,0.2]^{2}

for 10,000 iterations,

[-0.4,0.4]^{2}

for 15,000 iterations,

[-0.7,0.7]^{2}

for 20,000 iterations, and finally the full domain

[-1,1]^{2}

for 40,000 iterations.

Figure 7: 2D Helmholtz with

a_{1}=10

,

a_{2}=10

, and

k=1

: SSBroyden prediction, exact solution ( 39 ), absolute error, and pointwise absolute error

|u_{\rm pred}-u_{\rm exact}|

.

Figure 11

Figure 12

Figure 13

Figure 11: Stokes equation: Streamlines obtained using various optimization strategies: (a) Adam, (b) SOAP, (c) SOAP + SSBFGS (trust-region), (d) SOAP + SSBroyden (Armijo–Wolfe), (e) SOAP + SSBroyden (zoom), (f) Adam + SSBFGS (trust-region), (g) SOAP + SSBFGS (Wolfe), (h) SOAP + SSBFGS (zoom), and (i) NG. The NG achieves comparable accuracy to self-scaled optimizers while significantly reducing runtime. Panel (j) shows the Moffatt strength comparison. The reference solution is computed using a high-order spectral element method.

Figure 12: 2D viscous Burgers equation: Solution obtained using different optimization methods. (a) Best results from a quasi-Newton and line search method (SSBroyden with zoom line search), (b) NG and (c) SOAP optimization. In each panel, the left subfigure shows the PINN solution at

t=0.5

, the middle subfigure shows the spatial distribution of the absolute error (concentrated at the shock location), and the right subfigure shows the solution profile at

x=0.2

. The SOAP optimizer fails to converge at the shock location and exhibits slight overshoots in the shock region. Errors for all optimizer/line-search combinations (including SOAP) are reported in Table 9 .

Optimizer [# Iters.]	Relative $L_{2}$	# parameters	Runtime (s)
SOAP (21001)	$1.7\times 10^{-2}$	3881	414
Adam (1001) + SSBroyden with Zoom (20000)	$1.3\times 10^{-5}$	3881	927
Adam (1001) + SSBroyden with Wolfe (20000)	$5.3\times 10^{-7}$	3881	612
Adam (1001) + SSBFGS with Zoom (20000)	$9.0\times 10^{-7}$	3881	899
Adam (1001) + SSBFGS with Wolfe (20000)	$5.7\times 10^{-7}$	3881	624
NG (20000)	$5.8\times 10^{-7}$	3881	75

Table 9: Unsteady 2D viscous Burgers equation: Performance comparison of the considered optimizers. Columns report the relative

L_{2}

error, total number of parameters, and computational runtime in seconds.

Figure 13: Inviscid Burgers with Roe linearization: Predicted solutions at

t=1

with the LRPINN for both optimizers, together with a reference numerical solution obtained with a third-order WENO scheme. Left panel shows the solution for all

x\in\left[-1,1\right]

, whereas the right panel shows a zoomed-in view of the shock region (

x=0

).

Optimizer	Relative $L_{1}$	Relative $L_{2}$	# collocation points	Training time (s)
SSBroyden	$(1.2\pm 0.2)\times 10^{-2}$	$(7.8\pm 1.9)\times 10^{-2}$	625	20
	$(1.2\pm 0.2)\times 10^{-2}$	$(9.2\pm 2.5)\times 10^{-2}$	1250	21
	$(9.4\pm 1.8)\times 10^{-3}$	$(8.2\pm 2.0)\times 10^{-2}$	2500	22
	$(8.9\pm 1.9)\times 10^{-3}$	$(7.9\pm 1.3)\times 10^{-2}$	5000	26
	$(7.4\pm 1.0)\times 10^{-3}$	$(7.8\pm 1.4)\times 10^{-2}$	10000	31
NG	$(1.5\pm 0.6)\times 10^{-2}$	$(2.6\pm 1.8)\times 10^{-1}$	625	54

Table 10: Inviscid Burgers with Roe linearization: Performance comparison of the SSBroyden and NG optimizers as a function of the number of collocation points. The columns report the relative

L_{1}

and

L_{2}

errors, the number of collocation points

N_{\Omega}

, and the computational runtime in seconds. In all cases, the neural network has 6 hidden layers with 30 neurons each (

4801

trainable parameters). The training process consists of

5000

iterations for all cases. The collocation points are resampled every

10

iterations for SSBroyden and every iteration for NG. SSBroyden is combined with backtracking line search, as implemented in [ 19 ] . For each setting, the training procedure was repeated five times, and the reported errors correspond to the mean values across these runs, together with their associated standard deviations.

Figure 19

Figure 15: Inviscid Burgers equation using relaxation and entropy inequality: The solution of Burgers’ equation, incorporating an entropy inequality and flux relaxation, is obtained using a physics-informed neural network (PINN) trained with the SSBroyden optimizer and a zoom line-search strategy. Panel (a) compares the PINN solution with a DGSEM solution of order

N=3

, where

N

represents the polynomial degree used to approximate the solution within each element. To accurately resolve the shock, the DGSEM employs a shock indicator following [ 31 ] . For further comparison, the inviscid Burgers’ equation is also solved using a third-order WENO scheme [ 32 ] . The relative

L_{2}

error between the PINN and DGSEM solutions shown in panel (a) is

0.70\%

. Panel (b) shows the comparison between the PINN solution and the third-order WENO solution, yielding a relative

L_{2}

error of

0.7056\%

, while panel (c) presents the comparison with the analytical solution, resulting in a relative

L_{2}

error of

0.7172\%

. Panel (d) provides a zoomed-in view near the shock, indicated by the dashed box in panel (a) . This close-up demonstrates that the DGSEM method captures the shock more accurately than both the PINN and WENO schemes and closely matches the analytical solution. Consequently, the discrepancies near the shock dominate the errors observed in panels (a), (b), and (c) . Finally, panels (d), (e), and (f) present the pointwise

L_{1}

errors corresponding to the comparisons shown in panels (a), (b), and (c) , respectively, with the maximum error occurring at the shock location.

Figure 16: 1D Euler equations with Roe linearization: Neural network predictions for the density, velocity and pressure at

t=1

. The exact solution is also plotted for reference.

Figure 22

Figure 18: 1D Euler Equations with HLLC Flux: PINN architecture for 1D Euler equation using adaptive viscosity and HLLC flux.

Figure 19: Euler equations with HLLC flux vs analytical solution: Comparison of the 1D Euler solution for the Sod shock tube problem obtained using the PINN framework augmented with the HLLC flux against the analytical solution. The top row displays the density

\rho

, velocity

u

, pressure

p

, and local Mach number

Ma

(from left to right) at

t=0.15

. The bottom row presents the corresponding pointwise absolute error distributions in space. The results demonstrate excellent agreement between the analytical and PINN-predicted solutions, with accurate resolution of both the shock and the contact discontinuity.

Figure 20: Euler Equations with HLLC flux vs WENO solution: Comparison of the 1D Euler solution for the Sod shock tube problem obtained using the PINN framework augmented with the HLLC flux against a numerical reference solution computed using a WENO scheme [ 32 ] . The top row shows the density

\rho

, velocity

u

, pressure

p

, and local Mach number

Ma

(from left to right) at

t=0.15

. The bottom row presents the corresponding pointwise absolute error in space. The results indicate very good agreement between the WENO reference solution and the PINN predictions, with the shock and contact discontinuity captured accurately. Overall, the discrepancy between the PINN solution and the WENO reference solution is smaller than that observed with respect to the analytical solution, primarily due to the inherent dissipative nature of the WENO scheme.

Optimizer [# Iters.]	Relative $L_{1}$ ( $u,~v,~p$ )	Relative $L_{2}$ ( $u,~v,~p$ )	# parameters	Training time (s)
Adam (5001) + SSBroyden with Zoom (10000) + Analytical	$[1.3~2.2,~0.8]\times 10^{-3}$	$[6.1,~8.9,~2.1]\times 10^{-3}$	2601	460
Adam (5001) + SSBroyden with Zoom (10000) + WENO	$[0.8~1.6,~0.6]\times 10^{-3}$	$[3.4,~6.8,~1.7]\times 10^{-3}$	2601	460

Table 13: 1D Euler Equations with HLLC Flux: Relative

L_{1}

and

L_{2}

error comparison with analytical and WENO reference solutions for the 1D Euler equations using PINN with HLLC flux.

Figure 21: Stiff PK–PD ODE System: Reference vs. PINN (NG).Comparison of the exact solution and the PINN prediction obtained with the NG optimizer, showing accurate reconstruction of the multiscale stiff dynamics.

Figure 28

Parameter	Symbol	Value
Peak FP64 throughput	$\pi$	$67$ TFLOP/s
Peak memory bandwidth	$\beta$	$3.35$ TB/s
Ridge point	$I^{*}$	$10.15$ FLOP/byte

Table 15: NVIDIA H100 SXM hardware specifications used to construct the theoretical roofline model.

Figure 23: Roofline analysis for the SSBFGS and NG optimizers applied to the inviscid Burgers, Euler, and Stokes equations. The solid black line denotes the theoretical roofline, and the vertical dashed line marks the ridge point, which delineates the transition from the memory-bound to the compute-bound regime.

Curvature-Aware Optimization for High-Accuracy Physics-Informed Neural Networks

Abstract

Figures & tables

Appendix

Explore similar work

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

Lightweight Geometric Adaptation for Training Physics-Informed Neural Networks

SS-ESOAP: Self-Scaled Adaptive Preconditioning for Physics-Informed Learning