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

Motivation

• PINN optimization must address slow convergence, ill-conditioning, spectral bias, and failure modes arising in stiff, multiscale, oscillatory, and shock-dominated problems.These difficulties motivate methods that exploit curvature while remaining computationally and memory efficient.Source: §1 · pp. 1–2

Curvature-aware optimization

• Curvature-aware methods use Hessian-, Gauss–Newton-, or Fisher-based information to produce better-scaled updates, with Gram–Gauss–Newton approximating tractable NTK preconditioning.This can improve conditioning and promote more uniform error reduction across frequency modes.Source: §1 · p. 2

Benchmark scope

• The study benchmarks SSBFGS, SSBroyden, NG, SOAP [72], SPRING, and AdamW across Helmholtz, Stokes, Burgers, Euler, and stiff PK–PD problems.The suite spans elliptic, parabolic, and hyperbolic PDEs plus stiff ODE dynamics, including high-frequency solutions and shock-induced discontinuities.Source: §1 · pp. 2–3

Hyperbolic PDE methods

• The work introduces PINN methodologies for hyperbolic conservation laws, demonstrating convergence for the inviscid Burgers and compressible Euler equations.These methods target discontinuous solutions in high-speed-flow problems.Source: §1 · p. 3

Scaling and analysis

• New batch-training variants and a roofline analysis address the efficiency and scalability of quasi-Newton optimization for large-scale PINN training.The study also develops loss-landscape analyses in function and parameter spaces.Source: §1 · p. 3

Methodological extensions

• The paper proposes a higher-wavenumber Helmholtz methodology to mitigate spectral bias and incorporates Jacobi scaling into NG to improve optimization stability.It also demonstrates PINN applicability to stiff pharmacokinetic–pharmacodynamic ODE systems.Source: §1 · p. 3

2. Review of Considered Optimizers

• PINN training uses a nonlinear least-squares loss, motivating preconditioned updates because standard gradient methods perform inadequately; considered methods approximate three scalable curvature preconditioners.The optimization update uses a preconditioner P_k, with practical methods designed as fast approximations of idealized curvature-based preconditioners.Source: §2 · p. 4

2.1. AdaGrad methods

• AdaGrad-type methods range from Adam’s diagonal adaptivity to Shampoo’s Kronecker-factored within-layer correlations and SOAP’s Adam-style updates in a curvature eigenbasis.Adam [41] and AdamW use exponential moving averages, diagonal preconditioning, and momentum; Shampoo captures richer within-layer correlations, while SOAP [72] combines structured curvature with diagonal adaptivity.Source: §2.1 · pp. 5–6

2.2. Self-Scaled quasi-Newton methods

• Self-scaled quasi-Newton methods construct symmetric positive-definite inverse-Hessian approximations satisfying the secant condition, with BFGS [11, 23, 25, 64] and DFP [18, 24] as special cases.The self-scaled Broyden family [43] [2] introduces iteration-dependent scaling; τ_k ≠ 1 denotes self-scaling, whereas (τ_k, φ_k) = (1, 1) and (1, 0) recover BFGS and DFP.Source: §2.2 · pp. 6–7

Extensions to batched versions

• Because SSBFGS and SSBroyden are typically full-batch and noisy PINN gradients can violate the secant condition, stochastic variants update curvature only when a thresholded condition is satisfied.The condition preserves positive curvature, positive definiteness, and numerical stability; otherwise, updates may be skipped or damped, extending prior batch-training work with L-BFGS.Source: Extensions to batched versions · p. 7

2.3. Natural Gradient Methods

• Natural-gradient methods coincide with Gauss–Newton optimization for PINN least-squares objectives, using Levenberg–Marquardt regularization, Jacobi scaling, and an exact kernel trick that reduces complexity to O(N^2P) when N ≤ P.The kernel trick replaces a P-by-P solve with an N-by-N inversion, while NG implementations also support momentum through SPRING [26] and have reparameterization-invariant function-space interpretations [37].Source: §2.3 · pp. 8–9

3.1. Loss landscape in parameter space

