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

AI summary

Overview

Research area: Optimization algorithms for physics-informed neural networks (PINNs), sitting at the intersection of scientific machine learning, numerical PDE solving, and structured second-order optimization.

Technical level: Advanced. The paper assumes familiarity with quasi-Newton methods, Kronecker-factored preconditioning (Shampoo, K-FAC, SOAP), and PINN training objectives.

Scope: The paper introduces SS-eSOAP, a layerwise Kronecker-factored optimizer that adds a scalar secant-energy correction and an adaptive basis update to SOAP-style preconditioning, and evaluates it on matrix regression tasks and eight PDE benchmarks against first-order, structured, PINN-specific, and self-scaled quasi-Newton baselines.

(Note: the paper's title in the metadata renders the method as "SS-ESOAP", while the body consistently writes "SS-eSOAP".)

What This Paper Is About

PINN training objectives are often ill-conditioned: the loss mixes PDE residuals with boundary, initial-condition, and data terms, which produces gradient imbalance, numerical stiffness, and poor conditioning. First-order optimizers such as Adam plateau in high-accuracy regimes, while dense quasi-Newton methods improve local conditioning but are too expensive to store for deep networks. The goal is a scalable optimizer that recovers some of the step-size regulation benefit of self-scaled quasi-Newton methods while keeping memory comparable to Adam.

Key Contributions

1. A directional secant-energy correction for a Kronecker-factored metric, yielding a scalar factor τ_k that aligns the scale of the Kronecker preconditioner with the curvature observed along the most recent parameter displacement.
2. An adaptive basis-update mechanism paired with momentum reprojection and a variance-state transition: when an EShampoo-style off-diagonal criterion triggers a basis change, the momentum state is reprojected into the new coordinates while the coordinatewise second moment is downscaled by a factor γ rather than reprojected.
3. Two theoretical results with stated scope: a one-directional secant-matching result (Theorem 4.2) and a steady-state variance-mismatch bound across basis changes (Theorem 4.1).
4. An empirical comparison against first-order, structured, PINN-specific, and self-scaled quasi-Newton baselines, reporting wall-clock time, memory, multi-seed statistics, and physical error measurements (relative L² and H¹ errors over three seeds on four representative PDEs), plus identification of regimes where self-scaling helps and regimes where it hurts.

Main Findings

• Lowest residual on six of eight PDE benchmarks: SS-eSOAP achieves the lowest final residual on six, including Wave (4.12 × 10⁻⁵ versus 1.18 × 10⁻³ for the best baseline), Burgers (5.54 × 10⁻¹⁰ versus 1.12 × 10⁻⁹), Korteweg-de Vries (3.52 × 10⁻¹⁰ versus 8.12 × 10⁻¹⁰), Boussinesq (8.42 × 10⁻⁷ versus 2.85 × 10⁻⁶), Allen-Cahn (1.85 × 10⁻⁹ versus 3.15 × 10⁻⁹), and Lid-Driven Cavity (8.23 × 10⁻⁸ versus 2.10 × 10⁻⁷).
• Two benchmark exceptions: Purifying Shampoo performs better on Gray-Scott (5.12 × 10⁻⁹ versus 7.46 × 10⁻⁷ for SS-eSOAP), and SOAP performs better on Ginzburg-Landau (9.12 × 10⁻¹² versus 1.75 × 10⁻¹¹). The paper states explicitly that the gains are regime-dependent, not uniform.
• Wall-clock and memory on Boussinesq: SS-eSOAP reaches a residual of 10⁻⁵ in 4.1 hours with 9.2 GB peak VRAM on a single consumer GPU, while Adam does not reach this target within 14 hours (reported as DNC, > 14.0 hours) and uses 8.2 GB peak VRAM. The paper notes this is 2D Boussinesq with a target residual of 10⁻⁵.
• Large gains on stiff problems: On 1D Burgers, SS-eSOAP reaches 5.54 × 10⁻¹⁰, described as an improvement of about 3.9 orders of magnitude over Adam. On 2D Boussinesq it resolves singular blow-up profiles where Muon and Adam stagnate at 2.57 × 10⁻² and 5.12 × 10⁻⁴ respectively.
• Memory scaling: Per-layer storage is O(d_in d_out) for Adam versus O(d_in d_out + d_in² + d_out²) for SS-eSOAP, avoiding the O(d_in² d_out²) cost of dense quasi-Newton updates on a d_out × d_in matrix parameter.
• Self-scaling is applied rarely: The paper observes that τ_k = 1 is chosen for most training iterations, with self-scaling activated only infrequently, yet these occasional adjustments are sufficient for substantial improvements.
• Ablation results: On Allen-Cahn and Boussinesq, reprojecting the second moment produces large transient loss spikes, while a reset avoids the largest spikes but converges to a higher final loss. The Kronecker-adapted τ_k improves stability and convergence relative to unscaled and scalar SS-BFGS-style (Oren-Luenberger) variants.
• Trigger sensitivity: τ_trigger = 0 (updates every step) incurs prohibitive computational cost and τ_trigger ≥ 0.8 (lazy updates) leads to stale curvature estimates and poor convergence, with a broad stable region at τ_trigger ∈ [0.15, 0.3]. A fixed default of τ_trigger = 0.2 and γ ∈ {0.25, 0.5, 0.75} is used across all experiments.
• Residual is not solution accuracy: The paper notes that PDE residual alone does not establish solution accuracy, so three-seed relative L² and H¹ errors are reported on four representative PDEs; the error ranking agrees with the residual ranking on the stiff Burgers and Boussinesq cases, with remaining method-dependent differences in the full table.
• Matrix regression: On a matrix quadratic regression objective of the form ½‖AXB − C‖_F² with (m, n, p, q) = (500, 100, 1000, 250), SS-eSOAP converges faster and reaches lower residuals in the quadratic regime, the setting where the Hessian has an exact Kronecker form.

Methodology in Plain English

The optimizer builds on SOAP, which maintains left and right Kronecker factor matrices (L and R) for each weight matrix, eigendecomposes them to get an eigenbasis, and then runs Adam-style updates inside that rotated coordinate system. SS-eSOAP changes three things.

First, it adds a scalar correction. Standard self-scaling BFGS computes a factor that rescales the Hessian approximation to match curvature along the latest step, but doing this exactly requires storing a dense matrix over all parameters. The authors derive the analogous quantity for a Kronecker metric: a ratio of the trace of the gradient change times the parameter change to the trace of the parameter change transformed by the inverse Kronecker factors. This scalar, clipped at 1, is applied as an inverse square root to the eigenspace update. The theory shows this scalar is the unique solution of a clipped one-dimensional log-secant matching problem, and that when the unclipped solution is active, the scaled Kronecker quadratic form matches the true directional curvature up to exactly the directional secant error.

Second, it decides when to refresh the eigenbasis adaptively instead of on a fixed schedule. It monitors the relative off-diagonal mass of each Kronecker factor in the current basis — a heuristic proxy for how poorly the basis diagonalizes the tracked statistics — and recomputes the basis when the maximum ratio across the left and right factors exceeds a threshold.

Third, when the basis does change, it treats the two optimizer states differently. The momentum is reprojected into the new coordinates so it represents the same update direction. The coordinatewise second moment is not reprojected, because a similarity transformation assumes variance is preserved under rotation, which the authors argue is invalid when the principal directions shift substantially. Instead the second moment is multiplied by γ ∈ {0.25, 0.5, 0.75} chosen from the observed off-diagonal ratio; since V sits in the denominator, γ < 1 raises the immediate effective step size before new gradients refresh the state. The authors describe this transition explicitly as an empirical transient heuristic and state that the theoretical bound does not establish the optimality of γ < 1.

Evaluation covers matrix regression (including convex, strongly convex, and non-convex settings, with logistic regression and low-rank completion in the appendix) plus eight nonlinear evolutionary PDE benchmarks: Wave, Burgers, Allen-Cahn, Boussinesq, Korteweg-de Vries, Gray-Scott, Ginzburg-Landau, and Lid-Driven Cavity. Baselines are Adam, Muon, SOAP, and an adaptive Shampoo variant implemented from the Purifying Shampoo work, with L-BFGS, SS-BFGS, SS-Broyden, K-FAC for PINNs, NTK reweighting, and MultiAdam evaluated where memory and training requirements fit the benchmark.

Why This Matters

Impact on research. The paper targets the accuracy regime that PINN-based scientific discovery depends on — for example the search for finite-time singularities in the three-dimensional Euler and Navier-Stokes equations, where candidate self-similar profiles require low residuals before computer-assisted analysis can proceed. It shows that a structured, memory-efficient optimizer can close much of the gap with dense quasi-Newton methods on stiff problems, and it is candid that the advantage is regime-dependent rather than universal.

Real-world applications (plausible, based on the PDE families studied):

• Fluid dynamics and turbulence modeling, where the paper's Boussinesq and Lid-Driven Cavity benchmarks with singular blow-up profiles are canonical test problems.
• Geophysical and atmospheric flows, where Boussinesq-type approximations are standard.
• Reacting-flow and pattern-formation simulation, represented by the Gray-Scott benchmark (where the authors note a competing optimizer performs better).
• Wave propagation and dispersive systems, represented by the Korteweg-de Vries and Wave benchmarks.

Industry relevance. The headline efficiency result is practical: a fixed-accuracy target on a stiff 2D PDE reached in 4.1 hours with 9.2 GB peak VRAM on a single consumer GPU, versus the baseline optimizer failing to reach the same target in 14 hours. A training cost of roughly one Adam-sized state budget (8.2 GB to 9.2 GB in this comparison) makes curvature-aware PINN training feasible on hardware teams already have, rather than requiring dense quasi-Newton storage.

Future Directions

• Reducing the cubic-layer-width cost of the basis check and eigendecomposition, which currently limits use in wide layers.
• Selecting the variance-state transition (the choice of γ and the trigger threshold) from online diagnostics rather than fixed defaults, since the paper treats the transient heuristic as empirical and does not prove its optimality.
• Transferring the method to smoother physics architectures: all main PDE experiments use MLP or PirateNet backbones, and PINNsFormer, whose reported loss surface is smoother, remains unevaluated.
• Broadening the seed-level evidence: the three-seed study covers four representative PDEs, so broader conclusions require more runs. The paper also leaves open global convergence or sufficient-decrease guarantees for the full nonconvex algorithm.

Target Audience

Researchers and practitioners in scientific machine learning, numerical optimization, and computational physics who train PINNs to high accuracy and are willing to work with structured preconditioners. It is also relevant to optimizer researchers interested in how self-scaling ideas from quasi-Newton methods can be adapted to Kronecker-factored preconditioners, and to engineers who need to weigh wall-clock time against memory footprint when selecting a training method for stiff PDE problems. Readers without background in Kronecker factorization, Shampoo/SOAP, or the BFGS secant condition will find the methodology sections demanding.

Authors’ 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^{-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.