PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2604.05230v1cs.LGcs.AImath.NAmath.OC

TL;DR

PINN training requires efficient optimization for difficult PDEs and ODEs, especially when losses are ill-conditioned, stiff, oscillatory, or shock-dominated. The paper develops and evaluates curvature-aware optimizers, introduces hyperbolic-PDE methods and batch-training variants, and finds strong performance across challenging regimes while identifying memory and conditioning limits.

  • Problem

    PINN optimization must achieve fast, accurate, and robust convergence across differential-equation problems with ill-conditioned, stiff, oscillatory, and discontinuous solution landscapes.

  • Method

    The paper develops efficient Natural Gradient, self-scaled BFGS and Broyden, Jacobi-scaled, hyperbolic-PDE, and batch-training methods for PINNs.

  • Results

    Curvature-aware optimizers performed strongly across the benchmarks; for one 2D Helmholtz case, SSBFGS, SSBroyden, and NG reached errors of order 10^-9, with NG requiring 251 s.

  • Takeaways & Limitations

    The results support curvature-aware optimization as a practical approach for high-accuracy PINN training across diverse differential-equation regimes, with specialized methods needed for discontinuities and large-scale training.

  • Takeaways & Limitations

    NG can become numerically unstable with increased collocation density near drug administration, while dense SSBFGS has an O(n^2) memory footprint as network size grows.

Abstract

from arXiv · show

Efficient and robust optimization is essential for neural networks, enabling scientific machine learning models to converge rapidly to very high accuracy -- faithfully capturing complex physical behavior governed by differential equations. In this work, we present advanced optimization strategies to accelerate the convergence of physics-informed neural networks (PINNs) for challenging partial (PDEs) and ordinary differential equations (ODEs). Specifically, we provide efficient implementations of the Natural Gradient (NG) optimizer, Self-Scaling BFGS and Broyden optimizers, and demonstrate their performance on problems including the Helmholtz equation, Stokes flow, inviscid Burgers equation, Euler equations for high-speed flows, and stiff ODEs arising in pharmacokinetics and pharmacodynamics. Beyond optimizer development, we also propose new PINN-based methods for solving the inviscid Burgers and Euler equations, and compare the resulting solutions against high-order numerical methods to provide a rigorous and fair assessment. Finally, we address the challenge of scaling these quasi-Newton optimizers for batched training, enabling efficient and scalable solutions for large data-driven problems.

1. Introduction

The paper studies how curvature-aware and quasi-second-order optimization can improve PINN training across difficult differential-equation regimes. It combines optimizer analysis, new methods for hyperbolic problems, computational-scaling analysis, and batch-training strategies.

  • Optimization motivation: Curvature information can improve PINN optimization by rescaling ill-conditioned loss landscapes and promoting more uniform error reduction across solution modes.The paper links poor PINN training to ill-conditioned NTK spectra, slow convergence, and spectral bias.
  • Optimizer study: The study systematically compares SSBFGS, SSBroyden, NG, SOAP, and AdamW across challenging PINN problems spanning elliptic, parabolic, hyperbolic, and stiff ODE regimes.Benchmarks include Helmholtz, Stokes flow, viscous and inviscid Burgers equations, Euler equations, and PK–PD systems.
  • Hyperbolic PDE methods: The paper introduces PINN methodologies for hyperbolic conservation laws, demonstrating convergence for inviscid Burgers and compressible Euler equations with discontinuous solutions.These methods target shock-dominated problems where optimization alone is insufficient.
  • Learning-dynamics analysis: The work develops a loss-landscape framework in function and parameter spaces to analyze PDE learning dynamics and optimizer behavior.The study visualizes both parameter-space trajectories and function-space updates relative to prediction error.
  • Scalability: A roofline analysis characterizes quasi-Newton computational efficiency, while a new batch-training strategy targets scalability for large PINN workloads.The paper also addresses quasi-Newton limitations in large-scale and data-driven training.
  • Natural Gradient methods: Jacobi scaling is incorporated into Natural Gradient methods to improve numerical stability for ill-conditioned Gauss–Newton systems.The scaling equilibrates the Gramian and normalizes its diagonal entries.

