cs.LGAug 29, 2026

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

Authors: Guangyuan Wang, Mads Toftrup, Sebastian Loeschcke, Yixuan Wang, Anima Anandkumar

Organizations: McGill University · California Institute of Technology · Aarhus University · University of Copenhagen

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

Physics-informed neural networks (PINNs) often face ill-conditioned objectives that limit high-accuracy training. Dense quasi-Newton methods improve local conditioning but require expensive optimizer state, while Kronecker-factored methods such as SOAP scale to larger networks but rely on periodic basis updates. We introduce \method, which augments SOAP-style preconditioning with a scalar secant-energy correction adapted to Kronecker geometry and an adaptive basis update followed by variance-state downscaling. We characterize the directional secant matching induced by the scalar correction and give a bound on variance-state mismatch across basis changes. Across eight PDE benchmarks, \method attains the lowest final residual on six, including Burgers and Boussinesq, while SOAP-family baselines perform better on Gray-Scott and Ginzburg-Landau. On Boussinesq, \method reaches a residual of 10−510^{-5} in 4.1 hours with 9.2 GB peak VRAM, while Adam does not reach this target within 14 hours. Three-seed L2L^2 and H1H^1 errors on four representative PDEs support the link between lower residuals and improved solution accuracy. These results position \method as a scalable option for stiff, high-accuracy physics-informed training, rather than a uniform replacement for existing optimizers.

Figures & tables

Figure 1: Matrix Quadratic Regression Task. SS-eSOAP consistently achieves faster convergence and lower residuals in the quadratic regime where the self-scaling mechanism captures the spectrum.
BenchmarkAdamMuonSOAPPurifying ShampooSS-eSOAP (Ours)
Wave2.15×10−32.15\times 10^{-3}5.12×1045.12\times 10^{4}1.45×10−31.45\times 10^{-3}1.18×10−31.18\times 10^{-3}4.12×10−5\mathbf{4.12\times 10^{-5}}
Burgers4.21×10−64.21\times 10^{-6}3.85×10−43.85\times 10^{-4}2.15×10−92.15\times 10^{-9}1.12×10−91.12\times 10^{-9}5.54×10−10\mathbf{5.54\times 10^{-10}}
Allen-Cahn2.10×10−72.10\times 10^{-7}1.25×10−71.25\times 10^{-7}8.45×10−78.45\times 10^{-7}3.15×10−93.15\times 10^{-9}1.85×10−9\mathbf{1.85\times 10^{-9}}
Boussinesq5.12×10−45.12\times 10^{-4}2.57×10−22.57\times 10^{-2}9.15×10−69.15\times 10^{-6}2.85×10−62.85\times 10^{-6}8.42×10−7\mathbf{8.42\times 10^{-7}}
Korteweg-de Vries1.85×10−81.85\times 10^{-8}1.93×10−31.93\times 10^{-3}8.12×10−108.12\times 10^{-10}1.75×10−71.75\times 10^{-7}3.52×10−10\mathbf{3.52\times 10^{-10}}
Gray-Scott3.10×10−63.10\times 10^{-6}4.20×10−44.20\times 10^{-4}2.10×10−72.10\times 10^{-7}5.12×10−9\mathbf{5.12\times 10^{-9}}7.46×10−77.46\times 10^{-7}
Table 1: Final PDE residual after the fixed training budget. Lower is better. Each entry reports the selected run under the tuning protocol in Appendix C.4 . Bold marks the lowest value within each row. Residual is an optimization metric. Physical solution errors and three-seed statistics appear in Appendix 8 .
Figure 2: Ablation studies on 2D Boussinesq (top) and Allen-Cahn (bottom). (a) and (c): Stability analysis of the eigenbasis update. (b) and (d): Efficacy of the self-scaling curvature correction. The proposed Kronecker-adapted factor τk\tau_{k} (in blue) accelerates and stabilizes convergence compared to the unclipped counterpart, standard Oren-Luenberger scaling ( τOL\tau_{OL} , SS-BFGS style), and the unscaled baseline ( τ=1\tau=1 , SS-), validating structure-aware step-size regulation.
MetricAdamSS-eSOAP (Ours)
Storage Complexity (Per Layer)O(dindout)O(d_{\text{in}}d_{\text{out}})O(dindout+din2+dout2)O(d_{\text{in}}d_{\text{out}}+d_{\text{in}}^{2}+d_{\text{out}}^{2})
Peak VRAM Footprint8.2 GB9.2 GB
Time to Target SolutionDNC ( >14.0>14.0 hours)4.1 hours
Scalability (10M+ Parameters)TractableTractable
Table 2: Compute and memory comparison between Adam and SS-eSOAP on 2D Boussinesq, targeting residual 10−510^{-5} on a single consumer GPU.
Appendix figures & tables23 assets