• The loss landscape is visualized on a two-dimensional plane spanned by leading per-sample-gradient singular vectors around a reference parameter vector.Contour plots use projections onto dominant singular-vector pairs to study curvature and anisotropy around θ0.Source: §3.1 · p. 10
• The 1D Euler Sod problem has a rough loss landscape from shocks, whereas the elliptic Stokes problem has a much smoother landscape.Figure 2 compares random-direction projections with leading and next dominant singular-vector pairs for both PDEs.Source: §3.1 · p. 10, Figure 2 · p. 12

3.2. Loss landscape in function space

• NG updates correct prediction errors much better than the other considered optimizers, consistent with its derivation from an infinite-dimensional algorithm.Figure 3 compares NG, SSBroyden, and SOAP updates with current prediction errors for the 2D Helmholtz problem.Source: §3.2 · p. 10, Figure 3 · p. 13

4. Computational Experiments

• The benchmark suite targets stiff, oscillatory, multiscale, and shock-dominated PINN loss landscapes, with results generally computed in double precision and evaluated using relative L1, L2, or L∞ errors.Problems include Helmholtz, Burgers, Euler, Stokes, and PK–PD systems, with reference solutions taken from exact or high-resolution numerical solutions when available.Source: §4 · p. 14

4.1. 2D and 3D Helmholtz equation

• SPINNs consistently outperform vanilla PINNs on the challenging 3D Helmholtz case, reaching errors on the order of 10^-9.For 2D Helmholtz, trust-region variants outperform line-search variants, while progressive frequency or domain warm-up improves quasi-Newton convergence; NG converges directly in high-wavenumber cases.Source: §4.1 · p. 16, Table 7 · p. 20

4.2. 2D Stokes flow over a triangular wedge

• NG matches the accuracy of self-scaled optimizers on Stokes flow while running approximately 24 times faster and reproducing all four eddies, including the weakest with |u| = 10^-8.The most accurate self-scaled configurations use Adam–SSBFGS with trust-region or SOAP–SSBroyden with Armijo and Wolfe line searches.Source: §4.2 · p. 21, Figure 11 · p. 23

4.3. (2+1)D viscous Burgers equation

• The proposed viscous Burgers approach achieves an error of O(10^-7), improving on the reported O(10^-3) state of the art, while NG converges approximately 10–12 times faster than self-scaled optimizers.Errors concentrate near the discontinuity; SOAP struggles at the shock and exhibits slight overshoots there.Source: §4.3 · pp. 22–23, Figure 12 · p. 24

4.4. (1+1)D inviscid Burgers equation

• Because classical residual training becomes unstable and degenerate at inviscid Burgers shocks, the study adds conservation and entropy constraints to recover the physically admissible weak solution [36] [71].The inviscid Burgers section proposes two shock-stabilizing strategies; the Roe-linearization variant estimates shock regions during optimization, uses left and right network states, and hard-enforces boundary conditions through periodic coordinate features.Source: §4.4 · p. 26, §4.4.1 · p. 26

4.4.1. Inviscid Burgers equation with Roe linearization

• The Roe-based LRPINN captures the shock location and jump but produces a more smeared profile than WENO; SSBroyden remains stable with approximately 0.08 relative L2 error at comparatively low cost.NG can achieve slightly lower errors at larger interior-point counts such as NΩ = 5000, but with significantly higher runtime and greater variability at smaller sample sizes.Source: §4.4.1 · pp. 27–28, §4.4.2 · pp. 31–32

4.4.2. Inviscid Burgers equation using flux relaxation with entropy inequality

• Flux relaxation uses separate solution and flux networks, enforcing the algebraic flux constraint and Tadmor entropy inequality through the loss to obtain conservative, entropy-admissible solutions.The entropy inequality permits equality in smooth regions and dissipation across shocks, while violations are penalized in the loss.Source: §4.4.2 · p. 28
• The flux-relaxation PINN with an entropy inequality achieves relative L2 errors of 0.70% against DGSEM, 0.7056% against WENO, and 0.7172% against the analytical solution.The SSBroyden optimizer with zoom line-search is used, while DGSEM captures the shock more accurately than the PINN and WENO; errors are dominated by the shock region.Source: §4.4.2 · p. 29, Figure 15 · p. 31
• The relaxation–entropy formulation yields significantly improved shock-region accuracy over Roe linearization, whereas omitting conservation and entropy modifications prevents convergence to the true solution regardless of optimizer.The entropy-preserving method produces solutions closer to DGSEM and WENO references, while unconstrained PINNs can become nonphysical, unstable, or blow up.Source: §4.4.2 · pp. 30–31