2. Review of Considered Optimizers

The section reviews optimizer families used for PINN training, emphasizing scalable curvature-aware preconditioning and practical adaptations for numerical stability and large batches.

  • Optimization framework: PINN optimization methods use preconditioning to incorporate curvature information and produce better-scaled parameter updates than basic gradient descent.The reviewed methods approximate Hessian-, Gauss–Newton-, or related curvature structures through computationally tractable preconditioners.
  • AdaGrad methods: Adam and AdamW use diagonal exponential-moving-average preconditioners, providing scalable adaptive updates while neglecting cross-parameter correlations.This diagonal structure limits correction for intra-layer ill-conditioning.
  • AdaGrad methods: Shampoo and SOAP capture richer within-layer curvature through Kronecker-factored matrices, with SOAP applying Adam-style updates in an eigenbasis.These approximations trade dense curvature modeling for scalability and structured layerwise information.
  • Self-scaled quasi-Newton methods: Self-scaled quasi-Newton methods update inverse-Hessian approximations, while stochastic variants delay updates when noisy gradients violate the secant condition.The batched extension targets problems too large for a single GPU and seeks numerical stability under gradient noise.
  • Natural Gradient Methods: Gauss–Newton methods regularize and scale the Jacobian Gramian, using Levenberg–Marquardt damping and Jacobi scaling to improve stability.When the batch is smaller than the parameter dimension, the push-through identity reduces inversion cost to O(N^2P), linear in P.

3. Landscape Visualization

The study visualizes PINN loss landscapes in two-dimensional parameter subspaces and compares optimizer updates in parameter and function space. These views connect landscape geometry with the error-correction behavior of curvature-aware methods.

  • Loss landscape projections: Loss landscapes are evaluated on two-dimensional planes around a reference parameter vector using random directions and dominant singular-vector pairs.The loss is sampled as L(θ0 + αv1 + βv2), with α and β defining the plotted coordinates.
  • Loss landscape projections: The first two singular-vector directions capture the largest variation in gradient space and define the principal plane used for contour visualization.The construction uses the leading columns of V from the singular value decomposition of the gradient matrix.
  • Optimizer updates: All optimizers generate parameter-space steps θ_k+1 = θ_k + δ_k, with the update direction determined by each optimizer’s rule.For NG/Gauss–Newton, the direction uses a regularized curvature expression involving J_k, λ, and r_k.
  • Optimizer updates: Function-space pushforwards make optimizer updates easier to interpret, and NG visibly corrects the prediction error more effectively than SSBroyden and SOAP in the Helmholtz visualization.Figure 3 shows side-by-side optimizer updates and current prediction errors for the 2D Helmholtz problem.
  • Loss landscape projections: Figure 2 compares random-direction, leading-singular-vector, and next-pair projections for the 1D Euler Sod problem and the Stokes equation.The rows correspond to the two PDEs, while the columns represent the three projection choices.

4. Computational Experiments

The experiments evaluate curvature-aware and quasi-Newton optimizers across difficult PINN problems spanning oscillatory, high-frequency, elliptic, shock-dominated, and hyperbolic regimes. Results show optimizer- and formulation-dependent trade-offs between accuracy, robustness, and runtime, alongside new PINN methods for Burgers and Euler problems.

  • Helmholtz and Stokes: For 2D Helmholtz, trust-region variants outperform line-search variants for SSBFGS and SSBroyden, while progressive frequency warm-up improves their convergence speed and robustness.NG also converges without warm-up for the (6, 6) and (10, 10) cases.
  • Helmholtz and Stokes: SPINNs outperform standard PINNs on the challenging 3D Helmholtz case, achieving errors on the order of 10^-9 for (a1, a2, a3) = (4, 4, 3) and k = 1.The comparison uses the same evaluation protocol and reports lower solution error across the listed metrics.
  • Helmholtz and Stokes: NG reaches accuracy comparable to self-scaled optimizers on the 2D Helmholtz problem while running approximately 24 times faster, at about 150 seconds versus roughly one hour.The comparison identifies NG as the most computationally efficient among the reported configurations despite similar accuracy.
  • Inviscid Burgers: The proposed Burgers method combines separate solution and flux networks with an algebraic conservation constraint and Tadmor entropy inequality to enforce entropy-compliant solutions.The algebraic constraint alone is not sufficient to guarantee entropy compliance.
  • Inviscid Burgers: The inviscid Burgers formulations recover shock location and jump structure, while the relaxation–entropy inequality approach yields lower error than the Roe Linearization approach.The Roe-based method improves stability but retains some numerical diffusion near the shock.
  • Inviscid Burgers: SSBroyden remains stable across LRPINN configurations with approximately 0.08 relative L2 error, whereas NG becomes competitive at larger interior-point counts but is more variable with fewer points.The reported pattern favors SSBroyden for robustness and computational efficiency under the tested sampling conditions.

