SS-ESOAP: Self-Scaled Adaptive Preconditioning for Physics-Informed Learning
Abstract Overview
This paper introduces SS-eSOAP, a preconditioner for physics-informed neural network training that extends SOAP-style Kronecker-factored optimization with two additions: a scalar secant-energy correction and adaptive eigenbasis updates with variance-state downscaling. The method is designed for ill-conditioned, high-accuracy PINN regimes where first-order methods often plateau and dense quasi-Newton methods are too costly. The paper also provides theoretical analysis for directional secant matching and for mismatch in the variance state after basis changes. Empirically, the authors evaluate SS-eSOAP on matrix regression tasks and eight PDE benchmarks, with comparisons to Adam, SOAP, Purifying Shampoo, Muon, and other baselines.
Novelty
The distinctive aspect of the work is the combination of SOAP-style Kronecker preconditioning with a self-scaling secant-energy correction adapted to Kronecker geometry, rather than a dense quasi-Newton update. It also replaces fixed basis-refresh schedules with an adaptive trigger and uses variance-state downscaling instead of reprojecting the second-moment state after basis changes.
Results
Across eight PDE benchmarks, SS-eSOAP achieves the lowest final residual on six, while SOAP-family baselines remain stronger on Gray-Scott and Ginzburg-Landau. Reported gains include lower residuals on Wave, Burgers, Boussinesq, Korteweg-de Vries, Allen-Cahn, and Lid-Driven Cavity, and on Boussinesq it reaches a 10^-5 residual target in 4.1 hours with 9.2 GB peak VRAM whereas Adam does not reach that target within 14 hours. Three-seed L2 and H1 error measurements on four representative PDEs are reported to support that lower residuals correspond to better solution accuracy in the studied stiff cases.
Key Points
- SS-eSOAP augments SOAP-style Kronecker-factored preconditioning with a scalar self-scaling correction that matches directional secant energy along recent parameter displacements.
- The optimizer uses an adaptive basis-update trigger based on off-diagonal mass and applies momentum reprojection plus variance-state downscaling to improve stability after basis changes.
- On PDE benchmarks, the method performs best in several stiff high-accuracy regimes but is not uniformly superior, as Purifying Shampoo or SOAP perform better on some problems such as Gray-Scott and Ginzburg-Landau.