Supplementary material from the paper’s appendix.

Appendix

Method(Approx.) Precond.StorageComputationMini-batchDNN
Natural GradientF−1F^{-1}O(n4)O(n^{4})min⁡(O(PN2),O(NP2))\min(O(PN^{2}),O(NP^{2}))×\times×\times
BFGSH−1H^{-1}O(n2)O(n^{2})O(n4)O(n^{4})×\times\sqrt{}
L-BFGSH−1H^{-1}O(n2)O(n^{2})O(n3)O(n^{3})×\times\sqrt{}
SOAPH−1H^{-1}O(n2)O(n^{2})O(n3)O(n^{3})\sqrt{}\sqrt{}
SS-eSOAPτH−1\tau H^{-1}O(n2)O(n^{2})O(n3)O(n^{3})\sqrt{}\sqrt{}
ShampooHada −1/2H_{\text{ada }}^{-1/2}O(n2)O\left(n^{2}\right)O(n3)O\left(n^{3}\right)\sqrt{}\sqrt{}
Appendix
Table 3: Comparison of optimization methods showing preconditioner types, storage and computational complexity for an n×nn\times n weight matrix, where PP is the parameter count and NN is the batch size of collocation points, and practical compatibility with mini-batches and scalability with large neural networks.
Figure 3: Wave Equation Benchmark. (a) PDE residual loss trajectories (log scale) comparing SS-eSOAP (red) against Adam, Muon, SOAP, and Purifying Shampoo over 100k iterations. SS-eSOAP achieves the lowest final loss, while Muon exhibits significant instability. (b) Visualization of the predicted solution field and the corresponding absolute error map at the final training step, demonstrating high-precision recovery of the dynamics with errors on the magnitude of 10−610^{-6} .
Figure 4: Burgers Equation Benchmark. (a) PDE residual loss trajectories (log scale). SS-eSOAP (red) demonstrates superior convergence speed and stability, reaching a final residual orders of magnitude lower than Adam and Muon. (b) Visualization of the predicted solution and absolute error map. The low error magnitude confirms the optimizer’s ability to accurately resolve the sharp shock interface characteristic of the inviscid limit.
Figure 5: Korteweg-de Vries (KdV) Equation Benchmark. (a) PDE residual loss trajectories (log scale). SS-eSOAP (red) maintains a steep linear convergence rate, reaching a final residual of 10−1010^{-10} , significantly outperforming the baselines. (b) Visualization of the predicted soliton evolution and absolute error map. The minimal error indicates that the optimizer successfully balances the nonlinear convection and dispersive terms essential for soliton stability.
Figure 6: Gray-Scott and Ginzburg-Landau equations. Training loss trajectories of Adam, Muon, SOAP, and SS-eSOAP across different network sizes, with MLP as training backbone.
Figure 7: Lid-driven Cavity equation. Training loss trajectories of Adam, Muon, SOAP, and SS-eSOAP across different network sizes.
Figure 8: Training Dynamics across PDE Benchmarks. We compare the residual loss trajectories (log scale) of SS-eSOAP (red) against baselines: Adam, Muon, standard SOAP, and Purifying SOAP. Across most stiff PDE tasks, SS-eSOAP achieves the lowest final residual, often improving upon the next best method by up to 1 to 2 orders of magnitude while exhibiting fewer loss spikes.
Figure 9: Matrix Optimization Benchmarks. We compare SS-eSOAP against Adam, Muon, and Standard SOAP on (a) quadratic regression, (b) quadratic regression, and (c) low-rank completion. SS-eSOAP consistently achieves faster convergence and lower final residuals, particularly in the quadratic regime where the self-scaling mechanism effectively captures the spectrum.
 Training steps