4.5. (1+1)D Euler equations

The section develops PINN approaches for 1D Euler conservation laws, focusing on Roe linearization and HLLC fluxes to handle shocks and contact discontinuities. Both approaches produce accurate solutions, with HLLC emphasizing robustness and physical admissibility for strong discontinuities.

  • Motivation: The section introduces Roe linearization and HLLC flux formulations because Euler discontinuities can rapidly degrade the PINN loss, making optimizer choice alone insufficient.The two approaches target discontinuous solutions through modified numerical formulations within the PINN framework.
  • Roe linearization: The Roe strategy modifies the Euler Jacobian in shock regions while enforcing diagonalizability, consistency, and conservation for left and right states.Shock regions are identified using strongly negative spatial velocity derivatives.
  • Roe linearization: The Roe-based PINN accurately reproduces density, velocity, and pressure for the Sod shock tube, including sharp shocks, contact waves, and rarefaction transitions.The predictions show very few or no transition points in shocks and contact waves, with abrupt rarefaction changes.
  • HLLC flux: The HLLC formulation first trains a viscous approximation with viscosity as a network output, then switches to the inviscid conservative residual with HLLC numerical flux.The viscosity is driven toward zero through a soft loss constraint during warm-up.
  • HLLC flux: The HLLC-based PINN reproduces the principal Sod shock-tube flow features across density, velocity, pressure, and Mach number while comparing favorably with the exact solution.The evaluation also reports spatial absolute errors for the predicted profiles.
  • Comparison: Both SSBroyden and NG converge with the Roe-linearization approach, while Roe improves local conservation and HLLC provides stronger robustness and physical admissibility for strong shocks.The section characterizes the methods as complementary treatments of discontinuous Euler solutions.

4.6. Stiff PK–PD System with Discontinuous Drug Administration

The stiff PK–PD benchmark combines fast and slow pharmacokinetic and pharmacodynamic dynamics with discontinuous injection forcing. Among the optimizers, curvature-aware methods provide more accurate and stable solutions, with NG showing the strongest long-time error decay but sensitivity to collocation density.

  • Problem setup: The PK–PD system is stiff because fast distribution, slow elimination, delayed transit responses, nonlinear growth inhibition, and injection-driven transients coexist.The compartmental eigenvalues differ by several orders of magnitude, while chemotherapy injections create sharp localized forcing.
  • Optimization results: Adam has the largest errors across state variables, SOAP improves accuracy by more than one order of magnitude, and structured quasi-Newton methods provide more uniform error reduction.Classical BFGS improves late-time accuracy but is unstable across intermediate compartments.
  • Optimization results: The time-resolved errors show a sharp rise for every method immediately after drug injection, followed by smoother decay for SSBFGS and SSBroyden.NG exhibits the most consistent long-time error decay across the stiff multiscale dynamics.
  • Optimization results: Curvature-aware and geometry-informed optimizers significantly outperform first-order methods in accuracy and stability for the stiff PK–PD problem.NG reaches the lowest error levels in X_2–X_4 and maintains strong robustness after the transient phase.
  • Limitations: NG is particularly sensitive to increased collocation density near drug administration, which can ill-condition the Hessian or Fisher matrix and cause numerical instability.Redistributing collocation points also affects convergence speed and can cause stagnation for other second-order methods.

5. Performance Analysis of Optimizers using Roofline Models