4.5.1. Roe linearization approach

• Roe linearization accurately predicts density, velocity, and pressure in the Sod shock tube, capturing shocks and contact waves without spurious oscillations after viscosity removal.The method modifies the Jacobian in detected shock regions, enforces local conservation, and uses a viscous warm-up before the inviscid phase.Source: §4.5.1 · pp. 33–34, Figure 16 · p. 35

4.5.2. Harten–Lax–van Leer–Contact (HLLC) Flux approach

• The HLLC PINN uses a viscous warm-up followed by inviscid conservative residuals with HLLC flux, trained with Adam before switching to SSBroyden with zoom line search.Viscosity is an additional network output constrained toward zero, while HLLC restores contact waves and supports physically admissible states near strong discontinuities.Source: §4.5.2 · pp. 37–38
• The HLLC framework closely matches analytical and third-order WENO Sod-shock-tube solutions, accurately resolving shock and contact discontinuities with low error across flow variables.Errors are slightly lower against WENO because its dissipative discretization introduces additional numerical diffusion; runtime and relative L1/L2 errors are reported in Table 13.Source: §4.5.2 · pp. 38–39, Figure 19 · p. 38, Table 13 · p. 38

4.5. (1+1)D Euler equations

• Roe and HLLC formulations both address discontinuities in 1D Euler PINNs, while SSBroyden outperforms NG in computational cost and precision for Roe linearization.The Roe approach improves local conservation, whereas HLLC enhances robustness and physical admissibility for strong shocks.Source: §4.5 · p. 32, §4.5.2 · p. 39

Stiff PK–PD ODE System

• The benchmark combines fast and slow PK eigenmodes, delayed transit-compartment responses, Hill-type nonlinear inhibition with Ψ = 20, and injection-driven transients.These mechanisms jointly create stiffness, steep solution-manifold curvature, and non-smooth forcing over the treatment window.Source: §4.6 · pp. 40–41
• A 1000-point piecewise-uniform collocation grid concentrates sampling around the day-2 injection while retaining coverage of slower long-time pharmacodynamic relaxation.The allocation uses 300 points on [0, 1.9] days, 300 on [1.9, 4.0], and 400 on [4.0, 17.0].Source: §4.6 · p. 41
• Curvature-aware and geometry-informed optimizers outperform first-order methods in accuracy and stability, with NG showing the most consistent long-time error decay.SSBFGS and SSBroyden reduce errors more uniformly after the injection transient, whereas Adam retains the largest and most persistent errors.Source: §4.6 · pp. 41–42, Figure 22 · p. 42
• NG achieves the lowest errors in deeper transit states and tumor burden W, while SSBroyden provides uniformly low errors at competitive computational cost.Adam has the largest errors, SOAP improves accuracy by more than one order of magnitude, and classical BFGS performs unevenly across compartments.Source: §4.6 · p. 41, Table 14 · p. 42
• NG is sensitive to increased collocation density near drug administration, which can ill-condition the Hessian/Fisher matrix and cause numerical instability.For other second-order methods, reallocating collocation points affected convergence and stability, sometimes causing stagnation in stiffness-induced sharp local minima.Source: §4.6 · p. 42

Roofline Analysis