Optimizer30060090012001500
Adam1.21×10−4±0.081.21{\times}10^{-4}{\pm}0.085.42×10−6±0.055.42{\times}10^{-6}{\pm}0.052.15×10−7±0.022.15{\times}10^{-7}{\pm}0.021.51×10−7±0.011.51{\times}10^{-7}{\pm}0.011.12×10−7±0.011.12{\times}10^{-7}{\pm}0.01
Muon4.55×10−2±0.094.55{\times}10^{-2}{\pm}0.091.12×10−4±0.041.12{\times}10^{-4}{\pm}0.045.51×10−6±0.015.51{\times}10^{-6}{\pm}0.013.24×10−10±0.013.24{\times}10^{-10}{\pm}0.012.55×10−13±0.012.55{\times}10^{-13}{\pm}0.01
SOAP8.12×10−3±0.058.12{\times}10^{-3}{\pm}0.051.24×10−7±0.011.24{\times}10^{-7}{\pm}0.013.45×10−10±0.013.45{\times}10^{-10}{\pm}0.018.52×10−12±0.018.52{\times}10^{-12}{\pm}0.014.15×10−13±0.014.15{\times}10^{-13}{\pm}0.01
P. Shampoo7.55×10−3±0.067.55{\times}10^{-3}{\pm}0.069.81×10−7±0.019.81{\times}10^{-7}{\pm}0.012.15×10−10±0.012.15{\times}10^{-10}{\pm}0.015.41×10−12±0.015.41{\times}10^{-12}{\pm}0.012.22×10−13±0.012.22{\times}10^{-13}{\pm}0.01
SS-eSOAP1.25×10−8±0.02\mathbf{1.25{\times}10^{-8}{\pm}0.02}5.52×10−14±0.01\mathbf{5.52{\times}10^{-14}{\pm}0.01}8.21×10−17±0.01\mathbf{8.21{\times}10^{-17}{\pm}0.01}1.55×10−17±0.01\mathbf{1.55{\times}10^{-17}{\pm}0.01}4.32×10−18±0.01\mathbf{4.32{\times}10^{-18}{\pm}0.01}
Appendix
Table 4: Matrix logistic regression loss trajectory. Mean ±\pm standard error over random seeds.
OptimizerSearch Space
All MethodsLearning Rate η∈{10−4,5×10−4,10−3,5×10−3,10−2}\eta\in\{10^{-4},5\times 10^{-4},10^{-3},5\times 10^{-3},10^{-2}\}
Adamβ1=0.9\beta_{1}=0.9 , β2∈{0.95,0.99,0.999}\beta_{2}\in\{0.95,0.99,0.999\}
MuonMomentum ∈{0.9,0.95}\in\{0.9,0.95\} , Orthogonalization Freq. ∈{1,5,10}\in\{1,5,10\}
SOAPPrecondition Freq. F∈{10,50,100,200}F\in\{10,50,100,200\}
Purifying ShampooUpdate Freq. F∈{10,50,100}F\in\{10,50,100\} , Trigger τtrigger∈{0.1,0.2,0.5}\tau_{trigger}\in\{0.1,0.2,0.5\}
Appendix
Table 5: Hyperparameter grid search spaces for baseline optimizers.
ParameterSymbolValue
Momentumβ1\beta_{1}0.90.9
Second-Moment EMAβ2\beta_{2}0.950.95
Epsilonϵ\epsilon10−810^{-8}
Check IntervalIcheckI_{check}11
Adaptive Trigger Thresholdτtrigger\tau_{trigger}0.20.2
Variance Downscaling (Soft Reset)γ\gammaAdaptive {0.25,0.5,0.75}\{0.25,0.5,0.75\}
Appendix
Table 6: Fixed default hyperparameters for SS-eSOAP used across all experiments.
Figure 10: Loss around triggered updates for γ=1\gamma=1 , hard reset, and adaptive γ<1\gamma<1 .
Figure 11: Comparison with PINN-specific loss-balancing methods. MultiAdam, NTK reweighting, and SS-eSOAP trajectories under a matched benchmark, precision, architecture, and training budget.
Figure 12: Scaling on 1D Burgers. Network width versus seconds per 1,000 steps for Adam, L-BFGS (m=20)(m=20) , SS-BFGS, and SS-eSOAP .
MethodFinal Residual (×10−5)(\times 10^{-5})Relative L2L^{2} Error (×10−3)(\times 10^{-3})Status at W=1024\mathbf{W=1024}
SS-BFGS1.04±0.32\mathbf{1.04\pm 0.32}1.48±0.121.48\pm 0.12OOM
SS-Broyden3.15±0.583.15\pm 0.583.85±0.223.85\pm 0.22OOM
SS-eSOAP8.45±0.258.45\pm 0.251.12±0.05\mathbf{1.12\pm 0.05}∼10×\sim 10\times Adam time
Appendix
Table 7: Explicit self-scaled baselines at the tractable network width W=256W=256 . Results report mean ±\pm SD over three seeds. SS-BFGS and SS-Broyden run out of memory at W=1024W=1024 , while SS-eSOAP runs with approximately 10×10\times the Adam wall-clock time.
PDE BenchmarkMetricAdamL-BFGSNTKMultiAdamMuonSOAPSS-eSOAP
1D BurgersRel. L2L^{2}(4.12±0.82)×10−3(4.12\pm 0.82)\times 10^{-3}(1.05±0.15)×10−3(1.05\pm 0.15)\times 10^{-3}(2.84±0.51)×10−3(2.84\pm 0.51)\times 10^{-3}(1.95±0.35)×10−3(1.95\pm 0.35)\times 10^{-3}(5.67±0.85)×10−4(5.67\pm 0.85)\times 10^{-4}(3.21±0.48)×10−4(3.21\pm 0.48)\times 10^{-4}(8.45±0.34)×10−5\boldsymbol{(8.45\pm 0.34)\times 10^{-5}}
 Rel. H1H^{1}(8.75±1.75)×10−3(8.75\pm 1.75)\times 10^{-3}(3.62±0.54)×10−3(3.62\pm 0.54)\times 10^{-3}(5.11±1.02)×10−3(5.11\pm 1.02)\times 10^{-3}(4.02±0.60)×10−3(4.02\pm 0.60)\times 10^{-3}(1.15±0.17)×10−3(1.15\pm 0.17)\times 10^{-3}(8.90±1.33)×10−4(8.90\pm 1.33)\times 10^{-4}(1.92±0.08)×10−4\boldsymbol{(1.92\pm 0.08)\times 10^{-4}}