The roofline analysis identifies whether PINN optimizer kernels are limited by memory movement or floating-point computation. Euler and Stokes kernels are compute-bound, whereas inviscid Burgers is memory-bound, and dense SSBFGS faces a scaling bottleneck as networks grow.

  • Roofline model: The H100 roofline has a 67 TFLOP/s FP64 ceiling, a 3.35 TB/s memory-bandwidth ceiling, and a ridge point near 10.15 FLOP/byte.The ridge point separates memory-bound from compute-bound kernels.
  • Roofline results: Euler SSBFGS, Stokes SSBFGS, and Stokes NG operate to the right of the roofline ridge point in the compute-bound regime.Their performance is governed primarily by peak floating-point throughput rather than memory bandwidth.
  • Roofline results: The 1D inviscid Burgers kernel is memory-bound because its compact network performs relatively little computation per byte transferred.Its lower arithmetic intensity places it to the left of the ridge point.
  • Kernel comparisons: Euler SSBFGS achieves higher arithmetic intensity with an approximately 3,000-parameter network and a dense Hessian that largely fits in GPU L2 cache.The configuration reaches about 75% memory-bandwidth utilization.
  • Kernel comparisons: Stokes SSBFGS uses approximately 33,667 parameters and a roughly 9 GB dense Hessian, saturating compute and memory bandwidth while incurring substantial DRAM traffic.The larger Hessian lowers arithmetic intensity relative to the Euler case despite the larger problem size.
  • Scalability limitation: Dense SSBFGS has an O(n^2) memory footprint that becomes dominant as network size increases, constraining arithmetic intensity and compute efficiency.This scaling challenge pushes larger-network solvers toward the memory-bandwidth ceiling.

6. Reproducibility

The implementation is designed for reproducibility and flexible optimizer customization. JAX and Optimistix support modular choices of tolerances, norms, line searches, and descent updates, while the associated codebase is publicly available.

  • Implementation: All optimizer implementations use JAX to support efficient and flexible construction of optimizer variants.The authors provide a publicly accessible GitHub repository for the implementation.
  • Customization: Optimizer instances expose relative and absolute tolerances, convergence norms, line-search strategies, and Newton-type descent updates as configurable components.These settings allow SSBFGS and SSBroyden variants to be tailored to specific problems.

7. Summary

The paper evaluates curvature-aware optimizers across elliptic, parabolic, hyperbolic, and stiff ODE problems, finding improved PINN training and strong performance from NG, SSBFGS, and SSBroyden. It also analyzes hardware efficiency and introduces batch training for scalable quasi-Newton optimization.

  • Curvature-aware optimizers significantly improve PINN training over standard first-order approaches across the studied PDE and stiff PK–PD ODE problems.The comparison emphasizes convergence speed, robustness, accuracy, and scalability.
  • Elliptic problems: NG was often fastest for elliptic problems, while SSBFGS and SSBroyden achieved comparable final accuracy at times higher computational cost.For the 2D Helmholtz case with a1 = 1, a2 = 4, and k = 1, all three reached errors of the stated order, though the passage is truncated before the value.
  • Parabolic problems: For 2D viscous Burgers, curvature-aware methods again outperformed first-order training, with NG combining strong accuracy and lower computational cost.The quasi-Newton methods also delivered competitive performance.
  • Hyperbolic problems: For discontinuous hyperbolic PDEs, the paper recommends combining curvature-aware optimization with shock-aware, conservation-consistent PINN formulations because optimization alone is insufficient.Examples include Roe-based linearization and HLLC-flux constructions.
  • Scalability and hardware efficiency: Roofline analysis found major second-order kernels mostly compute-bound near the roofline ceiling, while small 1D inviscid Burgers problems remained memory-bound.The analysis also identifies memory-scaling limitations of dense second-order optimization.
  • Scalability and hardware efficiency: A batch training algorithm for SSBFGS and SSBroyden was validated on Helmholtz and Poisson equations, demonstrating improved stability and convergence rates for large-scale data-driven problems.

A. Details on Self-Scaled Quasi-Newton Methods

