# PDE-aware Optimizer for Physics-informed Neural Networks

arXiv: https://arxiv.org/abs/2507.08118 · HTML: https://mlsanity.com/paper/2507.08118/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 6 of its 9 references that are in the map.
- First posted on arXiv 10 Jul 2025.

## Builds on (most cited first, 6 of 6)

- [Optimizing Neural Networks with Kronecker-factored Approximate Curvature](https://mlsanity.com/paper/1503.05671/) (2015)
- [Shampoo: Preconditioned Stochastic Tensor Optimization](https://mlsanity.com/paper/1802.09568/) (2018)
- [ADAHESSIAN: An Adaptive Second Order Optimizer for Machine Learning](https://mlsanity.com/paper/2006.00719/) (2020)
- [Understanding and mitigating gradient pathologies in physics-informed neural networks](https://mlsanity.com/paper/2001.04536/) (2020)
- [Gradient Alignment in Physics-informed Neural Networks: A Second-Order Optimization Perspective](https://mlsanity.com/paper/2502.00604/) (2025)
- [Sub-Sampled Newton Methods I: Globally Convergent Algorithms](https://mlsanity.com/paper/1601.04737/) (2016)

## Abstract (arXiv)

Physics-Informed Neural Networks (PINNs) have emerged as a powerful framework for solving partial differential equations (PDEs) by embedding physical constraints into the loss function. However, standard optimizers such as Adam often struggle to balance competing loss terms, particularly in stiff or ill-conditioned systems. In this work, we propose a PDE-aware optimizer that adapts parameter updates based on the variance of per-sample PDE residual gradients. This method addresses gradient misalignment without incurring the heavy computational costs of second-order optimizers such as SOAP. We benchmark the PDE-aware optimizer against Adam and SOAP on 1D Burgers', Allen-Cahn and Korteweg-de Vries(KdV) equations. Across both PDEs, the PDE-aware optimizer achieves smoother convergence and lower absolute errors, particularly in regions with sharp gradients. Our results demonstrate the effectiveness of PDE residual-aware adaptivity in enhancing stability in PINNs training. While promising, further scaling on larger architectures and hardware accelerators remains an important direction for future research.