PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2608.29448v1cs.LGcs.AImath.OCstat.ML

TL;DR

High-accuracy PINN training is hindered by ill-conditioned objectives, while dense quasi-Newton methods are expensive and SOAP relies on periodic basis updates. SS-ESOAP combines Kronecker-adapted secant-energy scaling with adaptive basis updates and variance-state downscaling. Across eight PDE benchmarks it achieves the lowest final residual on six, but SOAP-family methods remain better on two benchmarks, so the method is specialized rather than universal.

  • Problem

    PINN objectives can be ill-conditioned in high-accuracy regimes, while dense quasi-Newton state is expensive and SOAP relies on periodic basis updates.

  • Method

    SS-ESOAP augments SOAP-style preconditioning with a Kronecker-adapted scalar secant correction, adaptive basis updates, and variance-state downscaling.

  • Results

    Across eight PDE benchmarks, SS-ESOAP attains the lowest final residual on six; on Boussinesq it reaches 10^-5 in 4.1 hours, while Adam misses the target within 14 hours.

  • Takeaways & Limitations

    SS-ESOAP is a scalable option for stiff, high-accuracy PINN training rather than a uniform replacement for existing optimizers.

  • Takeaways & Limitations

    The method lacks a global-convergence or sufficient-decrease proof, incurs cubic basis costs, and remains unevaluated on transformer-based physics models.

Abstract

from arXiv · show

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^{-5}$ in 4.1 hours with 9.2 GB peak VRAM, while Adam does not reach this target within 14 hours. Three-seed $L^2$ and $H^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.

1 Introduction

PINN training becomes difficult at high accuracy because composite objectives can be ill-conditioned, while dense curvature methods are costly for deep networks. SS-ESOAP adds Kronecker-compatible self-scaling and adaptive state updates to SOAP-style preconditioning, improving several stiff PDE benchmarks without uniform dominance.

  • PINN objectives combine PDE, boundary, and initial-condition terms, producing gradient imbalance, numerical stiffness, and ill-conditioned loss landscapes.
  • O(d_out^2 d_in^2) storage makes exact self-scaled quasi-Newton updates impractical for deep networks.
  • Kronecker-factored methods scale curvature statistics through tensor products of smaller matrices, but still face limitations in non-stationary regimes.
  • SS-ESOAP matches Kronecker metrics to recent secant energy and triggers basis updates using an off-diagonal criterion, followed by momentum reprojection and variance-state downscaling.
  • The scalar secant calculation costs O(d_in d_out), while basis checks and eigendecompositions retain cubic layer-width costs.

2 Background and Related Work

PINNs solve PDE-constrained learning with composite residual objectives, motivating structured and geometry-aware optimizers. Related methods trade dense curvature fidelity against scalable Kronecker structure, while SOAP’s fixed basis refresh and state reprojection create stability concerns addressed by adaptive alternatives.

  • PINNs represent PDE solutions with neural networks and train them by minimizing objectives containing residual, boundary, initial-condition, and observational-data terms.
  • Loss-balancing and geometry-aware methods address sharply different gradient scales and directions through reweighting, natural gradients, or structured preconditioning.
  • SS-ESOAP combines layerwise Kronecker statistics with a directional secant correction rather than reweighting loss components or solving a global parameter-space or residual-space system.
  • Structured Matrix Preconditioning: SOAP applies Adam in Kronecker-factor eigenbases, while Shampoo and Muon provide related structured matrix-scaling alternatives.
  • Limitations of SOAP, Purifying Shampoo, and Adaptive Eigenvalue Correction: Fixed-interval basis updates do not distinguish training stages or layer dynamics, and second-moment reprojection can produce unstable step sizes when principal directions shift.
  • Limitations of SOAP, Purifying Shampoo, and Adaptive Eigenvalue Correction: Purifying Shampoo separates eigenbasis and eigenvalue errors, motivating adaptive stopping criteria that SS-ESOAP adopts for its basis diagnostic.

3 Methodology: SS-ESOAP