The self-scaled quasi-Newton methods rely on positive-definite updates, curvature-preserving line searches, and parameter choices supported by convergence theory. The section also considers trust-region globalization and scaling choices used in PINN applications.

  • Convergence requirements: Positive definiteness of H_k or B_k is required so quasi-Newton search directions remain descent directions.This requirement constrains the scaling parameter η_k.
  • Convergence requirements: For convex objectives, self-scaled Broyden iterations converge superlinearly when Wolfe conditions hold, 0 < τ_k ≤ 1, and η_kτ_k ∈ [0, 1).
  • Scaling choices: The Oren–Luenberger scaling choice is competitive with unscaled BFGS, including in PINN applications, while variable τ_k and ϕ_k choices provide another scaling strategy.The passage identifies ϕ_k = 1 with the scaled BFGS case and τ_k = 1 with unscaled BFGS.
  • Scaling choices: Scaling choices introduced in prior work have produced very accurate results for a range of physics problems in PINNs.
  • Line searches: Armijo–Wolfe and zoom line searches enforce curvature conditions needed to preserve positive definiteness in quasi-Newton updates.The step length α_k is selected using sufficient-decrease and curvature conditions.
  • Trust-region globalization: A linear trust-region alternative computes a truncated Cauchy direction and updates its radius using actual-to-predicted reduction.The radius Δ_k is adjusted adaptively according to ρ_k.

B. Details on 2D Helmholtz problem

The appendix examines the 2D Helmholtz problem for specified coefficient and frequency settings and documents the common network and Fourier-feature configuration used in one benchmark table.

  • Problem settings: The detailed Helmholtz analysis considers a1 = 1 and a2 = 4 for K = 10 and K = 100.
  • Benchmark configuration: Table 16 evaluates the 2D Helmholtz problem with a1 = 1, a2 = 4, and K = 1.
  • Benchmark configuration: All optimizers in Table 16 use a fully connected network with four hidden layers and 30 neurons per layer.
  • Benchmark configuration: The Fourier feature mapping in Table 16 uses mode m = 1, producing four input features.

C. Stochastic SSBFGS and SSBroyden Methods

The appendix provides detailed pseudocode for stochastic SSBFGS and SSBroyden optimization methods. The algorithms are presented together in Algorithm 2.

  • Algorithm 2: Algorithm 2 contains detailed pseudocode for stochastic SSBFGS and SSBroyden optimizers.
  • Algorithm 2: The appendix focuses on stochastic variants of both self-scaled quasi-Newton methods.
  • Algorithm 2: The pseudocode is provided as an implementation reference for the two optimizers.

D. Proof of Theorem 2.1

The section establishes the update structure and supporting calculus for Algorithm 2, which applies lazy stochastic sSSBFGS and sSSBroyden to mini-batch training.

  • Proof: The proof defines a gradient function along the step, differentiates it using the Hessian, and applies integral and mean-value arguments under uniform positive definiteness.The argument assumes f ∈ C2(R^n) and uses continuity of the Hessian.
  • Algorithm 2: Algorithm 2 initializes the Hessian approximation as I_n and the previous gradient as zero before mini-batch iterations.The inputs include learning rate η, damping threshold τ, mini-batch size B, and iteration count K.
  • Algorithm 2: The parameter update uses a direction formed from the Hessian approximation and current gradient, then stores the current gradient for the next iteration.The update is given by Δθ_k ← −ηH_kg_k followed by θ_k+1 ← θ_k + Δθ_k.

E. Correctness of Algorithm 2

The section tests Algorithm 2 on a one-dimensional Poisson equation by comparing batch-trained sSSBFGS with Adam. sSSBFGS reaches FP32-level accuracy and loss values around 10^-7, whereas Adam does not reach the same level.

  • Batch-training comparison: sSSBFGS attains FP32-level accuracy while Adam does not converge to the same level, indicating machine-precision performance for sSSBFGS in this setting.The result uses the one-dimensional Poisson problem with Dirichlet boundary conditions u(0) = 0 and u(1) = 0.
  • Batch-training comparison: Loss values around 10^-7 show that sSSBFGS significantly outperforms Adam in FP32 batch training for the 1D Poisson equation.The comparison is reported in Figure 24.
Loading 2604.05230v1…