2D BoussinesqRel. L2L^{2}(1.25±0.25)×10−2(1.25\pm 0.25)\times 10^{-2}(6.41±0.96)×10−3(6.41\pm 0.96)\times 10^{-3}(9.33±1.86)×10−3(9.33\pm 1.86)\times 10^{-3}(7.15±1.07)×10−3(7.15\pm 1.07)\times 10^{-3}(2.14±0.32)×10−3(2.14\pm 0.32)\times 10^{-3}(1.85±0.27)×10−3(1.85\pm 0.27)\times 10^{-3}(4.12±0.16)×10−4\boldsymbol{(4.12\pm 0.16)\times 10^{-4}}
 Rel. H1H^{1}(3.41±0.68)×10−2(3.41\pm 0.68)\times 10^{-2}(1.12±0.16)×10−2(1.12\pm 0.16)\times 10^{-2}(2.05±0.41)×10−2(2.05\pm 0.41)\times 10^{-2}(1.56±0.23)×10−2(1.56\pm 0.23)\times 10^{-2}(5.88±0.88)×10−3(5.88\pm 0.88)\times 10^{-3}(4.76±0.71)×10−3(4.76\pm 0.71)\times 10^{-3}(9.35±0.37)×10−4\boldsymbol{(9.35\pm 0.37)\times 10^{-4}}