SS-ESOAP augments SOAP-style Kronecker preconditioning with a self-scaling curvature correction and adaptive basis management. Its scalar factor aligns preconditioner scale with recent displacement curvature, while state transitions stabilize changes in eigencoordinates.

  • SS-ESOAP integrates a self-scaling factor τ_k with a purifying-style adaptive update strategy.
  • Self-Scaling Curvature Correction: The correction uses vectorized parameter and gradient displacements to scale a Kronecker inverse-Hessian approximation through trace-based secant information.
  • Self-Scaling Curvature Correction: τ_k = 1 is selected for most iterations, with infrequent self-scaling adjustments aligning preconditioner scale to curvature along the latest displacement.
  • Adaptive Basis Updates and Variance-State Transition: An off-diagonal diagnostic measures how poorly the current basis diagonalizes each tracked Kronecker factor and triggers recomputation when the maximum ratio exceeds τ_trigger.
  • Adaptive Basis Updates and Variance-State Transition: The basis diagnostic is a heuristic proxy for preconditioning degradation rather than an exact measure of Hessian staleness.
  • Adaptive Basis Updates and Variance-State Transition: After a basis change, SS-ESOAP reprojects momentum but retains coordinatewise second-moment shape and downscales it by γ ∈ {0.25, 0.5, 0.75}.

4 Theoretical Analysis

The analysis establishes directional secant-energy matching for the Kronecker metric and bounds variance-state mismatch after basis changes.

  • The variance-state transition compares a scalar soft reset with the oracle diagonal second moment in the updated basis.
  • The variance mismatch is bounded by reset magnitude, basis rotation, and change in the gradient second-moment matrix.
  • Theorem 4.2 defines a Kronecker secant-energy correction using the displacement, gradient difference, and positive-definite Kronecker metric.
  • The scaling factor is the unique solution of a clipped one-dimensional log-secant matching problem.
  • When unclipped, the scaled Kronecker quadratic form differs from true directional curvature by exactly the directional secant error.

5 Numerical Experiments

Experiments evaluate SS-ESOAP on controlled matrix problems and diverse PDE benchmarks, including ablations of its adaptive basis and scaling mechanisms. It achieves the lowest residual on six of eight PINN benchmarks, with gains concentrated in stiff problems but not uniform across PDEs.

  • SS-ESOAP is evaluated on matrix regression, low-rank completion, and nonlinear evolutionary PDE tasks spanning structured curvature, stochasticity, and stiffness.
  • Matrix Regression Tasks: SS-ESOAP demonstrates faster convergence and lower residuals on the matrix quadratic task, where its self-scaling mechanism captures the spectrum.
  • Matrix Regression Tasks: SS-ESOAP tests exact Kronecker-curvature preconditioning on a strongly convex matrix quadratic regression objective.
  • PINN Benchmarks: SS-ESOAP achieves the lowest final residual on six of eight PINN benchmarks, including 5.54 × 10^-10 versus 1.12 × 10^-9 on Burgers.
  • PINN Benchmarks: 5.54 × 10^-10 is reached on Burgers, about 3.9 orders of magnitude better than Adam; Boussinesq reaches 8.42 × 10^-7 while Muon and Adam stagnate at 2.57 × 10^-2 and 5.12 × 10^-4.
  • PINN Benchmarks: The residual ranking agrees with relative L2 and H1 error rankings on stiff Burgers and Boussinesq cases across three seeds.
  • Component Analysis and Ablation Studies: Reprojection causes large transient loss spikes, whereas resetting avoids the largest spikes but converges to a higher final loss.
  • Component Analysis and Ablation Studies: The Kronecker-adapted scaling factor improves stability and convergence relative to unscaled and scalar SS-BFGS-style variants.

6 Limitations

The analysis does not establish global convergence, and the adaptive trigger and basis computations impose scalability limits. Performance is also not uniformly better than competing optimizers.

  • The analysis does not prove global convergence or sufficient decrease for the full nonconvex algorithm.
  • The basis trigger is heuristic, while its O(d^3) check and eigendecomposition costs limit use in wide layers.
  • Dense SS-BFGS or SS-Broyden may offer better directions on networks small enough to store their state, and residual-space natural-gradient methods remain alternatives.