• At 67 TFLOP/s FP64 compute, 3.35 TB/s bandwidth, and a 10.15 FLOP/byte ridge point, the H100 roofline separates memory- and compute-bound kernels.The ridge point marks the arithmetic-intensity threshold between the two regimes.Source: §5 · p. 43, Figure 23 · p. 44
• Euler SSBFGS, Stokes SSBFGS, and Stokes NG lie right of the ridge point, indicating compute-bound operation governed primarily by GPU floating-point throughput.Their positions indicate that memory traffic is not the principal bottleneck and bandwidth is efficiently utilized relative to computational workload.Source: §5 · pp. 43–44, Figure 23 · p. 44
• The 1D inviscid Burgers kernel is memory-bound because its compact network performs fewer computations per transferred byte, making execution dominated by data movement.Its low arithmetic intensity places it left of the ridge point, where memory access and transfer constrain performance.Source: §5 · p. 44
• Euler SSBFGS reaches about 75% memory-bandwidth utilization with approximately 3,000 parameters, while Stokes SSBFGS uses approximately 33,667 parameters and a roughly 9 GB Hessian.Euler’s smaller model enables cache residence and data reuse; Stokes nearly saturates compute and memory bandwidth but incurs substantial DRAM traffic.Source: §5 · p. 44
• For Gauss-Newton, AI = 391.4 FLOP/byte and achieved performance is 29.640 TFLOP/s, confirming compute-bound operation.The performance is measured as reliable FLOPs divided by wall-clock step time.Source: §5 · p. 44

Scalability Limitation

• Dense SSBFGS incurs O(n^2) memory growth as networks scale, driving solvers toward the bandwidth ceiling and limiting arithmetic intensity and GPU compute efficiency.The associated dense Hessian footprint becomes a dominant scalability constraint for second-order PINN optimization.Source: §5 · pp. 44–45

Implementation

• Publicly available JAX implementations provide efficient and flexible construction of optimizer variants.The authors state that the codebase is intended to facilitate wider adoption and reproducibility.Source: §6 · p. 45
• Optimizer instances expose tolerances, gradient-norm convergence criteria, line-search strategies, and Newton-type descent updates for problem-specific configuration.These components include relative and absolute tolerances, L1, L2, or L∞ norms, backtracking or trust-region searches, and the update from Equation (11).Source: §6 · p. 45
• Optimistix [58] and the AbstractQuasiNewton class enable modular customization of SSBFGS and SSBroyden variants.The paper illustrates SSBFGS with a trust-region line search and L∞-norm convergence criterion in Codelisting 1.Source: §6 · p. 45

Code Availability

• The associated GitHub repository, https://github.com/CrunchOptimizer, is publicly accessible and will be updated upon manuscript acceptance.Source: §6 · p. 45

Overall findings

• Curvature-aware optimizers significantly improved PINN training, with SSBFGS, SSBroyden, and NG especially effective for ill-conditioned problems and very high accuracy.Optimization alone was insufficient for discontinuous inviscid Burgers and Euler problems, which required shock-aware and conservation-consistent formulations.Source: §7 · p. 46

Elliptic problems

• For elliptic problems, NG was often fastest; on 2D Helmholtz, NG reached order 10−9 error in 251 s versus 319 s for SSBroyden and 564 s for SSBFGS.SOAP stagnated at order 10−4, while in Stokes NG was approximately 24× more computationally efficient than self-scaled optimizers at matched accuracy.Source: §7 · p. 46

Parabolic problems

• For parabolic problems, curvature-aware methods produced high-accuracy solutions, with NG combining strong accuracy and lower computational cost while quasi-Newton methods remained competitive and robust.The natural-gradient method is recommended when fast convergence is the primary objective; SSBFGS and SSBroyden are alternatives when robustness and very high final accuracy matter.Source: §7 · p. 46

Hyperbolic problems

• For hyperbolic problems, accurate solutions required shock-aware, conservation-consistent methods including Roe linearization, entropy-constrained relaxation, and HLLC fluxes.NG and Broyden-family methods behaved similarly after the formulation addressed shocks and contact discontinuities; optimization alone was generally insufficient.Source: §7 · p. 46

Stiff PK–PD ODE

• In the stiff PK–PD system, NG achieved relative L2 errors below 10−3 in deeper transit states, while SSBroyden and SSBFGS balanced accuracy across compartments.The system combined multi-timescale dynamics, extreme nonlinear saturation Ψ = 20, and discontinuous forcing; classical BFGS was uneven and Adam struggled, particularly in intermediate compartments.Source: §7 · p. 46