Allen-CahnRel. L2L^{2}(8.90±1.78)×10−3(8.90\pm 1.78)\times 10^{-3}(2.45±0.36)×10−3(2.45\pm 0.36)\times 10^{-3}(5.72±1.14)×10−3(5.72\pm 1.14)\times 10^{-3}(4.11±0.61)×10−3(4.11\pm 0.61)\times 10^{-3}(1.02±0.15)×10−3(1.02\pm 0.15)\times 10^{-3}(9.15±1.37)×10−4(9.15\pm 1.37)\times 10^{-4}(2.78±0.11)×10−4\boldsymbol{(2.78\pm 0.11)\times 10^{-4}}
 Rel. H1H^{1}(1.56±0.31)×10−2(1.56\pm 0.31)\times 10^{-2}(5.89±0.88)×10−3(5.89\pm 0.88)\times 10^{-3}(9.81±1.96)×10−3(9.81\pm 1.96)\times 10^{-3}(8.34±1.25)×10−3(8.34\pm 1.25)\times 10^{-3}(2.75±0.41)×10−3(2.75\pm 0.41)\times 10^{-3}(2.10±0.31)×10−3(2.10\pm 0.31)\times 10^{-3}(6.41±0.25)×10−4\boldsymbol{(6.41\pm 0.25)\times 10^{-4}}
Appendix
Table 8: Relative L2L^{2} and H1H^{1} errors against high-fidelity PDE reference solutions. Values are mean ±\pm standard deviation over three independent seeds. Lower is better, with the best result in each row shown in bold. NTK denotes NTK-based loss reweighting.
IterationOff-Diagonal MassGradient CosinePreconditioning StateTrigger
10,00010{,}000 (post)0.020.020.880.88Stable, alignedNo
12,50012{,}5000.180.180.720.72Minor degradationNo
14,50014{,}5000.420.420.350.35Moderate oscillationNo
15,00015{,}000 (pre)0.65\mathbf{0.65}0.05\mathbf{0.05}Severe oscillation, stale basisYes
15,01015{,}010 (post)0.030.030.850.85Stable, restoredNo
Appendix
Table 9: Diagnostic relation between off-diagonal mass and optimization behavior around one triggered basis update. This diagnostic uses τtrigger=0.5\tau_{\mathrm{trigger}}=0.5 . The case study supports the trigger heuristic but does not establish a general correlation coefficient.
Figure 13: Interaction between directional scaling and the inner optimizer. Gray-Scott trajectories for the original inner Adam update, removal of τk\tau_{k} , and the Momentum-SGD inner update. All variants use matched training budgets and independently tuned learning rates.
OperationStandard SOAPSS-eSOAP (No Trigger)SS-eSOAP (Triggered)
Precond. UpdateO(m2)O(m^{2})O(m2)O(m^{2})O(m2)O(m^{2})
ProjectionsO(m3)O(m^{3})O(m3)O(m^{3})O(m3)O(m^{3})
Eigen. DecompO(m3)O(m^{3}) (fixed interval)0O(m3)O(m^{3})
Stability Op.O(m3)O(m^{3}) (Reprojection)O(m3)O(m^{3}) (Check)O(m2)O(m^{2}) (Downscale)
Self-ScalingN/AO(m2)O(m^{2})O(m2)O(m^{2})
Appendix
Table 10: Complexity Comparison per Iteration (Assuming m=nm=n )
  2D BoussinesqAllen-Cahn