7 Conclusion

SS-ESOAP combines directional secant-energy matching, adaptive basis updates, and variance-state transitions for high-accuracy PINN optimization. It performs best on most tested PDEs but is not uniformly superior, leaving computational and transfer questions open.

  • SS-ESOAP combines directional secant-energy matching with adaptive basis updates and a variance-state transition.
  • SS-ESOAP attains the lowest final residual on six of eight PDE benchmarks and improves wall-clock time to a fixed target on Boussinesq.
  • Three-seed physical errors on four PDEs support the solution quality of the low-residual fits.
  • Gray-Scott and Ginzburg-Landau show that the added scaling is not uniformly beneficial.
  • Open questions include reducing cubic basis-check costs, selecting state transitions from online diagnostics, and transferring the method to smoother physics architectures.

A.2 Proof of Theorem 4.2

The analysis establishes the directional behavior of the scalar secant correction and situates SS-ESOAP as a curvature-corrected, adaptively updated SOAP-style method. Its implementation uses eigenspace operations and state updates designed for efficient Kronecker preconditioning.

  • A.2 Proof of Theorem 4.2: When the unclipped correction is active, scaled Kronecker curvature matches true directional curvature exactly under an exact secant relation.
  • A.2 Proof of Theorem 4.2: Under an approximate secant relation, the curvature mismatch equals the directional secant error; if the unconstrained scale exceeds one, τk is clipped to 1.
  • B.1 Relation to Idealized and Warm-Started Shampoo: SS-ESOAP triggers basis recomputation when off-diagonal mass exceeds a threshold, refining the eigenbasis only when curvature changes significantly.
  • B.1 Relation to Idealized and Warm-Started Shampoo: The scalar τk corrects eigenvalues globally so the inverse-Hessian approximation matches the secant condition along the latest displacement.
  • B.1 Relation to Idealized and Warm-Started Shampoo: After basis changes, SS-ESOAP reprojections momentum but downscales the coordinatewise second moment, avoiding full variance reprojection.
  • A.2 Proof of Theorem 4.2: Rotated eigenspace quantities permit computing the correction with matrix-vector products and elementwise operations without explicitly forming inverse matrices.

C.3 Full Details of Matrix Optimization Tasks

The matrix experiments test SS-ESOAP in structured convex, stochastic, and nonconvex settings, while additional diagnostics examine basis transitions and PINN optimizer comparisons. Results generally favor SS-ESOAP but expose benchmark- and transition-dependent behavior.

  • Quadratic Regression: The quadratic regression Hessian has Kronecker structure, allowing SS-ESOAP to preserve residual curvature information while isolating spectral skew from A and B.
  • Logistic Regression: SS-ESOAP maintains consistent descent in stochastic logistic regression as sigmoid predictions saturate, unlike Adam’s deceleration when gradients vanish.
  • Low-rank Matrix Completion: SS-ESOAP achieves the fastest descent and lowest final residual in low-rank matrix completion, while Muon plateaus after rapid initial progress.
  • PINN Optimizer Comparisons: 4.0 × 10^-5 is the SS-ESOAP residual in the matched PINN comparison, below MultiAdam near 3.0 × 10^-4 and NTK reweighting near 7.0 × 10^-5.
  • Dense and limited-memory quasi-Newton baselines: At width W = 256, SS-BFGS obtains the lowest residual while SS-ESOAP obtains the lowest relative L2 error; at W = 1024, SS-ESOAP incurs roughly 10× Adam’s wall-clock cost.
  • Basis-update diagnostics: The basis-update diagnostic rises from off-diagonal mass 0.02 to 0.65 as cosine similarity falls from 0.88 to 0.05, then both recover after updating.

E Computational Complexity Analysis