Scalability and implementation

• Roofline analysis found main second-order kernels near the compute roofline, whereas small 1D inviscid Burgers problems were memory-bound; a validated batched SSBFGS/SSBroyden algorithm improved stability and convergence.The batch method was tested on Helmholtz and Poisson equations, supporting application to large-scale data-driven problems despite dense second-order memory-scaling limits.Source: §7 · p. 46

• Positive-definite matrix updates are required to preserve descent directions, constraining the self-scaling parameters.The parameters (τ_k, ϕ_k) cannot be chosen arbitrarily, and the update formulas for H_k or B_k must preserve positive definiteness.Source: A. Details on Self-Scaled Quasi-Newton… · p. 50
• Under Wolfe line searches and convex objectives, self-scaled Broyden iterations converge superlinearly when 0 < τ_k ≤ 1 and η_kτ_k ∈ [0, 1).This convergence result is attributed to Al-Baali.Source: A. Details on Self-Scaled Quasi-Newton… · p. 50
• Scaled BFGS using the Oren–Luenberger factor is competitive with unscaled BFGS and has produced accurate results in PINN applications [43].The scaled choice corresponds to ϕ_k = 1, while unscaled BFGS corresponds to τ_k = 1; the cited studies also report accurate results across physics problems.Source: A. Details on Self-Scaled Quasi-Newton… · pp. 50–51
• Inexact Armijo–Wolfe backtracking and zoom line searches enforce curvature conditions needed to preserve positive definiteness.The curvature requirement is y_k^T s_k > 0, with strong Wolfe conditions enforced through the zoom procedure.Source: A. Details on Self-Scaled Quasi-Newton… · p. 51
• A linear trust-region method provides an alternative globalization mechanism by truncating the Cauchy direction and adapting the radius from actual-to-predicted reduction.The trust-region radius Δ_k is updated using the ratio ρ_k, with a detailed treatment referenced in Chapter 4 of.Source: A. Details on Self-Scaled Quasi-Newton… · p. 51

• The 2D Helmholtz analysis considers a_1 = 1 and a_2 = 4 for both K = 10 and K = 100.Source: B. Details on 2D Helmholtz problem · p. 51
• The optimizer comparisons use a fully connected network with four hidden layers and 30 neurons per layer; the Fourier feature mapping uses m = 1 and four input features.Table 16 specifies this setup for K = 1.Source: Table 16 · p. 52

• The appendix presents detailed pseudocode for the stochastic SSBFGS and SSBroyden optimizers in Algorithm 2.Source: C. Stochastic SSBFGS and SSBroyden… · p. 52

Proof of Theorem 2.1

• The proof differentiates the gradient along the search path, applies the Fundamental Theorem of Calculus and an integral Mean Value Theorem, then uses Hessian continuity.It defines g(t) as the gradient along x_k−1 + t s_k−1 and obtains g′(t) = ∇^2f(x_k−1 + t s_k−1)s_k−1.Source: D. Proof of Theorem 2.1 · p. 52
• Uniform positive definiteness assumption (18) supplies the final bound needed to complete the theorem’s proof.Source: D. Proof of Theorem 2.1 · p. 52

Lazy Stochastic sSSBFGS and sSSBroyden for Mini-Batch Training

• Algorithm 2 trains with mini-batches by initializing H_0 to I_n and g_0 to zero, then updating parameters using Δθ_k ← −ηH_kg_k.The procedure takes learning rate η, damping threshold τ, mini-batch size B, and K iterations, and stores the previous gradient after each update.Source: D. Proof of Theorem 2.1 · p. 53

E. Correctness of Algorithm 2

• In FP32 batch training, sSSBFGS reaches loss values on the order of 10^-7 and attains machine precision, whereas Adam does not converge to the same accuracy.The comparison uses the one-dimensional Poisson equation and its Dirichlet boundary conditions.Source: E. Correctness of Algorithm 2 · p. 53, Figure 24 · p. 54