ConfigUpdate Freq.Final LossTime (h)Final LossTime (h)
τ=0.0\tau=0.0 (Ideal)100% (Every Step)1.2×10−71.2\times 10^{-7}14.28.5×10−98.5\times 10^{-9}6.8
τ=0.2\tau=0.2 ( Ours )∼\sim 15% (Adaptive)8.4×10−7\mathbf{8.4\times 10^{-7}}4.19.1×10−9\mathbf{9.1\times 10^{-9}}1.9
τ=0.8\tau=0.8 (Lazy)<< 5% (Rare)3.5×10−43.5\times 10^{-4}3.84.2×10−54.2\times 10^{-5}1.7
τ=∞\tau=\infty (Frozen)0% (Never)1.8×10−21.8\times 10^{-2}3.62.1×10−32.1\times 10^{-3}1.6
Appendix
Table 11: Ablation of the eigenbasis update trigger threshold τtrigger\tau_{\text{trigger}} on the Boussinesq and Allen-Cahn benchmarks. The “Ideal” setting ( τ=0.0\tau=0.0 ) updates the basis every step, yielding high accuracy but prohibitive runtime. The “Lazy” and “Frozen” settings fail to capture evolving curvature. Our proposed threshold ( τ=0.2\tau=0.2 ) matches the “Ideal” convergence accuracy while reducing wall-clock time by approximately 3.5×3.5\times , effectively finding the optimal efficiency-accuracy trade-off.
 2D Boussinesq (Target: 10−510^{-5} )1D Burgers (Target: 10−610^{-6} )
OptimizerTime to Target (h)Peak VRAM (GB)Time to Target (h)Peak VRAM (GB)
AdamDNC ( >> 14.0)8.2DNC ( >> 14.0)8.2
MuonDNC ( >> 14.0)8.5DNC ( >> 14.0)8.5
SOAP11.29.18.19.1
Purifying Shampoo8.59.16.39.1
SS-ESOAP (Ours)4.19.23.49.2
Appendix
Table 12: Wall-clock time and memory efficiency comparison to reach target high-precision residual thresholds on the stiff 2D Boussinesq and 1D Burgers equations. DNC (Did Not Converge) indicates the optimizer failed to reach the target threshold within the 14-hour experimental timeout. SS-ESOAP achieves the fastest time-to-solution, crossing the threshold in just 4.1 hours on Boussinesq, while maintaining a memory footprint strictly comparable to first-order baselines and avoiding the prohibitive O(N2)\mathcal{O}(N^{2}) memory cost of full SS-BFGS.
ProblemK-FACSS-ESOAP (ours)
1D Burgers2.45×10−82.45\times 10^{-8}5.54×10−105.54\times 10^{-10}
2D Boussinesq3.12×10−53.12\times 10^{-5}8.42×10−78.42\times 10^{-7}
Lid-Driven Cavity6.85×10−76.85\times 10^{-7}8.23×10−88.23\times 10^{-8}
Appendix
Table 13: Empirical comparison against the scalable K-FAC, which shares our structural efficiency.
MethodValue
Adam4.21×10−54.21\times 10^{-5}
SOAP2.15×10−82.15\times 10^{-8}
SS-Broyden1.25×10−121.25\times 10^{-12}
SS-ESOAP (ours)3.42×10−103.42\times 10^{-10}
Appendix
Table 14: To validate that our approximation preserves self-scaling benefits, we benchmarked a small-MLP 1D Burgers task where SS-Broyden is tractable. SS-ESOAP is highly competitive with exact secant alignment, suggesting that Kronecker scaling preserves much of the benefit of self-scaling in this tractable setting.

Explore similar work

CardsList

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

    Jul 2, 2026Joseph Webb, Sadok Jerad, Coralia CartisGauss-Newton OptimizationPDE Solving

    Mathematical Institute, University of Oxford, Oxford, United Kingdom

  2. Scale-PINN: Learning Efficient Physics-Informed Neural Networks Through Sequential Correction

    Feb 23, 2026Pao-Hsiung Chiu, Jian Cheng Wong, Chin Chun Ooi +3Neural PDE SolversPhysics-Informed ML

    Institute of High Performance Computing, A*STAR, Singapore. · Centre for Frontier AI Research, A*STAR, Singapore. · School of Mechanical Engineering, Tianjin University, China. +1

  3. Lightweight Geometric Adaptation for Training Physics-Informed Neural Networks

    Apr 16, 2026Kang An, Chenhao Si, Shiqian Ma +1Neural Network OptimizationSecond-Order Optimization

    School of Data Science The Chinese University of Hong Kong, Shenzhen Shenzhen, China