SS-ESOAP adds inexpensive secant-scaling and downscaling operations but incurs cubic basis checks and eigendecompositions. This overhead is more suitable for typical PINN widths than for very large models, while adaptive skipping can reduce decompositions.

  • Per-iteration costs: O(mn) secant-scaling operations are negligible relative to the standard O(m^2n) projection cost.
  • Basis-management overhead: O(m^3) off-diagonal checks and triggered eigendecompositions dominate the added basis-management cost.
  • Adaptive scheduling: Adaptive triggering can skip eigendecompositions when off-diagonal mass remains small, potentially reducing amortized cost relative to fixed schedules.
  • State-transition cost: O(mn) downscaling replaces O(m^3) variance reprojection when the basis changes.
  • Model-width dependence: For PINN layers with m ≤ 256, the cubic term is described as negligible, whereas for models with dmodel = 4096 it is prohibitive.

F Computational Efficiency Analysis

SS-ESOAP uses Kronecker structure to reduce curvature-related costs while preserving high-precision convergence, especially in stiff regimes. Its adaptive basis threshold balances accuracy and runtime, but self-scaling can underperform on smoother benchmarks.

  • Structural efficiency: O(d_out^2) per layer replaces O((d_in d_out)^2) curvature-estimation cost through Kronecker structure.This reduces the cost relative to dense self-scaled quasi-Newton updates.
  • Convergence speed: 4.1 hours is the time SS-ESOAP takes to cross the 10^-5 residual threshold on stiff 2D Boussinesq, while Adam stagnates after 10+ hours.In easier regimes, Adam can reduce loss rapidly and match SS-ESOAP’s time to an intermediate loss.
  • Scalability and compute trade-off: Over 400 TB of VRAM would be required for full SS-BFGS on the largest configuration, whereas SS-ESOAP uses megabytes of curvature storage per layer.The comparison concerns a PirateNet configuration with approximately 10^7 parameters.
  • Ablation: 3.5× lower wall-clock time is achieved by the proposed τ = 0.2 basis-update threshold versus updating the eigenbasis every step.The threshold matches the Ideal setting’s convergence accuracy while avoiding the Ideal setting’s prohibitive runtime.
  • Limitations and trade-offs: SOAP achieves marginally lower final residuals than SS-ESOAP on Gray-Scott and Ginzburg-Landau.The paper attributes this trade-off to the conservative damping imposed by the self-scaling factor τk.
  • Limitations and trade-offs: Ultra-stiff Burgers and Boussinesq problems require conservative scaling, whereas Gray-Scott and Ginzburg-Landau have smoother, highly oscillatory landscapes.The paper links these differing optimization landscapes to the regime-dependent benefit of self-scaling.

J Relation to Natural-Gradient Methods

SS-ESOAP differs from natural-gradient and Gauss-Newton approaches by using layerwise Kronecker statistics and a directional secant correction rather than global function-space or residual-space solves. The paper positions it as structurally efficient but does not claim superiority without matched experiments.

  • Comparison: Natural-gradient and Gauss-Newton methods use function-space geometry, while K-FAC approximates that geometry with layerwise Kronecker blocks.Related approaches include energy natural gradients, Gauss-Newton natural gradients, and ANaGRAM.
  • Comparison: 12.8 million parameters have been demonstrated for PINNs using dual formulations that move the solve from parameter space to residual space.This provides a scale reference for related natural-gradient-style methods.
  • SS-ESOAP design: SS-ESOAP maintains layerwise gradient second moments, performs no global residual-space solve, and adds a directional secant-energy correction.Its curvature storage scales with layerwise Kronecker factors, while basis checks and eigendecompositions remain cubic in layer width.
  • Scope: The paper does not claim lower wall-clock cost or higher accuracy than modern natural-gradient methods without a matched experiment.Dense self-scaled methods remain stronger references when their state and line-search costs are tractable.
  • Secant-scaling validation: SS-ESOAP is highly competitive with exact secant alignment on a small-MLP 1D Burgers task where SS-Broyden is tractable.The comparison suggests that Kronecker scaling preserves much of self-scaling’s benefit in this setting.
Loading 2608.29448v1…