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

arXiv cs.LG Papers

Summary

SS-eSOAP introduces a self-scaled adaptive preconditioning method for physics-informed neural networks, improving optimization by reducing residuals and achieving higher accuracy on various PDE benchmarks compared to baseline methods.

arXiv:2608.29448v1 Announce Type: new 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.
Original Article
View Cached Full Text

Cached at: 09/01/26, 01:15 PM

# SS-eSOAP: Self-Scaled Adaptive Preconditioning for Physics-Informed Learning
Source: [https://arxiv.org/html/2608.29448](https://arxiv.org/html/2608.29448)
Mads ToftrupAffiliation:Aarhus UniversitySebastian LoeschckeAffiliation:University of CopenhagenYixuan WangAffiliation:California Institute of TechnologyAnima AnandkumarAffiliation:California Institute of Technology

###### 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 introduceSS\-eSOAP, 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,SS\-eSOAPattains the lowest final residual on six, including Burgers and Boussinesq, while SOAP\-family baselines perform better on Gray\-Scott and Ginzburg\-Landau\. On Boussinesq,SS\-eSOAPreaches a residual of10−510^\{\-5\}in 4\.1 hours with 9\.2 GB peak VRAM, while Adam does not reach this target within 14 hours\. Three\-seedL2L^\{2\}andH1H^\{1\}errors on four representative PDEs support the link between lower residuals and improved solution accuracy\. These results positionSS\-eSOAPas a scalable option for stiff, high\-accuracy physics\-informed training, rather than a uniform replacement for existing optimizers\.

## 1Introduction

High\-precision optimization is central to scientific machine learning\. One example is the search for finite\-time singularities in the three\-dimensional Euler and Navier\-Stokes equations\. PINNs have been used to reconstruct candidate self\-similar profiles in unbounded domains\[[35](https://arxiv.org/html/2608.29448#bib.bib29),[34](https://arxiv.org/html/2608.29448#bib.bib28),[33](https://arxiv.org/html/2608.29448#bib.bib26)\]\. Such candidates require low residuals before computer\-assisted analysis\. PINN objectives are often ill\-conditioned near this accuracy regime, where standard first\-order training may plateau\[[31](https://arxiv.org/html/2608.29448#bib.bib12),[17](https://arxiv.org/html/2608.29448#bib.bib14),[25](https://arxiv.org/html/2608.29448#bib.bib13),[29](https://arxiv.org/html/2608.29448#bib.bib31),[36](https://arxiv.org/html/2608.29448#bib.bib7)\]\.

More broadly, PINN optimization presents challenges that differ from standard supervised learning\. PINN losses combine PDE residuals with boundary and initial\-condition terms, which can induce gradient imbalance, numerical stiffness, and ill\-conditioned loss landscapes\[[31](https://arxiv.org/html/2608.29448#bib.bib12),[17](https://arxiv.org/html/2608.29448#bib.bib14),[25](https://arxiv.org/html/2608.29448#bib.bib13)\]\. As a result, standard first\-order optimizers such as Adam often converge slowly or plateau in high\-accuracy regimes, motivating curvature\-aware methods for PINN training\[[27](https://arxiv.org/html/2608.29448#bib.bib15)\]\.

Second\-order methods use curvature information via Hessian\-approximation updates, enabling more efficient descent directions than first\-order methods in many non\-convex problems\. Quasi\-Newton methods such as BFGS\[[5](https://arxiv.org/html/2608.29448#bib.bib9)\]and L\-BFGS\[[18](https://arxiv.org/html/2608.29448#bib.bib22)\]build these approximations from gradient differences, avoiding explicit second derivatives while retaining favorable local convergence\. Self\-scaled variants, such asSS\-BFGS, rescale the Hessian approximation to better match the local curvature spectrum\[[23](https://arxiv.org/html/2608.29448#bib.bib19),[2](https://arxiv.org/html/2608.29448#bib.bib21),[22](https://arxiv.org/html/2608.29448#bib.bib20),[27](https://arxiv.org/html/2608.29448#bib.bib15)\]\. However, these methods operate on the flattened vector, requiringO⁡\(din2​dout2\)O\(d\_\{\\mathrm\{in\}\}^\{2\}d\_\{\\mathrm\{out\}\}^\{2\}\)for a singledout×dind\_\{\\mathrm\{out\}\}\\times d\_\{\\mathrm\{in\}\}matrix parameter, making exact self\-scaled quasi\-Newton updates impractical for deep networks\.

Kronecker\-factored preconditioners provide a scalable alternative\. Methods such as K\-FAC\[[20](https://arxiv.org/html/2608.29448#bib.bib23)\], Shampoo\[[8](https://arxiv.org/html/2608.29448#bib.bib24)\], and SOAP\[[28](https://arxiv.org/html/2608.29448#bib.bib25)\]approximate curvature or gradient covariance through tensor products of smaller matrices, enabling structured preconditioning without maintaining a dense Hessian approximation\. Rather than enforcing a quasi\-Newton secant condition, these methods use second\-order statistics such as empirical Fisher information or uncentered gradient covariance to rescale the gradient\. However, existing Kronecker\-factored methods still face limitations in non\-stationary regimes\. Standard Shampoo and SOAP typically update the preconditioner eigenspace at fixed intervals, which may not match the rate at which curvature statistics evolve during PINN training\. Moreover, accumulated second\-moment statistics are often reprojected after basis updates, which can be unstable when the basis rotates substantially\[[7](https://arxiv.org/html/2608.29448#bib.bib8)\]\.

In this work, we proposeSS\-eSOAP, a layerwise Kronecker\-factored optimizer for high\-accuracy PINN training\. The method adds two operations to SOAP\-style preconditioning\. First, a scalar factor matches the Kronecker metric to the secant energy observed along the latest parameter displacement\. Second, an EShampoo\-style off\-diagonal criterion triggers a basis update\[[7](https://arxiv.org/html/2608.29448#bib.bib8)\]\. At each triggered update,SS\-eSOAPreprojects the momentum state and downscales the coordinatewise second moment instead of reprojecting that second moment\.

For adout×dind\_\{\\mathrm\{out\}\}\\times d\_\{\\mathrm\{in\}\}matrix parameter,SS\-eSOAPstoresO⁡\(din2\+dout2\)O\(d\_\{\\mathrm\{in\}\}^\{2\}\+d\_\{\\mathrm\{out\}\}^\{2\}\)Kronecker factors in addition to Adam\-style states\. The scalar secant calculation costsO⁡\(din​dout\)O\(d\_\{\\mathrm\{in\}\}d\_\{\\mathrm\{out\}\}\)after the required eigenspace quantities are available, orO⁡\(din​dout\)O\(d\_\{\\mathrm\{in\}\}d\_\{\\mathrm\{out\}\}\)when counted over all matrix entries\. The basis check and eigendecomposition retain cubic layer\-width costs, analyzed in Appendix[F](https://arxiv.org/html/2608.29448#A6)\.

Across eight PDE benchmarks,SS\-eSOAPrecords the lowest final residual on six\. On Boussinesq, it reaches10−510^\{\-5\}in 4\.1 hours, while Adam does not reach the same target within a 14\-hour run\. The gains are regime\-dependent\. Purifying Shampoo performs better on Gray\-Scott, and SOAP performs better on Ginzburg\-Landau\. We therefore studySS\-eSOAPas a specialized optimizer for stiff, high\-accuracy regimes\.

Our contributions are:

- •We derive a directional secant\-energy correction for a Kronecker\-factored metric\.
- •We pair an adaptive basis trigger with momentum reprojection and a variance\-state transition designed for abrupt basis changes\.
- •We state the scope of two theoretical results: a one\-directional secant\-matching result and a steady\-state variance\-mismatch bound\.
- •We compare against first\-order, structured, PINN\-specific, and self\-scaled quasi\-Newton baselines, with wall\-clock, memory, multi\-seed, and physical\-error measurements\.
- •We identify regimes where self\-scaling helps and regimes where it hurts\.

## 2Background and Related Work

### 2\.1Physics\-Informed Neural Networks \(PINNs\)

Physics\-informed neural networks \(PINNs\)\[[24](https://arxiv.org/html/2608.29448#bib.bib10),[14](https://arxiv.org/html/2608.29448#bib.bib11)\]solve PDE\-constrained learning problems by representing the solution field with a neural network and enforcing the governing equations through differentiable residual losses\. Formally, given a PDE of the formℱ​\[u\]​\(x,t\)=0\\mathcal\{F\}\[u\]\(x,t\)=0over a domainΩ\\Omega, a PINN approximates the solutionu⁡\(x,t\)u\(x,t\)by a neural networkuθ​\(x,t\)u\_\{\\theta\}\(x,t\)\. Automatic differentiation evaluatesℱ⁡\[uθ\]\\mathcal\{F\}\[u\_\{\\theta\}\], and the network is trained by minimizing a composite objective

ℒ⁡\(θ\)=λr​ℒr​\(θ\)\+λb​ℒb​\(θ\)\+λd​ℒd​\(θ\),\\displaystyle\\mathcal\{L\}\(\\theta\)=\\lambda\_\{r\}\\mathcal\{L\}\_\{r\}\(\\theta\)\+\\lambda\_\{b\}\\mathcal\{L\}\_\{b\}\(\\theta\)\+\\lambda\_\{d\}\\mathcal\{L\}\_\{d\}\(\\theta\),whereℒr\\mathcal\{L\}\_\{r\}penalizes the PDE residual at collocation points,ℒb\\mathcal\{L\}\_\{b\}enforces boundary or initial conditions, andℒd\\mathcal\{L\}\_\{d\}incorporates observational data when available\.

### 2\.2Optimization for Physics\-Informed Learning

PINN objectives combine residual, boundary, initial\-condition, and data terms\. Their gradients may differ sharply in scale and direction\[[31](https://arxiv.org/html/2608.29448#bib.bib12),[32](https://arxiv.org/html/2608.29448#bib.bib6),[17](https://arxiv.org/html/2608.29448#bib.bib14)\]\. Loss\-balancing methods address this issue through learning\-rate annealing, NTK\-based reweighting, or parameterwise normalization\[[31](https://arxiv.org/html/2608.29448#bib.bib12),[32](https://arxiv.org/html/2608.29448#bib.bib6),[37](https://arxiv.org/html/2608.29448#bib.bib5)\]\. Other work targets the geometry of the full objective\. Examples include NysNewton\-CG\[[25](https://arxiv.org/html/2608.29448#bib.bib13)\], PDE preconditioning\[[19](https://arxiv.org/html/2608.29448#bib.bib4)\], energy and Gauss\-Newton natural gradients\[[21](https://arxiv.org/html/2608.29448#bib.bib36),[11](https://arxiv.org/html/2608.29448#bib.bib37)\], ANaGRAM\[[26](https://arxiv.org/html/2608.29448#bib.bib32)\], dual natural gradients\[[12](https://arxiv.org/html/2608.29448#bib.bib33)\], and K\-FAC for PINNs\[[6](https://arxiv.org/html/2608.29448#bib.bib35)\]\. Recent work also studies gradient alignment under SOAP\[[29](https://arxiv.org/html/2608.29448#bib.bib31)\], self\-scaled BFGS and Broyden methods\[[16](https://arxiv.org/html/2608.29448#bib.bib3),[10](https://arxiv.org/html/2608.29448#bib.bib2)\]\.

SS\-eSOAPtargets a different point in this design space\. It does not reweight the PINN loss components or solve a global parameter\-space or residual\-space system\. It combines layerwise Kronecker statistics with a directional secant correction\. Dense self\-scaled methods remain stronger references when their state and line\-search costs are tractable\.

### 2\.3Structured Matrix Preconditioning

Shampoo approximates full\-matrix AdaGrad with Kronecker factors\[[8](https://arxiv.org/html/2608.29448#bib.bib24)\]\. For a matrix parameterWt∈ℝm×nW\_\{t\}\\in\\mathbb\{R\}^\{m\\times n\}and gradientGtG\_\{t\}, it maintains

Lt=β2​Lt−1\+\(1−β2\)​Gt​Gt⊤,Rt=β2​Rt−1\+\(1−β2\)​Gt⊤​Gt\.L\_\{t\}=\\beta\_\{2\}L\_\{t\-1\}\+\(1\-\\beta\_\{2\}\)G\_\{t\}G\_\{t\}^\{\\top\},\\qquad R\_\{t\}=\\beta\_\{2\}R\_\{t\-1\}\+\(1\-\\beta\_\{2\}\)G\_\{t\}^\{\\top\}G\_\{t\}\.\(1\)SOAP applies Adam in the eigenbasis of these factors\[[28](https://arxiv.org/html/2608.29448#bib.bib25)\]\. Muon instead orthogonalizes matrix updates and is related to a zero\-decay Shampoo limit\[[13](https://arxiv.org/html/2608.29448#bib.bib27),[4](https://arxiv.org/html/2608.29448#bib.bib38)\]\.

SOAP commonly refreshes its basis on a fixed schedule\. Purifying Shampoo separates eigenvalue and eigenbasis errors and develops an adaptive stopping criterion for warm\-started basis updates\[[7](https://arxiv.org/html/2608.29448#bib.bib8)\]\.SS\-eSOAPadopts this off\-diagonal basis diagnostic\. Its new elements are the Kronecker\-adapted secant correction and the second\-moment transition used after a triggered basis change\.

Muon\[[13](https://arxiv.org/html/2608.29448#bib.bib27)\]is a recent optimization algorithm designed to approximate second\-order scaling for large\-scale neural network training at a cost comparable to first\-order methods\. The method uses Newton\-Schulz iterations to orthogonalize the gradient asG′=U​V⊤G^\{\\prime\}=UV^\{\\top\}, whereG=U​Σ​V⊤G=U\\Sigma V^\{\\top\}is the SVD of the matrix gradient of a linear layer\. Additionally, it can be seen as a special case of Shampoo withβ2=0\\beta\_\{2\}=0as\(GG⊤\)−1/4G\(G⊤G\)−1/4=UVT\(GG^\{\\top\}\)^\{\-1/4\}G\(G^\{\\top\}G\)^\{\-1/4\}=UV^\{T\}\[[4](https://arxiv.org/html/2608.29448#bib.bib38)\]\.

### 2\.4Limitations of SOAP, Purifying Shampoo, and Adaptive Eigenvalue Correction

Standard implementations of SOAP update the eigenbasis\(QL,QR\)\(Q\_\{L\},Q\_\{R\}\)at fixed intervals\. This static schedule fails to distinguish between different stages of training or layers with varying dynamics\. Furthermore, when the basis is updated, standard approaches reproject the accumulated second\-moment estimatorVkV\_\{k\}\. We argue that reprojectingVkV\_\{k\}via similarity transformationVkn​e​w=\(QLn​e​w\)T​QL​Vk​QRT​QRn​e​wV\_\{k\}^\{new\}=\(Q\_\{L\}^\{new\}\)^\{T\}Q\_\{L\}V\_\{k\}Q\_\{R\}^\{T\}Q\_\{R\}^\{new\}assumes variance preservation under rotation, which is invalid when the principal directions of the Hessian shift significantly, leading to unstable step sizes\.

Recent analysis of the Shampoo algorithm has identified that its practical success relies heavily on two heuristics:stale preconditioning\(updating the eigenbasis at fixed, infrequent intervals\) andlearning rate grafting\(forcing the update magnitude to match Adam’s to mitigate spectral errors\)\.\[[7](https://arxiv.org/html/2608.29448#bib.bib8)\]formalized these issues by decomposing the preconditioner update into its eigenvector \(basis\) and eigenvalue \(curvature\) components, termed the “Purifying” framework\.

### 2\.5Self\-Scaling in Quasi\-Newton Iteration Algorithms

A general class of quasi\-Newton iteration algorithms can be cast under the self\-scaled Broyden formula\[[2](https://arxiv.org/html/2608.29448#bib.bib21)\]\. If we define the auxiliary variables

sk\\displaystyle\{s\}\_\{k\}=Θk\+1−Θk,yk=∇J\(Θk\+1\)−∇J\(Θk\),vk=yk⋅Hk​yk\[skyk⋅sk−Hk​ykyk⋅Hk​yk\],\\displaystyle=\{\\Theta\}\_\{k\+1\}\-\{\\Theta\}\_\{k\},\\quad\{y\}\_\{k\}=\\nabla J\(\{\\Theta\}\_\{k\+1\}\)\-\\nabla J\(\{\\Theta\}\_\{k\}\),\\quad\{v\}\_\{k\}=\\sqrt\{\{y\}\_\{k\}\\cdot H\_\{k\}\{y\}\_\{k\}\}\\left\[\\frac\{\{s\}\_\{k\}\}\{\{y\}\_\{k\}\\cdot\{s\}\_\{k\}\}\-\\frac\{H\_\{k\}\{y\}\_\{k\}\}\{\{y\}\_\{k\}\\cdot H\_\{k\}\{y\}\_\{k\}\}\\right\],then, the next approximation of the inverse Hessian matrix at each iteration can be calculated as in\[[3](https://arxiv.org/html/2608.29448#bib.bib17),[1](https://arxiv.org/html/2608.29448#bib.bib18)\]by

Hk\+1=1τk​\[Hk−Hk​yk⊗Hk​ykyk⋅Hk​yk\+ϕk​vk⊗vk\]\+sk⊗skyk⋅sk,H\_\{k\+1\}=\\frac\{1\}\{\\tau\_\{k\}\}\\left\[H\_\{k\}\-\\frac\{H\_\{k\}\{y\}\_\{k\}\\otimes H\_\{k\}\{y\}\_\{k\}\}\{\{y\}\_\{k\}\\cdot H\_\{k\}\{y\}\_\{k\}\}\+\\phi\_\{k\}\{v\}\_\{k\}\\otimes\{v\}\_\{k\}\\right\]\+\\frac\{\{s\}\_\{k\}\\otimes\{s\}\_\{k\}\}\{\{y\}\_\{k\}\\cdot\{s\}\_\{k\}\},where we define

τk=min⁡\{1,yk⋅skαk​sk⋅Hk−1​sk\}\.\\displaystyle\\tau\_\{k\}=\\min\\left\\\{1,\\frac\{\{y\}\_\{k\}\\cdot\{s\}\_\{k\}\}\{\\alpha\_\{k\}\{s\}\_\{k\}\\cdot H\_\{k\}^\{\-1\}s\_\{k\}\}\\right\\\}\.\(2\)Here,⊗\\otimesdenotes the tensor product of two vectors andτk,ϕk\\tau\_\{k\},\\phi\_\{k\}are respectively the scaling and the updating parameters, which often change between iterations\. Forτk=1\\tau\_\{k\}=1andϕk=1\\phi\_\{k\}=1, one recovers the standard BFGS algorithm\. To efficiently compute the scaling parameterτk\\tau\_\{k\},\[[27](https://arxiv.org/html/2608.29448#bib.bib15)\]setHk−1sk=−αk∇J\(Θk\)H\_\{k\}^\{\-1\}\{s\}\_\{k\}=\-\\alpha\_\{k\}\\nabla J\(\{\\Theta\}\_\{k\}\)so the explicit dependence ofτk\\tau\_\{k\}onHk−1H\_\{k\}^\{\-1\}disappears\. This is feasible since the step lengthαk\\alpha\_\{k\}and the gradient∇J​\(Θk\)\\nabla J\\left\(\\Theta\_\{k\}\\right\)are available at iterationkk\. Thus, the calculation ofτk\\tau\_\{k\}only involves vector multiplications inO⁡\(n\)O\(n\)\.

The original scaling scheme of\[[23](https://arxiv.org/html/2608.29448#bib.bib19)\]isτk\(1\)=−yk⋅skαk​sk⋅∇J​\(Θk\)\\tau\_\{k\}^\{\(1\)\}=\\frac\{\-\{y\}\_\{k\}\\cdot\{s\}\_\{k\}\}\{\\alpha\_\{k\}\{s\}\_\{k\}\\cdot\\nabla J\(\\Theta\_\{k\}\)\}, which reduces the condition number of the transformed Hessian in the new space\. However, this choice was shown to be inferior in terms of performance compared to the standard BFGS, when combined with an inexact line search computation ofαk\\alpha\_\{k\}, as confirmed by\[[22](https://arxiv.org/html/2608.29448#bib.bib20)\]\. The choice of taking the minimum value of both, as adapted by\[[27](https://arxiv.org/html/2608.29448#bib.bib15)\], is essentially a switch between standard BFGS \(τk=1\\tau\_\{k\}=1\) and self\-scaled BFGS which ensures theoretically super\-linear convergence with inexact line searches\.

## 3Methodology:SS\-eSOAP

We proposeSS\-eSOAP, which integrates a self\-scaling factorτk\\tau\_\{k\}and a purifying\-style adaptive update strategy\. Algorithm[1](https://arxiv.org/html/2608.29448#alg1)presents a simplified overview, while Algorithm[2](https://arxiv.org/html/2608.29448#alg2)in the appendix provides the full pseudocode\.

Algorithm 1SS\-eSOAPInput:

W0W\_\{0\},

η\\eta,

β1\\beta\_\{1\},

β2\\beta\_\{2\},

τtrigger\\tau\_\{\\mathrm\{trigger\}\},

L0,R0←IL\_\{0\},R\_\{0\}\\leftarrow I,

QL,QR←IQ\_\{L\},Q\_\{R\}\\leftarrow I,

M~0,V0←0\\widetilde\{M\}\_\{0\},V\_\{0\}\\leftarrow 0
for

k=1,2,…k=1,2,\\dotsdo

Compute gradient

GkG\_\{k\}and

Lk,RkL\_\{k\},R\_\{k\}via \([1](https://arxiv.org/html/2608.29448#S2.E1)\)

if

k\(modinterval\)=0k\\pmod\{\\text\{interval\}\}=0then

Compute off\-diagonal ratio via \([4](https://arxiv.org/html/2608.29448#S3.E4)\)

if

max⁡\{ρ⁡\(Lk,QL\),ρ⁡\(Rk,QR\)\}\>τtrigger\\max\\\{\\rho\(L\_\{k\},Q\_\{L\}\),\\rho\(R\_\{k\},Q\_\{R\}\)\\\}\>\\tau\_\{\\mathrm\{trigger\}\}then

QLold←QLQ\_\{L\}^\{\\mathrm\{old\}\}\\leftarrow Q\_\{L\},

QRold←QRQ\_\{R\}^\{\\mathrm\{old\}\}\\leftarrow Q\_\{R\}
Recompute

QL,QRQ\_\{L\},Q\_\{R\}from

Lk,RkL\_\{k\},R\_\{k\}
M~k−1←\(QL⊤​QLold\)​M~k−1​\(\(QRold\)⊤​QR\)\\widetilde\{M\}\_\{k\-1\}\\leftarrow\(Q\_\{L\}^\{\\top\}Q\_\{L\}^\{\\mathrm\{old\}\}\)\\widetilde\{M\}\_\{k\-1\}\(\(Q\_\{R\}^\{\\mathrm\{old\}\}\)^\{\\top\}Q\_\{R\}\)
Vk−1←γ​Vk−1V\_\{k\-1\}\\leftarrow\\gamma V\_\{k\-1\}
endif

endif

Project

Gk′=QLT​Gk​QRG^\{\\prime\}\_\{k\}=Q\_\{L\}^\{T\}G\_\{k\}Q\_\{R\}and update

Mk′,VkM^\{\\prime\}\_\{k\},V\_\{k\}\(Adam updates in eigenspace\)

Compute

τk\\tau\_\{k\}via \([3](https://arxiv.org/html/2608.29448#S3.E3)\) and precondition

Nk′←τk−1/2Mk′/\(Vk\+ϵ\)N^\{\\prime\}\_\{k\}\\leftarrow\\tau\_\{k\}^\{\-1/2\}M^\{\\prime\}\_\{k\}/\(\\sqrt\{V\_\{k\}\}\+\\epsilon\)
Update

Wk←Wk−1−η​QL​Nk′​QRTW\_\{k\}\\leftarrow W\_\{k\-1\}\-\\eta Q\_\{L\}N^\{\\prime\}\_\{k\}Q\_\{R\}^\{T\}
endfor

### 3\.1Self\-Scaling Curvature Correction

Inspired by Self\-Scaling BFGS \(SS\-BFGS\)\[[27](https://arxiv.org/html/2608.29448#bib.bib15)\], we introduce a scalarτk\\tau\_\{k\}to the inverse Hessian approximation to ensure the preconditioner spans the correct spectrum of the true Hessian\. In the Kronecker structure, we approximate the inverse Hessian as\(H^kS​S\)−1=τk−1​\(Lk−1⊗Rk−1\)\(\\hat\{H\}\_\{k\}^\{SS\}\)^\{\-1\}=\\tau\_\{k\}^\{\-1\}\(L\_\{k\}^\{\-1\}\\otimes R\_\{k\}^\{\-1\}\)\. Leveraging the secant condition properties, we define the scaling factorτk\\tau\_\{k\}using the trace of the updates\. Letsk=vec​\(Wk−Wk−1\)s\_\{k\}=\\text\{vec\}\(W\_\{k\}\-W\_\{k\-1\}\)andyk=vec​\(Gk−Gk−1\)y\_\{k\}=\\text\{vec\}\(G\_\{k\}\-G\_\{k\-1\}\)then for the case of a Kronecker approximation of the Hessian in \([2](https://arxiv.org/html/2608.29448#S2.E2)\) we get:

τk=min⁡\{1,Tr⁡\(\(Gk−Gk−1\)T​\(Wk−Wk−1\)\)Tr⁡\(\(Wk−Wk−1\)T​Lk−1​\(Wk−Wk−1\)​Rk−1\)\}\.\\tau\_\{k\}=\\min\\left\\\{1,\\frac\{\\operatorname\{Tr\}\(\(G\_\{k\}\-G\_\{k\-1\}\)^\{T\}\(W\_\{k\}\-W\_\{k\-1\}\)\)\}\{\\operatorname\{Tr\}\(\(W\_\{k\}\-W\_\{k\-1\}\)^\{T\}L\_\{k\}^\{\-1\}\(W\_\{k\}\-W\_\{k\-1\}\)R\_\{k\}^\{\-1\}\)\}\\right\\\}\.\(3\)We observe thatτk=1\\tau\_\{k\}=1is chosen for most training iterations, with self\-scaling activated only infrequently\. Nevertheless, these occasional adjustments are sufficient to yield substantial improvements in optimization performance\. This formulation ensures that the scale of the preconditioner aligns with the curvature along the most recent displacement\. This is applied in the eigenspace update asN′k=τk−1/2M′k/Vk\+ϵN^\{\\prime\}\_\{k\}=\\tau\_\{k\}^\{\-1/2\}\{M^\{\\prime\}\_\{k\}\}/\{\\sqrt\{V\_\{k\}\}\+\\epsilon\}, whereMk′M^\{\\prime\}\_\{k\}is the momentum in the eigenspace\.

### 3\.2Adaptive Basis Updates and Variance\-State Transition

To address the stability issues of fixed intervals and reprojection, we implement an eigenvalue correction to adaptively control the approximation error of the eigenspace\.

Adaptive eigenbasis trigger\.Following the stopping criterion studied by\[[7](https://arxiv.org/html/2608.29448#bib.bib8)\], we monitor the relative off\-diagonal mass of each Kronecker factor in its current basis,

ρ⁡\(Lk,QL\)=‖QL⊤​Lk​QL−diag⁡\(QL⊤​Lk​QL\)‖F‖QL⊤​Lk​QL‖F\.\\rho\(L\_\{k\},Q\_\{L\}\)=\\frac\{\\\|Q\_\{L\}^\{\\top\}L\_\{k\}Q\_\{L\}\-\\operatorname\{diag\}\(Q\_\{L\}^\{\\top\}L\_\{k\}Q\_\{L\}\)\\\|\_\{F\}\}\{\\\|Q\_\{L\}^\{\\top\}L\_\{k\}Q\_\{L\}\\\|\_\{F\}\}\.\(4\)The statistic measures how poorly the current basis diagonalizes the tracked Kronecker factor\. It is a heuristic proxy for preconditioning degradation, not an exact Hessian\-staleness measure\. We recompute a basis when the maximum ratio for the left and right factors exceedsτtrigger\\tau\_\{\\mathrm\{trigger\}\}\. Appendix Table[9](https://arxiv.org/html/2608.29448#A4.T9)compares this statistic with consecutive gradient cosine similarity around a triggered update\.

State transition after a basis update\.When the basis changes, we first reproject the momentum so that it represents the same update direction in the new coordinates\. We do not reproject the coordinatewise second moment\. Instead, we retain its coordinatewise shape and applyVk−1←γ​Vk−1V\_\{k\-1\}\\leftarrow\\gamma V\_\{k\-1\}, withγ∈\{0\.25,0\.5,0\.75\}\\gamma\\in\\\{0\.25,0\.5,0\.75\\\}selected from the observed off\-diagonal ratio\. BecauseVVappears in the denominator of the Adam\-style update,γ<1\\gamma<1raises the immediate effective step magnitude before new gradients refresh the state\. We treat this transition as an empirical transient heuristic\. Theorem[4\.1](https://arxiv.org/html/2608.29448#S4.Thmtheorem1)gives a steady\-state mismatch bound and does not establish the optimality ofγ<1\\gamma<1\. Appendix Figure[10](https://arxiv.org/html/2608.29448#A4.F10)reports the targeted transient ablation\.

## 4Theoretical Analysis

We analyze self\-scaling and variance downscaling theoretically, with proofs in Appendix[A](https://arxiv.org/html/2608.29448#A1)\.

### 4\.1Spectral Stability Analysis of Variance Adaptation

###### Theorem 4\.1\(Variance\-state stability under basis changes\)\.

LetCt=𝔼⁡\[gt​gt⊤\]C\_\{t\}=\\mathbb\{E\}\[g\_\{t\}g\_\{t\}^\{\\top\}\]denote the population gradient second\-moment matrix at steptt\. For an orthogonal basisQ∈O⁡\(d\)Q\\in O\(d\), define the diagonal second\-moment operator𝒟Q​\(C\):=diag⁡\(Q⊤​C​Q\)\\mathcal\{D\}\_\{Q\}\(C\):=\\operatorname\{diag\}\\left\(Q^\{\\top\}CQ\\right\)\. LetVt=𝒟Qt​\(Ct\)V\_\{t\}=\\mathcal\{D\}\_\{Q\_\{t\}\}\\left\(C\_\{t\}\\right\)andVt\+1∗=𝒟Qt\+1​\(Ct\+1\)V\_\{t\+1\}^\{\*\}=\\mathcal\{D\}\_\{Q\_\{t\+1\}\}\\left\(C\_\{t\+1\}\\right\), whereVtV\_\{t\}is the diagonal second\-moment state in the old basisQtQ\_\{t\}, andVt\+1∗V\_\{t\+1\}^\{\*\}is the oracle diagonal second\-moment state in the new basisQt\+1Q\_\{t\+1\}\. Define the basis transition matrixU=Qt\+1⊤​QtU=Q\_\{t\+1\}^\{\\top\}Q\_\{t\}\. Consider the scalar soft\-reset updateVscale=γ​Vt,0<γ≤1V\_\{\\text\{scale \}\}=\\gamma V\_\{t\},0<\\gamma\\leq 1\. Then

‖Vscale−Vt\+1∗‖F≤\(1−γ\)​‖Vt‖F\+2​‖Ct‖2​‖I−U‖F\+‖Ct\+1−Ct‖F\.\\left\\\|V\_\{\\text\{scale \}\}\-V\_\{t\+1\}^\{\*\}\\right\\\|\_\{F\}\\leq\(1\-\\gamma\)\\left\\\|V\_\{t\}\\right\\\|\_\{F\}\+2\\left\\\|C\_\{t\}\\right\\\|\_\{2\}\\\|I\-U\\\|\_\{F\}\+\\left\\\|C\_\{t\+1\}\-C\_\{t\}\\right\\\|\_\{F\}\.In particular, if the gradient second moment evolves slowly, so that‖Ct\+1−Ct‖F≤Δt\\left\\\|C\_\{t\+1\}\-C\_\{t\}\\right\\\|\_\{F\}\\leq\\Delta\_\{t\}, then

‖Vscale−Vt\+1∗‖F≤\(1−γ\)​‖Vt‖F\+2​‖Ct‖2​‖I−U‖F\+Δt\.\\left\\\|V\_\{\\text\{scale \}\}\-V\_\{t\+1\}^\{\*\}\\right\\\|\_\{F\}\\leq\(1\-\\gamma\)\\left\\\|V\_\{t\}\\right\\\|\_\{F\}\+2\\left\\\|C\_\{t\}\\right\\\|\_\{2\}\\\|I\-U\\\|\_\{F\}\+\\Delta\_\{t\}\.

### 4\.2Self\-Scaling Spectral Analysis

###### Theorem 4\.2\(Kronecker secant\-energy matching\)\.

LetSk=Wk−Wk−1S\_\{k\}=W\_\{k\}\-W\_\{k\-1\},Yk=∇ℒ​\(Wk\)−∇ℒ​\(Wk−1\)Y\_\{k\}=\\nabla\\mathcal\{L\}\(W\_\{k\}\)\-\\nabla\\mathcal\{L\}\(W\_\{k\-1\}\), and letKk≻0K\_\{k\}\\succ 0denote the Kronecker metric whose quadratic form issk⊤​Kk​sk=Tr⁡\(Sk⊤​Lk−1​Sk​Rk−1\)s\_\{k\}^\{\\top\}K\_\{k\}s\_\{k\}=\\operatorname\{Tr\}\\\!\\left\(S\_\{k\}^\{\\top\}L\_\{k\}^\{\-1\}S\_\{k\}R\_\{k\}^\{\-1\}\\right\), wheresk=vec⁡\(Sk\)s\_\{k\}=\\operatorname\{vec\}\(S\_\{k\}\)\. Assumesk≠0s\_\{k\}\\neq 0andsk⊤​yk\>0s\_\{k\}^\{\\top\}y\_\{k\}\>0\. Defineak=sk⊤​Kk​ska\_\{k\}=s\_\{k\}^\{\\top\}K\_\{k\}s\_\{k\}andck=sk⊤​ykc\_\{k\}=s\_\{k\}^\{\\top\}y\_\{k\}\. Then the scaling factor

τk=min⁡\{1,ckak\}=min⁡\{1,sk⊤​yksk⊤​Kk​sk\}\\displaystyle\\tau\_\{k\}=\\min\\left\\\{1,\\frac\{c\_\{k\}\}\{a\_\{k\}\}\\right\\\}=\\min\\left\\\{1,\\frac\{s\_\{k\}^\{\\top\}y\_\{k\}\}\{s\_\{k\}^\{\\top\}K\_\{k\}s\_\{k\}\}\\right\\\}is the unique solution of the clipped one\-dimensional log\-secant matching problem

τk=arg⁡min0<τ≤1⁡\[log⁡sk⊤​yksk⊤​\(τ​Kk\)​sk\]2\.\\displaystyle\\tau\_\{k\}=\\arg\\min\_\{0<\\tau\\leq 1\}\\left\[\\log\\frac\{s\_\{k\}^\{\\top\}y\_\{k\}\}\{s\_\{k\}^\{\\top\}\(\\tau K\_\{k\}\)s\_\{k\}\}\\right\]^\{2\}\.Ifyk=Hk​sk\+rky\_\{k\}=H\_\{k\}s\_\{k\}\+r\_\{k\}, then, whenever the unclipped solution is active, the scaled Kronecker quadratic form satisfiessk⊤​\(τk​Kk\)​sk=sk⊤​Hk​sk\+sk⊤​rks\_\{k\}^\{\\top\}\(\\tau\_\{k\}K\_\{k\}\)s\_\{k\}=s\_\{k\}^\{\\top\}H\_\{k\}s\_\{k\}\+s\_\{k\}^\{\\top\}r\_\{k\}\. Thus the mismatch between the scaled Kronecker secant energy and the true directional curvature is exactly the directional secant errorsk⊤​rks\_\{k\}^\{\\top\}r\_\{k\}\.

## 5Numerical Experiments

We evaluateSS\-eSOAPon a suite of challenging optimization tasks, including matrix regression, low\-rank completion tasks and nonlinear evolutionary PDEs\. The matrix tasks provide controlled settings where the curvature structure is known, allowing us to isolate the effect of self\-scaling on the optimizer dynamics\. In contrast, the nonlinear PDE benchmarks evaluate its robustness against the curvature anisotropy and stiffness typically encountered in frontier AI for Science applications\[[31](https://arxiv.org/html/2608.29448#bib.bib12),[17](https://arxiv.org/html/2608.29448#bib.bib14),[25](https://arxiv.org/html/2608.29448#bib.bib13)\]\.

### 5\.1Baselines

The main comparisons include Adam\[[15](https://arxiv.org/html/2608.29448#bib.bib16)\], Muon\[[13](https://arxiv.org/html/2608.29448#bib.bib27)\], SOAP\[[28](https://arxiv.org/html/2608.29448#bib.bib25)\], and the adaptive Shampoo variant implemented from\[[7](https://arxiv.org/html/2608.29448#bib.bib8)\]\. We also evaluate L\-BFGS, SS\-BFGS, SS\-Broyden, K\-FAC for PINNs, NTK reweighting, and MultiAdam where their memory and training requirements fit the benchmark\[[18](https://arxiv.org/html/2608.29448#bib.bib22),[16](https://arxiv.org/html/2608.29448#bib.bib3),[6](https://arxiv.org/html/2608.29448#bib.bib35),[32](https://arxiv.org/html/2608.29448#bib.bib6),[37](https://arxiv.org/html/2608.29448#bib.bib5)\]\. Appendix[C\.4](https://arxiv.org/html/2608.29448#A3.SS4)reports each search space, selected setting, precision, stopping rule, hardware, and seed\.

### 5\.2Matrix Regression Tasks

To evaluate the robustness and versatility ofSS\-eSOAPacross varying optimization landscapes, we conduct a comprehensive suite of matrix optimization experiments\. These tasks span the spectrum of convexity and gradient stochasticity, ranging from strongly convex problems and convex objectives to challenging non\-convex settings\. This diverse testbed allows us to rigorously benchmarkSS\-eSOAPagainst baselines in both deterministic and stochastic regimes\. Detailed analysis is in Appendix[C\.3](https://arxiv.org/html/2608.29448#A3.SS3)\.

Figure 1:Matrix Quadratic Regression Task\.SS\-eSOAPconsistently achieves faster convergence and lower residuals in the quadratic regime where the self\-scaling mechanism captures the spectrum\.Quadratic regression\.We consider a matrix quadratic regression objectivef⁡\(X\):=12​‖A​X​B−C‖F2f\(X\):=\\frac\{1\}\{2\}\\\|AXB\-C\\\|\_\{\\mathrm\{F\}\}^\{2\}, whereX∈ℝm×n,A∈ℝp×m,B∈ℝn×qX\\in\\mathbb\{R\}^\{m\\times n\},A\\in\\mathbb\{R\}^\{p\\times m\},B\\in\\mathbb\{R\}^\{n\\times q\}andC∈ℝp×qC\\in\\mathbb\{R\}^\{p\\times q\}\. Then its gradient is∇f​\(X\)=A⊤​\(A​X​B−C\)​B⊤\\nabla f\(X\)=A^\{\\top\}\(AXB\-C\)B^\{\\top\}, its Hessian is∇2f​\(X\)=\(B​B⊤\)⊗\(A⊤​A\)∈ℝm​n×m​n\\nabla^\{2\}f\(X\)=\(BB^\{\\top\}\)\\otimes\(A^\{\\top\}A\)\\in\\mathbb\{R\}^\{mn\\times mn\}\. The Hessian structure implies that ideal preconditioning isolates the residual condition numberκ⁡\(E\)\\kappa\(E\)from the constant spectral skew ofAAandBB\. We set\(m,n,p,q\)=\(500,100,1000,250\)\(m,n,p,q\)=\(500,100,1000,250\)so thatffis strongly convex\.

As illustrated in Figure[1](https://arxiv.org/html/2608.29448#S5.F1),SS\-eSOAPdemonstrates a clear advantage in navigating the structured curvature of the quadratic objective\.

The quadratic task isolates the setting where the Hessian has an exact Kronecker form\. It therefore tests whether the scalar correction recovers the directional scale predicted by the analysis without the nonlinear and multi\-objective effects of a PINN\. Logistic regression and low\-rank completion results appear in Appendix[C\.3](https://arxiv.org/html/2608.29448#A3.SS3)\.

### 5\.3PINN Benchmarks

We applySS\-eSOAPto standard PINN benchmarks including the Allen\-Cahn, 1D Burgers, and 2D Boussinesq equations, assessing performance across varying degrees of physical complexity\.

Table[1](https://arxiv.org/html/2608.29448#S5.T1)shows thatSS\-eSOAPachieves the lowest final residual on six of eight PINN benchmarks\. The largest gains appear on Wave, whereSS\-eSOAPreaches4\.12×10−54\.12\\times 10^\{\-5\}compared to1\.18×10−31\.18\\times 10^\{\-3\}for the best baseline, and on Boussinesq, where it reaches8\.42×10−78\.42\\times 10^\{\-7\}compared to2\.85×10−62\.85\\times 10^\{\-6\}\.SS\-eSOAPalso improves over the strongest baselines on Burgers \(5\.54×10−105\.54\\times 10^\{\-10\}vs\.1\.12×10−91\.12\\times 10^\{\-9\}\), Korteweg\-de Vries \(3\.52×10−103\.52\\times 10^\{\-10\}vs\.8\.12×10−108\.12\\times 10^\{\-10\}\), and Lid\-Driven Cavity \(8\.23×10−88\.23\\times 10^\{\-8\}vs\.2\.10×10−72\.10\\times 10^\{\-7\}\)\. The exceptions are Gray\-Scott, where Purifying Shampoo obtains5\.12×10−95\.12\\times 10^\{\-9\}compared to7\.46×10−77\.46\\times 10^\{\-7\}forSS\-eSOAP, and Ginzburg\-Landau, where SOAP obtains9\.12×10−129\.12\\times 10^\{\-12\}compared to1\.75×10−111\.75\\times 10^\{\-11\}forSS\-eSOAP\. Overall, these results indicate thatSS\-eSOAPis most beneficial in regimes where additional scale calibration and damping improve the stability of structured preconditioning, while the gains are problem\-dependent and not uniform across all PDEs\.

1D Burgers and 2D Boussinesq\.To validate our optimizer’s capability in this frontier regime, we adopt the rigorous high\-precision formulations for the 1D Burgers and 2D Boussinesq equations established by\[[33](https://arxiv.org/html/2608.29448#bib.bib26)\]\. In these stiff settings, standard first\-order methods typically stall due to spectral cliffs\. However,SS\-eSOAPreaches substantially lower residuals, driving the residual down to5\.54×10−105\.54\\times 10^\{\-10\}precision for the Burgers equation, which is an improvement of about 3\.9 orders of magnitude over Adam\. Similarly, on the 2D Boussinesq equation,SS\-eSOAPachieves a final residual of8\.42×10−78\.42\\times 10^\{\-7\}, successfully resolving the singular blow\-up profiles where baselines like Muon and Adam stagnate at2\.57×10−22\.57\\times 10^\{\-2\}and5\.12×10−45\.12\\times 10^\{\-4\}respectively\.

Additional benchmarks\.Additional results for Lid\-Driven Cavity, Wave, Korteweg\-de Vries, Gray\-Scott, and Ginzburg\-Landau are provided in Appendix Figure[8](https://arxiv.org/html/2608.29448#A3.F8)\. Detailed problem formulations and hyperparameter settings are given in Appendix[C\.2](https://arxiv.org/html/2608.29448#A3.SS2)\.

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](https://arxiv.org/html/2608.29448#A3.SS4)\. Bold marks the lowest value within each row\. Residual is an optimization metric\. Physical solution errors and three\-seed statistics appear in Appendix[8](https://arxiv.org/html/2608.29448#A4.T8)\.PDE residual alone does not establish solution accuracy\. We therefore compare predictions with high\-fidelity reference solutions on four representative PDEs\. Appendix Table[8](https://arxiv.org/html/2608.29448#A4.T8)reports relativeL2L^\{2\}andH1H^\{1\}errors over three seeds for general\-purpose, PINN\-specific, and quasi\-Newton baselines\. The error ranking agrees with the residual ranking on the stiff Burgers and Boussinesq cases, while the full table records the remaining method\-dependent differences\.

### 5\.4Component Analysis and Ablation Studies

To isolate the impact of the specific mechanisms introduced inSS\-eSOAP, namely variance downscaling, Kronecker\-aware self\-scaling, and the adaptive trigger, we conduct targeted ablation studies on the stiff Allen\-Cahn and 2D Boussinesq benchmarks\.

Variance\-state transition\.Figure[2](https://arxiv.org/html/2608.29448#S5.F2)compares downscaling, reprojection, and a reset on Allen\-Cahn and Boussinesq\. Reprojection produces large transient loss spikes in these runs\. A reset avoids the largest spikes but converges to a higher final loss\. The targetedγ\\gammaexperiment in Appendix Figure[10](https://arxiv.org/html/2608.29448#A4.F10)separates behavior immediately before and after a triggered basis update\.

Kronecker\-adapted secant scaling\.We compare the directional correction in Equation \([3](https://arxiv.org/html/2608.29448#S3.E3)\) with an unscaled update and a scalar correction that omits the Kronecker metric\. Figure[2](https://arxiv.org/html/2608.29448#S5.F2)reports the resulting trajectories\. This experiment tests directional scale matching\. It does not compare a full dense SS\-BFGS update, which appears only in the tractable small\-network experiment in Appendix Table[7](https://arxiv.org/html/2608.29448#A4.T7)\.

Figure[2](https://arxiv.org/html/2608.29448#S5.F2)summarizes the ablations for the proposed structural choices\. Figures[2\(a\)](https://arxiv.org/html/2608.29448#S5.F2.sf1),[2\(c\)](https://arxiv.org/html/2608.29448#S5.F2.sf3)show that adaptive variance downscaling stabilizes basis updates\. Figures[2\(b\)](https://arxiv.org/html/2608.29448#S5.F2.sf2),[2\(d\)](https://arxiv.org/html/2608.29448#S5.F2.sf4)show that the Kronecker\-adapted scaling factorτk\\tau\_\{k\}improves stability and convergence relative to unscaled and scalar SS\-BFGS\-style variants\. Together, these results suggest thatSS\-eSOAPcaptures part of the step\-size regulation benefit of self\-scaled quasi\-Newton methods without requiring dense Hessian cross\-terms\.

Efficiency of the adaptive trigger\.Finally, we analyze the sensitivity of the eigenbasis update thresholdτtrigger\\tau\_\{\\text\{trigger\}\}\. As shown in Table[11](https://arxiv.org/html/2608.29448#A7.T11)in the Appendix, extreme values degrade performance:τtrigger=0\\tau\_\{\\text\{trigger\}\}=0\(updates every step\) incurs prohibitive computational cost, whileτtrigger≥0\.8\\tau\_\{\\text\{trigger\}\}\\geq 0\.8\(lazy updates\) leads to stale curvature estimates and poor convergence\. However, we observe a broad stable operating regionτtrigger∈\[0\.15,0\.3\]\\tau\_\{\\text\{trigger\}\}\\in\[0\.15,0\.3\]where performance remains consistent\. We use the fixed default values \(τtrigger=0\.2\\tau\_\{\\text\{trigger\}\}=0\.2and discreteγ∈\{0\.25,0\.5,0\.75\}\\gamma\\in\\\{0\.25,0\.5,0\.75\\\}\) across all experiments in this paper, ranging from simple regression tasks to the stiff Burgers and Boussinesq PDEs\.SS\-eSOAPachieves strong results across diverse problems without problem\-specific tuning, suggesting that its adaptive mechanisms are robust across scales and stiffness regimes\.

\(a\)Stability ablation, adaptive eigenbasis, Boussinesq\.\(b\)Self\-scaling, Boussinesq\.\(c\)Stability ablation, adaptive eigenbasis, Allen\-Cahn\.\(d\)Self\-scaling, Allen\-Cahn\.
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 \(τO​L\\tau\_\{OL\}, SS\-BFGS style\), and the unscaled baseline \(τ=1\\tau=1, SS\-\), validating structure\-aware step\-size regulation\.
### 5\.5Computational Efficiency Summary

WhileSS\-eSOAPintroduces a modest per\-step computational overhead compared to first\-order methods, its curvature\-aware updates drastically reduce the total iterations required for convergence on stiff problems\. As summarized in Table[2](https://arxiv.org/html/2608.29448#S5.T2),SS\-eSOAPmaintains a memory footprint highly competitive with Adam by leveraging Kronecker factorization\. This effectively bypasses the prohibitive storage costs of full quasi\-Newton methods\. Ultimately, this structural efficiency translates to a substantial reduction in time\-to\-solution for high\-precision PDE training\. A comprehensive analysis detailing structural efficiency and memory tradeoffs is provided in Appendix[F](https://arxiv.org/html/2608.29448#A6)\.

Table 2:Compute and memory comparison between Adam andSS\-eSOAPon 2D Boussinesq, targeting residual10−510^\{\-5\}on a single consumer GPU\.

## 6Limitations

Our analysis does not prove global convergence or sufficient decrease for the full nonconvex algorithm\. The basis trigger uses the tracked Kronecker factors as a heuristic proxy for preconditioning degradation\. ItsO⁡\(d3\)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 modern residual\-space natural\-gradient methods form a strong alternative\. The empirical gains are not uniform: SOAP\-family methods perform better on Gray\-Scott and Ginzburg\-Landau\. All main PDE experiments use MLP or PirateNet backbones\. Transfer to transformer\-based physics models such as PINNsFormer, whose reported loss surface is smoother\[[38](https://arxiv.org/html/2608.29448#bib.bib1)\], remains unevaluated\. The three\-seed study covers four representative PDEs, so broader seed\-level conclusions require more runs\.

## 7Conclusion

We introducedSS\-eSOAP, a Kronecker\-factored optimizer that combines directional secant\-energy matching with adaptive basis updates and a variance\-state transition\. The method attains the lowest final residual on six of eight PDE benchmarks and improves wall\-clock time to a fixed target on the stiff Boussinesq case\. Three\-seed physical errors on four PDEs support the solution quality of these low\-residual fits\. Results on Gray\-Scott and Ginzburg\-Landau show that the added scaling is not uniformly beneficial\. The main open questions are how to reduce the cubic basis\-check cost, how to select the state transition from online diagnostics, and how the method transfers to smoother physics architectures\.

## References

- \[1\]M\. Al\-Baali and H\. Khalfan\(2005\)Wide interval for efficient self\-scaling quasi\-newton algorithms\.Optimization Methods and Software20\(6\),pp\. 679–691\.External Links:[Document](https://dx.doi.org/10.1080/10556780410001709448)Cited by:[§2\.5](https://arxiv.org/html/2608.29448#S2.SS5.p1.2)\.
- \[2\]M\. Al\-Baali\(1992\)Highly efficient broyden methods of minimization with variable parameter\.Optimization Methods and Software1\(4\),pp\. 301–310\.External Links:[Document](https://dx.doi.org/10.1080/10556789208805528)Cited by:[§1](https://arxiv.org/html/2608.29448#S1.p3.1),[§2\.5](https://arxiv.org/html/2608.29448#S2.SS5.p1.1)\.
- \[3\]M\. Al\-Baali\(1993\)Variational quasi\-newton methods for unconstrained optimization\.Journal of Optimization Theory and Applications77\(1\),pp\. 127–143\.External Links:[Document](https://dx.doi.org/10.1007/BF00940782)Cited by:[§2\.5](https://arxiv.org/html/2608.29448#S2.SS5.p1.2)\.
- \[4\]J\. Bernstein and L\. Newhouse\(2024\)Old optimizer, new norm: an anthology\.External Links:2409\.20325,[Link](https://arxiv.org/abs/2409.20325)Cited by:[§2\.3](https://arxiv.org/html/2608.29448#S2.SS3.p1.2),[§2\.3](https://arxiv.org/html/2608.29448#S2.SS3.p3.1)\.
- \[5\]C\. G\. BROYDEN\(1970\)The convergence of a class of double\-rank minimization algorithms 1\. general considerations\.IMA Journal of Applied Mathematics6\(1\),pp\. 76–90\.External Links:ISSN 0272\-4960,[Document](https://dx.doi.org/10.1093/imamat/6.1.76),[Link](https://doi.org/10.1093/imamat/6.1.76),https://academic\.oup\.com/imamat/article\-pdf/6/1/76/2233756/6\-1\-76\.pdfCited by:[§1](https://arxiv.org/html/2608.29448#S1.p3.1)\.
- \[6\]F\. Dangel, J\. Müller, and M\. Zeinhofer\(2024\)Kronecker\-factored approximate curvature for physics\-informed neural networks\.InAdvances in Neural Information Processing Systems \(NeurIPS\),Cited by:[Appendix J](https://arxiv.org/html/2608.29448#A10.p1.1),[§2\.2](https://arxiv.org/html/2608.29448#S2.SS2.p1.1),[§5\.1](https://arxiv.org/html/2608.29448#S5.SS1.p1.1)\.
- \[7\]R\. Eschenhagen, A\. Defazio, T\. Lee, R\. E\. Turner, and H\. M\. Shi\(2025\)Purifying shampoo: investigating shampoo’s heuristics by decomposing its preconditioner\.arXiv preprint arXiv:2506\.03595\.Cited by:[§B\.1](https://arxiv.org/html/2608.29448#A2.SS1.p1.1),[§B\.1](https://arxiv.org/html/2608.29448#A2.SS1.p2.1),[§B\.2](https://arxiv.org/html/2608.29448#A2.SS2.p1.1),[§1](https://arxiv.org/html/2608.29448#S1.p4.1),[§1](https://arxiv.org/html/2608.29448#S1.p5.1),[§2\.3](https://arxiv.org/html/2608.29448#S2.SS3.p2.1),[§2\.4](https://arxiv.org/html/2608.29448#S2.SS4.p2.1),[§3\.2](https://arxiv.org/html/2608.29448#S3.SS2.p2.1),[§5\.1](https://arxiv.org/html/2608.29448#S5.SS1.p1.1)\.
- \[8\]V\. Gupta and et al\.\(2018\)Shampoo: preconditioned stochastic tensor optimization\.InProceedings of the 35th International Conference on Machine Learning \(ICML\),pp\. 1842–1850\.External Links:[Link](https://proceedings.mlr.press/v80/gupta18a.html)Cited by:[§1](https://arxiv.org/html/2608.29448#S1.p4.1),[§2\.3](https://arxiv.org/html/2608.29448#S2.SS3.p1.1)\.
- \[9\]A\. Guzmán\-Cordero, F\. Dangel, G\. Goldshlager, and M\. Zeinhofer\(2025\)Improving energy natural gradient descent through Woodbury, momentum, and randomization\.InAdvances in Neural Information Processing Systems \(NeurIPS\),Cited by:[Appendix J](https://arxiv.org/html/2608.29448#A10.p1.1)\.
- \[10\]A\. Jnini, E\. Kiyani, K\. Shukla, J\. F\. Urban, N\. A\. Daryakenari, J\. Müller, M\. Zeinhofer, and G\. E\. Karniadakis\(2026\)Curvature\-aware optimization for high\-accuracy physics\-informed neural networks\.Computer Methods in Applied Mechanics and Engineering462,pp\. 119289\.Cited by:[§2\.2](https://arxiv.org/html/2608.29448#S2.SS2.p1.1)\.
- \[11\]A\. Jnini, F\. Vella, and M\. Zeinhofer\(2024\)Gauss\-newton natural gradient descent for physics\-informed computational fluid dynamics\.Computers & Fluids\.Cited by:[Appendix J](https://arxiv.org/html/2608.29448#A10.p1.1),[§2\.2](https://arxiv.org/html/2608.29448#S2.SS2.p1.1)\.
- \[12\]A\. Jnini and F\. Vella\(2025\)Dual natural gradient descent for scalable training of physics\-informed neural networks\.Transactions on Machine Learning Research\.Cited by:[Appendix J](https://arxiv.org/html/2608.29448#A10.p1.1),[§2\.2](https://arxiv.org/html/2608.29448#S2.SS2.p1.1)\.
- \[13\]K\. Jordan, Y\. Jin, V\. Boza, J\. You, F\. Cesista, L\. Newhouse, and J\. Bernstein\(2024\)Muon: an optimizer for hidden layers in neural networks\.Note:[https://kellerjordan\.github\.io/posts/muon/](https://kellerjordan.github.io/posts/muon/)Accessed: 2026\-01\-27Cited by:[§2\.3](https://arxiv.org/html/2608.29448#S2.SS3.p1.2),[§2\.3](https://arxiv.org/html/2608.29448#S2.SS3.p3.1),[§5\.1](https://arxiv.org/html/2608.29448#S5.SS1.p1.1)\.
- \[14\]G\. E\. Karniadakis, I\. G\. Kevrekidis, L\. Lu, P\. Perdikaris, S\. Wang, and L\. Yang\(2021\)Physics\-informed machine learning\.Nature Reviews Physics3,pp\. 422 – 440\.External Links:[Link](https://api.semanticscholar.org/CorpusID:236407461)Cited by:[§2\.1](https://arxiv.org/html/2608.29448#S2.SS1.p1.1)\.
- \[15\]D\. P\. Kingma and J\. Ba\(2015\)Adam: a method for stochastic optimization\.InInternational Conference on Learning Representations \(ICLR\),External Links:[Link](https://arxiv.org/abs/1412.6980)Cited by:[§5\.1](https://arxiv.org/html/2608.29448#S5.SS1.p1.1)\.
- \[16\]E\. Kiyani, K\. Shukla, J\. F\. Urbán, J\. Darbon, and G\. E\. Karniadakis\(2025\)Optimizing the optimizer for physics\-informed neural networks and kolmogorov\-arnold networks\.Computer Methods in Applied Mechanics and Engineering446,pp\. 118308\.Cited by:[§2\.2](https://arxiv.org/html/2608.29448#S2.SS2.p1.1),[§5\.1](https://arxiv.org/html/2608.29448#S5.SS1.p1.1)\.
- \[17\]A\. Krishnapriyan, A\. Gholami, S\. Zhe, R\. Kirby, and M\. W\. Mahoney\(2021\)Characterizing possible failure modes in physics\-informed neural networks\.Advances in neural information processing systems34,pp\. 26548–26560\.Cited by:[§1](https://arxiv.org/html/2608.29448#S1.p1.1),[§1](https://arxiv.org/html/2608.29448#S1.p2.1),[§2\.2](https://arxiv.org/html/2608.29448#S2.SS2.p1.1),[§5](https://arxiv.org/html/2608.29448#S5.p1.1)\.
- \[18\]D\. C\. Liu and J\. Nocedal\(1989\)On the limited memory bfgs method for large scale optimization\.Mathematical Programming45\(1\-3\),pp\. 503–528\.External Links:[Link](https://api.semanticscholar.org/CorpusID:5681609)Cited by:[§1](https://arxiv.org/html/2608.29448#S1.p3.1),[§5\.1](https://arxiv.org/html/2608.29448#S5.SS1.p1.1)\.
- \[19\]S\. Liu, C\. Su, J\. Yao, Z\. Hao, H\. Su, Y\. Wu, and J\. Zhu\(2024\)Preconditioning for physics\-informed neural networks\.arXiv preprint arXiv:2402\.00531\.Cited by:[§2\.2](https://arxiv.org/html/2608.29448#S2.SS2.p1.1)\.
- \[20\]J\. Martens and R\. Grosse\(2016\)Kronecker\-factored approximate curvature for deep networks\.InInternational Conference on Machine Learning \(ICML\),External Links:[Link](https://arxiv.org/abs/1503.05671)Cited by:[§1](https://arxiv.org/html/2608.29448#S1.p4.1)\.
- \[21\]J\. Müller and M\. Zeinhofer\(2023\)Achieving high accuracy with PINNs via energy natural gradients\.InInternational Conference on Machine Learning \(ICML\),Cited by:[Appendix J](https://arxiv.org/html/2608.29448#A10.p1.1),[§2\.2](https://arxiv.org/html/2608.29448#S2.SS2.p1.1)\.
- \[22\]J\. Nocedal and Y\. Yuan\(1993\)Analysis of a self\-scaling quasi\-Newton method\.Mathematical Programming61\(1\),pp\. 19–37\.Cited by:[§1](https://arxiv.org/html/2608.29448#S1.p3.1),[§2\.5](https://arxiv.org/html/2608.29448#S2.SS5.p2.1)\.
- \[23\]S\. S\. Oren and D\. G\. Luenberger\(1974\)Self\-scaling variable metric \(ssvm\) algorithms\.Management Science20\(5\),pp\. 845–862\.External Links:[Document](https://dx.doi.org/10.1287/mnsc.20.5.845)Cited by:[§1](https://arxiv.org/html/2608.29448#S1.p3.1),[§2\.5](https://arxiv.org/html/2608.29448#S2.SS5.p2.1)\.
- \[24\]M\. Raissi, P\. Perdikaris, and G\.E\. Karniadakis\(2019\)Physics\-informed neural networks: a deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations\.Journal of Computational Physics378,pp\. 686–707\.External Links:ISSN 0021\-9991,[Document](https://dx.doi.org/https%3A//doi.org/10.1016/j.jcp.2018.10.045),[Link](https://www.sciencedirect.com/science/article/pii/S0021999118307125)Cited by:[§2\.1](https://arxiv.org/html/2608.29448#S2.SS1.p1.1)\.
- \[25\]P\. Rathore, W\. Lei, Z\. Frangella, L\. Lu, and M\. Udell\(2024\)Challenges in training pinns: a loss landscape perspective\.arXiv preprint arXiv:2402\.01868\.Cited by:[§1](https://arxiv.org/html/2608.29448#S1.p1.1),[§1](https://arxiv.org/html/2608.29448#S1.p2.1),[§2\.2](https://arxiv.org/html/2608.29448#S2.SS2.p1.1),[§5](https://arxiv.org/html/2608.29448#S5.p1.1)\.
- \[26\]N\. Schwencke and C\. Furtlehner\(2025\)ANaGRAM: a natural gradient relative to adapted model for efficient PINNs learning\.InThe Thirteenth International Conference on Learning Representations \(ICLR\),Cited by:[Appendix J](https://arxiv.org/html/2608.29448#A10.p1.1),[§2\.2](https://arxiv.org/html/2608.29448#S2.SS2.p1.1)\.
- \[27\]J\. F\. Urbán, P\. Stefanou, and J\. A\. Pons\(2025\)Unveiling the optimization process of physics informed neural networks: how accurate and competitive can pinns be?\.Journal of Computational Physics523,pp\. 113656\.Cited by:[§1](https://arxiv.org/html/2608.29448#S1.p2.1),[§1](https://arxiv.org/html/2608.29448#S1.p3.1),[§2\.5](https://arxiv.org/html/2608.29448#S2.SS5.p1.4),[§2\.5](https://arxiv.org/html/2608.29448#S2.SS5.p2.1),[§3\.1](https://arxiv.org/html/2608.29448#S3.SS1.p1.1)\.
- \[28\]N\. Vyas, D\. Morwani, R\. Zhao, and et al\.\(2024\)SOAP: improving and stabilizing shampoo using adam\.arXiv preprint arXiv:2409\.11321\.External Links:[Link](https://arxiv.org/abs/2409.11321)Cited by:[§1](https://arxiv.org/html/2608.29448#S1.p4.1),[§2\.3](https://arxiv.org/html/2608.29448#S2.SS3.p1.2),[§5\.1](https://arxiv.org/html/2608.29448#S5.SS1.p1.1)\.
- \[29\]S\. Wang, A\. K\. Bhartari, B\. Li, and P\. Perdikaris\(2025\)Gradient alignment in physics\-informed neural networks: a second\-order optimization perspective\.arXiv preprint arXiv:2502\.00604\.Cited by:[§1](https://arxiv.org/html/2608.29448#S1.p1.1),[§2\.2](https://arxiv.org/html/2608.29448#S2.SS2.p1.1)\.
- \[30\]S\. Wang, B\. Li, Y\. Chen, and P\. Perdikaris\(2024\)Piratenets: physics\-informed deep learning with residual adaptive networks\.Journal of Machine Learning Research25\(402\),pp\. 1–51\.Cited by:[§C\.1](https://arxiv.org/html/2608.29448#A3.SS1.p1.1),[§C\.1](https://arxiv.org/html/2608.29448#A3.SS1.p1.3)\.
- \[31\]S\. Wang, Y\. Teng, and P\. Perdikaris\(2021\)Understanding and mitigating gradient flow pathologies in physics\-informed neural networks\.SIAM Journal on Scientific Computing43\(5\),pp\. A3055–A3081\.Cited by:[§1](https://arxiv.org/html/2608.29448#S1.p1.1),[§1](https://arxiv.org/html/2608.29448#S1.p2.1),[§2\.2](https://arxiv.org/html/2608.29448#S2.SS2.p1.1),[§5](https://arxiv.org/html/2608.29448#S5.p1.1)\.
- \[32\]S\. Wang, X\. Yu, and P\. Perdikaris\(2022\)When and why pinns fail to train: a neural tangent kernel perspective\.Journal of Computational Physics449,pp\. 110768\.Cited by:[§2\.2](https://arxiv.org/html/2608.29448#S2.SS2.p1.1),[§5\.1](https://arxiv.org/html/2608.29448#S5.SS1.p1.1)\.
- \[33\]Y\. Wang, Z\. Liu, Z\. Li, A\. Anandkumar, and T\. Y\. Hou\(2025\)High precision pinns in unbounded domains: application to singularity formulation in pdes\.arXiv preprint arXiv:2506\.19243\.Cited by:[§C\.2](https://arxiv.org/html/2608.29448#A3.SS2.p1.3),[§C\.2](https://arxiv.org/html/2608.29448#A3.SS2.p3.1),[§1](https://arxiv.org/html/2608.29448#S1.p1.1),[§5\.3](https://arxiv.org/html/2608.29448#S5.SS3.p3.1)\.
- \[34\]Y\. Wang, M\. Bennani, J\. Martens, S\. Racanière, S\. Blackwell, A\. Matthews, S\. Nikolov, G\. Cao\-Labora, D\. S\. Park, M\. Arjovsky,et al\.\(2025\)Discovery of unstable singularities\.arXiv preprint arXiv:2509\.14185\.Cited by:[§1](https://arxiv.org/html/2608.29448#S1.p1.1)\.
- \[35\]Y\. Wang, C\. Lai, J\. Gómez\-Serrano, and T\. Buckmaster\(2023\)Asymptotic self\-similar blow\-up profile for three\-dimensional axisymmetric euler equations using neural networks\.Physical Review Letters130\(24\),pp\. 244002\.Cited by:[§C\.2](https://arxiv.org/html/2608.29448#A3.SS2.p2.1),[§1](https://arxiv.org/html/2608.29448#S1.p1.1)\.
- \[36\]C\. Xu, D\. Liu, A\. Nassereldine, and J\. Xiong\(2025\)FP64 is all you need: rethinking failure modes in physics\-informed neural networks\.Advances in Nerual Information Processing Systems\.Cited by:[§1](https://arxiv.org/html/2608.29448#S1.p1.1)\.
- \[37\]J\. Yao, C\. Su, Z\. Hao, S\. Liu, H\. Su, and J\. Zhu\(2023\)Multiadam: parameter\-wise scale\-invariant optimizer for multiscale training of physics\-informed neural networks\.InInternational conference on machine learning,pp\. 39702–39721\.Cited by:[§2\.2](https://arxiv.org/html/2608.29448#S2.SS2.p1.1),[§5\.1](https://arxiv.org/html/2608.29448#S5.SS1.p1.1)\.
- \[38\]Z\. Zhao, X\. Ding, and B\. A\. Prakash\(2024\)PINNsformer: a transformer\-based framework for physics\-informed neural networks\.InThe Twelfth International Conference on Learning Representations,External Links:[Link](https://openreview.net/forum?id=DO2WFXU1Be)Cited by:[§6](https://arxiv.org/html/2608.29448#S6.p1.1)\.

## Appendix

## Appendix AProofs of Main Theorems

### A\.1Proof of Theorem[4\.1](https://arxiv.org/html/2608.29448#S4.Thmtheorem1)

###### Proof\.

By the triangle inequality,

‖γ​Vt−Vt\+1∗‖F≤‖\(γ−1\)​Vt‖F\+‖Vt−Vt\+1∗‖F\.\\left\\\|\\gamma V\_\{t\}\-V\_\{t\+1\}^\{\*\}\\right\\\|\_\{F\}\\leq\\left\\\|\(\\gamma\-1\)V\_\{t\}\\right\\\|\_\{F\}\+\\left\\\|V\_\{t\}\-V\_\{t\+1\}^\{\*\}\\right\\\|\_\{F\}\.Since0<γ≤10<\\gamma\\leq 1,

‖\(γ−1\)​Vt‖F=\(1−γ\)​‖Vt‖F\.\\left\\\|\(\\gamma\-1\)V\_\{t\}\\right\\\|\_\{F\}=\(1\-\\gamma\)\\left\\\|V\_\{t\}\\right\\\|\_\{F\}\.It remains to bound the second term,‖Vt−Vt\+1∗‖F\\left\\\|V\_\{t\}\-V\_\{t\+1\}^\{\*\}\\right\\\|\_\{F\}\. DefineMt:=Qt⊤​Ct​QtM\_\{t\}:=Q\_\{t\}^\{\\top\}C\_\{t\}Q\_\{t\}\. SinceU=Qt\+1⊤​QtU=Q\_\{t\+1\}^\{\\top\}Q\_\{t\}, we haveQt\+1=Qt​U⊤Q\_\{t\+1\}=Q\_\{t\}U^\{\\top\}\. Therefore,

Qt\+1⊤​Ct​Qt\+1=U​Qt⊤​Ct​Qt​U⊤=U​Mt​U⊤\.Q\_\{t\+1\}^\{\\top\}C\_\{t\}Q\_\{t\+1\}=UQ\_\{t\}^\{\\top\}C\_\{t\}Q\_\{t\}U^\{\\top\}=UM\_\{t\}U^\{\\top\}\.Using the definitions ofVtV\_\{t\}andVt\+1∗V\_\{t\+1\}^\{\*\},

Vt=diag⁡\(Qt⊤​Ct​Qt\)=diag⁡\(Mt\),Vt\+1∗=diag⁡\(Qt\+1⊤​Ct\+1​Qt\+1\)\.V\_\{t\}=\\operatorname\{diag\}\\left\(Q\_\{t\}^\{\\top\}C\_\{t\}Q\_\{t\}\\right\)=\\operatorname\{diag\}\\left\(M\_\{t\}\\right\),\\quad V\_\{t\+1\}^\{\*\}=\\operatorname\{diag\}\\left\(Q\_\{t\+1\}^\{\\top\}C\_\{t\+1\}Q\_\{t\+1\}\\right\)\.Hence,

‖Vt−Vt\+1∗‖F=\\displaystyle\\left\\\|V\_\{t\}\-V\_\{t\+1\}^\{\*\}\\right\\\|\_\{F\}=‖diag⁡\(Qt⊤​Ct​Qt\)−diag⁡\(Qt\+1⊤​Ct\+1​Qt\+1\)‖F\\displaystyle\\left\\\|\\operatorname\{diag\}\\left\(Q\_\{t\}^\{\\top\}C\_\{t\}Q\_\{t\}\\right\)\-\\operatorname\{diag\}\\left\(Q\_\{t\+1\}^\{\\top\}C\_\{t\+1\}Q\_\{t\+1\}\\right\)\\right\\\|\_\{F\}≤\\displaystyle\\leq‖diag⁡\(Qt⊤​Ct​Qt\)−diag⁡\(Qt\+1⊤​Ct​Qt\+1\)‖F\\displaystyle\\left\\\|\\operatorname\{diag\}\\left\(Q\_\{t\}^\{\\top\}C\_\{t\}Q\_\{t\}\\right\)\-\\operatorname\{diag\}\\left\(Q\_\{t\+1\}^\{\\top\}C\_\{t\}Q\_\{t\+1\}\\right\)\\right\\\|\_\{F\}\+‖diag⁡\(Qt\+1⊤​Ct​Qt\+1\)−diag⁡\(Qt\+1⊤​Ct\+1​Qt\+1\)‖F\.\\displaystyle\+\\left\\\|\\operatorname\{diag\}\\left\(Q\_\{t\+1\}^\{\\top\}C\_\{t\}Q\_\{t\+1\}\\right\)\-\\operatorname\{diag\}\\left\(Q\_\{t\+1\}^\{\\top\}C\_\{t\+1\}Q\_\{t\+1\}\\right\)\\right\\\|\_\{F\}\.For the first term, usingMt=Qt⊤​Ct​QtM\_\{t\}=Q\_\{t\}^\{\\top\}C\_\{t\}Q\_\{t\}andQt\+1⊤​Ct​Qt\+1=U​Mt​U⊤Q\_\{t\+1\}^\{\\top\}C\_\{t\}Q\_\{t\+1\}=UM\_\{t\}U^\{\\top\},

‖diag⁡\(Qt⊤​Ct​Qt\)−diag⁡\(Qt\+1⊤​Ct​Qt\+1\)‖F=‖diag⁡\(Mt\)−diag⁡\(U​Mt​U⊤\)‖F\.\\left\\\|\\operatorname\{diag\}\\left\(Q\_\{t\}^\{\\top\}C\_\{t\}Q\_\{t\}\\right\)\-\\operatorname\{diag\}\\left\(Q\_\{t\+1\}^\{\\top\}C\_\{t\}Q\_\{t\+1\}\\right\)\\right\\\|\_\{F\}=\\left\\\|\\operatorname\{diag\}\\left\(M\_\{t\}\\right\)\-\\operatorname\{diag\}\\left\(UM\_\{t\}U^\{\\top\}\\right\)\\right\\\|\_\{F\}\.The diagonal projection is non\-expansive in Frobenius norm, so

‖diag⁡\(Mt\)−diag⁡\(U​Mt​U⊤\)‖F≤‖Mt−U​Mt​U⊤‖F\\left\\\|\\operatorname\{diag\}\\left\(M\_\{t\}\\right\)\-\\operatorname\{diag\}\\left\(UM\_\{t\}U^\{\\top\}\\right\)\\right\\\|\_\{F\}\\leq\\left\\\|M\_\{t\}\-UM\_\{t\}U^\{\\top\}\\right\\\|\_\{F\}Now,

Mt−U​Mt​U⊤=\(I−U\)​Mt​U⊤\+Mt​\(I−U⊤\)\.M\_\{t\}\-UM\_\{t\}U^\{\\top\}=\(I\-U\)M\_\{t\}U^\{\\top\}\+M\_\{t\}\\left\(I\-U^\{\\top\}\\right\)\.Therefore,

‖Mt−U​Mt​U⊤‖F\\displaystyle\\left\\\|M\_\{t\}\-UM\_\{t\}U^\{\\top\}\\right\\\|\_\{F\}≤‖\(I−U\)​Mt​U⊤‖F\+‖Mt​\(I−U⊤\)‖F\\displaystyle\\leq\\left\\\|\(I\-U\)M\_\{t\}U^\{\\top\}\\right\\\|\_\{F\}\+\\left\\\|M\_\{t\}\\left\(I\-U^\{\\top\}\\right\)\\right\\\|\_\{F\}≤‖I−U‖F​‖Mt‖2\+‖Mt‖2​‖I−U⊤‖F\.\\displaystyle\\leq\\\|I\-U\\\|\_\{F\}\\left\\\|M\_\{t\}\\right\\\|\_\{2\}\+\\left\\\|M\_\{t\}\\right\\\|\_\{2\}\\left\\\|I\-U^\{\\top\}\\right\\\|\_\{F\}\.Since‖I−U⊤‖F=‖I−U‖F\\left\\\|I\-U^\{\\top\}\\right\\\|\_\{F\}=\\\|I\-U\\\|\_\{F\}, and orthogonal similarity preserves spectral norm,

‖Mt‖2=‖Qt⊤​Ct​Qt‖2=‖Ct‖2,\\left\\\|M\_\{t\}\\right\\\|\_\{2\}=\\left\\\|Q\_\{t\}^\{\\top\}C\_\{t\}Q\_\{t\}\\right\\\|\_\{2\}=\\left\\\|C\_\{t\}\\right\\\|\_\{2\},we obtain

‖Mt−U​Mt​U⊤‖F≤2​‖Ct‖2​‖I−U‖F\.\\left\\\|M\_\{t\}\-UM\_\{t\}U^\{\\top\}\\right\\\|\_\{F\}\\leq 2\\left\\\|C\_\{t\}\\right\\\|\_\{2\}\\\|I\-U\\\|\_\{F\}\.For the second term, observe that

‖diag⁡\(Qt\+1⊤​Ct​Qt\+1\)−diag⁡\(Qt\+1⊤​Ct\+1​Qt\+1\)‖F\\displaystyle\\left\\\|\\operatorname\{diag\}\\left\(Q\_\{t\+1\}^\{\\top\}C\_\{t\}Q\_\{t\+1\}\\right\)\-\\operatorname\{diag\}\\left\(Q\_\{t\+1\}^\{\\top\}C\_\{t\+1\}Q\_\{t\+1\}\\right\)\\right\\\|\_\{F\}=‖diag⁡\(Qt\+1⊤​\(Ct−Ct\+1\)​Qt\+1\)‖F\.\\displaystyle=\\left\\\|\\operatorname\{diag\}\\left\(Q\_\{t\+1\}^\{\\top\}\\left\(C\_\{t\}\-C\_\{t\+1\}\\right\)Q\_\{t\+1\}\\right\)\\right\\\|\_\{F\}\.Again, diagonal projection is non\-expansive in Frobenius norm, and Frobenius norm is orthogonally invariant\. Hence,

‖diag⁡\(Qt\+1⊤​\(Ct−Ct\+1\)​Qt\+1\)‖F≤‖Ct−Ct\+1‖F\.\\left\\\|\\operatorname\{diag\}\\left\(Q\_\{t\+1\}^\{\\top\}\\left\(C\_\{t\}\-C\_\{t\+1\}\\right\)Q\_\{t\+1\}\\right\)\\right\\\|\_\{F\}\\leq\\left\\\|C\_\{t\}\-C\_\{t\+1\}\\right\\\|\_\{F\}\.Combining the preceding bounds gives

‖Vt−Vt\+1∗‖F≤2​‖Ct‖2​‖I−U‖F\+‖Ct\+1−Ct‖F\.\\left\\\|V\_\{t\}\-V\_\{t\+1\}^\{\*\}\\right\\\|\_\{F\}\\leq 2\\left\\\|C\_\{t\}\\right\\\|\_\{2\}\\\|I\-U\\\|\_\{F\}\+\\left\\\|C\_\{t\+1\}\-C\_\{t\}\\right\\\|\_\{F\}\.Finally,

‖Vscale−Vt\+1∗‖F\\displaystyle\\left\\\|V\_\{\\text\{scale \}\}\-V\_\{t\+1\}^\{\*\}\\right\\\|\_\{F\}=‖γ​Vt−Vt\+1∗‖F\\displaystyle=\\left\\\|\\gamma V\_\{t\}\-V\_\{t\+1\}^\{\*\}\\right\\\|\_\{F\}≤\(1−γ\)​‖Vt‖F\+‖Vt−Vt\+1∗‖F\\displaystyle\\leq\(1\-\\gamma\)\\left\\\|V\_\{t\}\\right\\\|\_\{F\}\+\\left\\\|V\_\{t\}\-V\_\{t\+1\}^\{\*\}\\right\\\|\_\{F\}≤\(1−γ\)​‖Vt‖F\+2​‖Ct‖2​‖I−U‖F\+‖Ct\+1−Ct‖F\\displaystyle\\leq\(1\-\\gamma\)\\left\\\|V\_\{t\}\\right\\\|\_\{F\}\+2\\left\\\|C\_\{t\}\\right\\\|\_\{2\}\\\|I\-U\\\|\_\{F\}\+\\left\\\|C\_\{t\+1\}\-C\_\{t\}\\right\\\|\_\{F\}If‖Ct\+1−Ct‖F≤Δt\\left\\\|C\_\{t\+1\}\-C\_\{t\}\\right\\\|\_\{F\}\\leq\\Delta\_\{t\}, then the simplified bound follows immediately

‖Vscale−Vt\+1∗‖F≤\(1−γ\)​‖Vt‖F\+2​‖Ct‖2​‖I−U‖F\+Δt\.\\left\\\|V\_\{\\text\{scale \}\}\-V\_\{t\+1\}^\{\*\}\\right\\\|\_\{F\}\\leq\(1\-\\gamma\)\\left\\\|V\_\{t\}\\right\\\|\_\{F\}\+2\\left\\\|C\_\{t\}\\right\\\|\_\{2\}\\\|I\-U\\\|\_\{F\}\+\\Delta\_\{t\}\.This completes the proof\. ∎

### A\.2Proof of Theorem[4\.2](https://arxiv.org/html/2608.29448#S4.Thmtheorem2)

###### Proof\.

Let

ak=sk⊤​Kk​sk,ck=sk⊤​yk\.\\displaystyle a\_\{k\}=s\_\{k\}^\{\\top\}K\_\{k\}s\_\{k\},\\qquad c\_\{k\}=s\_\{k\}^\{\\top\}y\_\{k\}\.BecauseKk≻0K\_\{k\}\\succ 0andsk≠0s\_\{k\}\\neq 0, we haveak\>0a\_\{k\}\>0\. By the curvature assumption,ck\>0c\_\{k\}\>0\. Consider

J⁡\(τ\)=\[log⁡cksk⊤​\(τ​Kk\)​sk\]2=\[log⁡ckτ​ak\]2\.\\displaystyle J\(\\tau\)=\\left\[\\log\\frac\{c\_\{k\}\}\{s\_\{k\}^\{\\top\}\(\\tau K\_\{k\}\)s\_\{k\}\}\\right\]^\{2\}=\\left\[\\log\\frac\{c\_\{k\}\}\{\\tau a\_\{k\}\}\\right\]^\{2\}\.Let

u=log⁡ckak,z=log⁡τ\.\\displaystyle u=\\log\\frac\{c\_\{k\}\}\{a\_\{k\}\},\\qquad z=\\log\\tau\.The constraint0<τ≤10<\\tau\\leq 1is equivalent toz≤0z\\leq 0\. Therefore

J⁡\(τ\)=\(u−z\)2\.\\displaystyle J\(\\tau\)=\(u\-z\)^\{2\}\.The unique minimizer overz≤0z\\leq 0is the Euclidean projection ofuuonto\(−∞,0\]\(\-\\infty,0\], namely

z⋆=min⁡\{u,0\}\.\\displaystyle z^\{\\star\}=\\min\\\{u,0\\\}\.Hence

τk⋆=exp⁡\(z⋆\)=min⁡\{1,ckak\}=min⁡\{1,sk⊤​yksk⊤​Kk​sk\}\.\\displaystyle\\tau\_\{k\}^\{\\star\}=\\exp\(z^\{\\star\}\)=\\min\\left\\\{1,\\frac\{c\_\{k\}\}\{a\_\{k\}\}\\right\\\}=\\min\\left\\\{1,\\frac\{s\_\{k\}^\{\\top\}y\_\{k\}\}\{s\_\{k\}^\{\\top\}K\_\{k\}s\_\{k\}\}\\right\\\}\.This proves that it is the unique clipped log\-secant matching solution\. Finally, suppose

yk=Hk​sk\+rk\.\\displaystyle y\_\{k\}=H\_\{k\}s\_\{k\}\+r\_\{k\}\.Then

sk⊤​yk=sk⊤​Hk​sk\+sk⊤​rk\.\\displaystyle s\_\{k\}^\{\\top\}y\_\{k\}=s\_\{k\}^\{\\top\}H\_\{k\}s\_\{k\}\+s\_\{k\}^\{\\top\}r\_\{k\}\.When the unclipped solution is active,τk=ck/ak\\tau\_\{k\}=c\_\{k\}/a\_\{k\}, and therefore

sk⊤​\(τk​Kk\)​sk=τk​ak=ck=sk⊤​Hk​sk\+sk⊤​rk\.\\displaystyle s\_\{k\}^\{\\top\}\(\\tau\_\{k\}K\_\{k\}\)s\_\{k\}=\\tau\_\{k\}a\_\{k\}=c\_\{k\}=s\_\{k\}^\{\\top\}H\_\{k\}s\_\{k\}\+s\_\{k\}^\{\\top\}r\_\{k\}\.Thus, under an exact secant relation, the scaled Kronecker quadratic form matches the true directional curvature exactly; under an approximate secant relation, the mismatch is precisely the directional secant error\. If the unconstrained solution exceeds one, the constraint0<τ≤10<\\tau\\leq 1selects the closest admissible log\-scale match, namelyτk=1\\tau\_\{k\}=1\. ∎

## Appendix BAlgorithms Pseudocode

To connect ourSS\-eSOAPpresented in Algorithm[2](https://arxiv.org/html/2608.29448#alg2)with the established algorithms in the second\-order optimizer literature, we provide the following technical context\.

### B\.1Relation to Idealized and Warm\-Started Shampoo

Our algorithm can be viewed as a curvature\-corrected approximation of Idealized Eigenvalue\-Corrected Shampoo \(Algorithm 1 in\[[7](https://arxiv.org/html/2608.29448#bib.bib8)\]\)\.

In the idealized setting, the exact eigendecomposition of the preconditionersLtL\_\{t\}andRtR\_\{t\}is computed at every step to perform the updateWt\+1=Wt−ηLt−1/2GtRt−1/2W\_\{t\+1\}=W\_\{t\}\-\\eta L\_\{t\}^\{\-1/2\}G\_\{t\}R\_\{t\}^\{\-1/2\}\. Standard SOAP approximates this efficiently using Warm\-started QR iteration \(Algorithm 4 in\[[7](https://arxiv.org/html/2608.29448#bib.bib8)\]\), which performs a single step of power iteration followed by QR decomposition to update the eigenbasis\.

SS\-eSOAPimproves upon this approximation in two critical ways\.

First, instead of a fixed schedule or a single QR step, we use the Purifying trigger \(Step 2 in Algorithm[2](https://arxiv.org/html/2608.29448#alg2)\)\. We monitor the off\-diagonal mass of the preconditioner in the current basis,ρ⁡\(A,Q\)\\rho\(A,Q\), which serves as a proxy for the approximation error of the stale eigenbasis\. We only trigger a computationally expensive re\-decomposition when this error exceedsτtrigger\\tau\_\{\\text\{trigger\}\}, ensuring the eigenbasis is refined only when the curvature landscape shifts significantly\.

Second, we introduce the scalarτk\\tau\_\{k\}\(Step 4\), which acts as a global correction factor for the eigenvalues, ensuring the trace of the inverse Hessian approximation matches the secant condition along the most recent displacement\.

### B\.2Relation to EShampoo and Grafting

We highlight two key advantages of our strategy over EShampoo \(Algorithm 2 in\[[7](https://arxiv.org/html/2608.29448#bib.bib8)\]\) and Shampoo with Adam grafting\.

Stability via Downscaling vs\. Reprojection\.EShampoo typically employs a fixed eigenbasis computation frequencyFF\. When the basis is updated, standard approaches often reproject the optimizer state \(e\.g\., the second momentVtV\_\{t\}\) into the new basis via similarity transformation\. As detailed in our analysis, this is numerically unstable: a significant rotation of the eigenbasis can map historical low\-variance directions to high\-curvature directions in the new basis, causing overly large updates\.SS\-eSOAPreplaces reprojection with variance downscaling \(Step 2, Soft Reset\)\. By keeping the historicalVtV\_\{t\}but reducing its magnitude byγ\\gamma, we effectively increase the damping uniformly, preserving optimization history while conservatively forgetting the precise directional variance that is no longer valid\.

Intrinsic vs\. Extrinsic Scaling\.Shampoo with Adam grafting attempts to stabilize second\-order updates by forcing the update magnitude or direction to align with Adam\. This is an ad\-hoc extrinsic constraint\. In contrast,SS\-eSOAPachieves intrinsic stability via the self\-scaling factorτk\\tau\_\{k\}\. This factor naturally damps the update when the local curvature is high \(large denominator in \([3](https://arxiv.org/html/2608.29448#S3.E3)\)\) and accelerates it in flat regions without the need for grafting or manual tuning of the epsilon parameter\.

Algorithm 2SS\-eSOAP, full pseudocodeRequire:

W0∈ℝm×nW\_\{0\}\\in\\mathbb\{R\}^\{m\\times n\},

η\>0\\eta\>0,

β1,β2∈\(0,1\)\\beta\_\{1\},\\beta\_\{2\}\\in\(0,1\),

ϵ\>0\\epsilon\>0,

λ≥0\\lambda\\geq 0,

τmin∈\(0,1\]\\tau\_\{\\min\}\\in\(0,1\]\.

Require:

IcheckI\_\{\\mathrm\{check\}\},

τtrigger\\tau\_\{\\mathrm\{trigger\}\},

TwarmT\_\{\\mathrm\{warm\}\}\.

Initialize:

M~0=V0=𝟎\\widetilde\{M\}\_\{0\}=V\_\{0\}=\\mathbf\{0\},

L0=ϵ​ImL\_\{0\}=\\epsilon I\_\{m\},

R0=ϵ​InR\_\{0\}=\\epsilon I\_\{n\},

QL=ImQ\_\{L\}=I\_\{m\},

QR=InQ\_\{R\}=I\_\{n\}\.

for

t=1,…,Tt=1,\\ldots,Tdo

1\. Gradient and Kronecker factors

Gt←∇Wℒ​\(Wt−1\)G\_\{t\}\\leftarrow\\nabla\_\{W\}\\mathcal\{L\}\(W\_\{t\-1\}\)
Lt←β2​Lt−1\+\(1−β2\)​Gt​Gt⊤L\_\{t\}\\leftarrow\\beta\_\{2\}L\_\{t\-1\}\+\(1\-\\beta\_\{2\}\)G\_\{t\}G\_\{t\}^\{\\top\},

Rt←β2​Rt−1\+\(1−β2\)​Gt⊤​GtR\_\{t\}\\leftarrow\\beta\_\{2\}R\_\{t\-1\}\+\(1\-\\beta\_\{2\}\)G\_\{t\}^\{\\top\}G\_\{t\}
2\. Adaptive eigenbasis update

if

t\>Twarmt\>T\_\{\\mathrm\{warm\}\}and

tmodIcheck=0t\\bmod I\_\{\\mathrm\{check\}\}=0then

ρ⁡\(A,Q\)←‖Q⊤​A​Q−diag⁡\(Q⊤​A​Q\)‖F‖Q⊤​A​Q‖F\+ϵ\\displaystyle\\rho\(A,Q\)\\leftarrow\\frac\{\\\|Q^\{\\top\}AQ\-\\operatorname\{diag\}\(Q^\{\\top\}AQ\)\\\|\_\{F\}\}\{\\\|Q^\{\\top\}AQ\\\|\_\{F\}\+\\epsilon\},

ρt←max⁡\{ρ⁡\(Lt,QL\),ρ⁡\(Rt,QR\)\}\\hskip 9\.24994pt\\rho\_\{t\}\\leftarrow\\max\\\{\\rho\(L\_\{t\},Q\_\{L\}\),\\rho\(R\_\{t\},Q\_\{R\}\)\\\}
if

ρt\>τtrigger\\rho\_\{t\}\>\\tau\_\{\\mathrm\{trigger\}\}then

\(QLold,QRold\)←\(QL,QR\)\(Q\_\{L\}^\{\\mathrm\{old\}\},Q\_\{R\}^\{\\mathrm\{old\}\}\)\\leftarrow\(Q\_\{L\},Q\_\{R\}\)
QL←eigvec⁡\(Lt\)Q\_\{L\}\\leftarrow\\operatorname\{eigvec\}\(L\_\{t\}\),

QR←eigvec⁡\(Rt\)Q\_\{R\}\\leftarrow\\operatorname\{eigvec\}\(R\_\{t\}\)
M~t−1←\(QL⊤​QLold\)​M~t−1​\(\(QRold\)⊤​QR\)\\widetilde\{M\}\_\{t\-1\}\\leftarrow\(Q\_\{L\}^\{\\top\}Q\_\{L\}^\{\\mathrm\{old\}\}\)\\widetilde\{M\}\_\{t\-1\}\(\(Q\_\{R\}^\{\\mathrm\{old\}\}\)^\{\\top\}Q\_\{R\}\)
γt←0\.25\\gamma\_\{t\}\\leftarrow 0\.25if

ρt\>0\.8\\rho\_\{t\}\>0\.8,

0\.50\.5if

ρt\>0\.5\\rho\_\{t\}\>0\.5, and

0\.750\.75otherwise

Vt−1←γt​Vt−1V\_\{t\-1\}\\leftarrow\\gamma\_\{t\}V\_\{t\-1\}
endif

endif

3\. Eigenspace statistics

G~t←QL⊤​Gt​QR\\widetilde\{G\}\_\{t\}\\leftarrow Q\_\{L\}^\{\\top\}G\_\{t\}Q\_\{R\}
M~t←β1​M~t−1\+\(1−β1\)​G~t\\widetilde\{M\}\_\{t\}\\leftarrow\\beta\_\{1\}\\widetilde\{M\}\_\{t\-1\}\+\(1\-\\beta\_\{1\}\)\\widetilde\{G\}\_\{t\},

Vt←β2​Vt−1\+\(1−β2\)​G~t⊙2\\hskip 9\.24994ptV\_\{t\}\\leftarrow\\beta\_\{2\}V\_\{t\-1\}\+\(1\-\\beta\_\{2\}\)\\widetilde\{G\}\_\{t\}^\{\\odot 2\}
4\. Self\-scaling correction

if

t=1t=1then

τt←1\\tau\_\{t\}\\leftarrow 1
else

St←Wt−1−Wt−2S\_\{t\}\\leftarrow W\_\{t\-1\}\-W\_\{t\-2\},

Yt←Gt−Gt−1\\hskip 9\.24994ptY\_\{t\}\\leftarrow G\_\{t\}\-G\_\{t\-1\}
ct←Tr⁡\(Yt⊤​St\)c\_\{t\}\\leftarrow\\operatorname\{Tr\}\(Y\_\{t\}^\{\\top\}S\_\{t\}\),

at←Tr⁡\(St⊤​Lt−1​St​Rt−1\)\\hskip 9\.24994pta\_\{t\}\\leftarrow\\operatorname\{Tr\}\(S\_\{t\}^\{\\top\}L\_\{t\}^\{\-1\}S\_\{t\}R\_\{t\}^\{\-1\}\)
if

ct\>0c\_\{t\}\>0and

at\>0a\_\{t\}\>0then

τt←min⁡\{1,max⁡\{τmin,ct/at\}\}\\displaystyle\\tau\_\{t\}\\leftarrow\\min\\\{1,\\max\\\{\\tau\_\{\\min\},c\_\{t\}/a\_\{t\}\\\}\\\}
else

τt←1\\tau\_\{t\}\\leftarrow 1
endif

endif

5\. Parameter update

M^t←M~t/\(1−β1t\)\\widehat\{M\}\_\{t\}\\leftarrow\\widetilde\{M\}\_\{t\}/\(1\-\\beta\_\{1\}^\{t\}\),

V^t←Vt/\(1−β2t\)\\hskip 9\.24994pt\\widehat\{V\}\_\{t\}\\leftarrow V\_\{t\}/\(1\-\\beta\_\{2\}^\{t\}\)
U~t←τt−1/2M^t⊘\(V^t\+ϵ\)\\widetilde\{U\}\_\{t\}\\leftarrow\\tau\_\{t\}^\{\-1/2\}\\widehat\{M\}\_\{t\}\\oslash\(\\sqrt\{\\widehat\{V\}\_\{t\}\}\+\\epsilon\)
Ut←QL​U~t​QR⊤U\_\{t\}\\leftarrow Q\_\{L\}\\widetilde\{U\}\_\{t\}Q\_\{R\}^\{\\top\},

Wt←\(1−η​λ\)​Wt−1−η​Ut\\hskip 9\.24994ptW\_\{t\}\\leftarrow\(1\-\\eta\\lambda\)W\_\{t\-1\}\-\\eta U\_\{t\}
endfor

### B\.3Efficient Computation ofτk\\tau\_\{k\}

Similar to \([2](https://arxiv.org/html/2608.29448#S2.E2)\), we avoid direct computation of the inverse preconditionerHk−1H\_\{k\}^\{\-1\}, as it isO⁡\(n3\)O\(n^\{3\}\)and can be numerically unstable when the curvature matrices are ill\-conditioned\. In our setting, the preconditioner admits a Kronecker\-factored eigendecomposition

Hk\\displaystyle H\_\{k\}=Lk⊗Rk,\\displaystyle=L\_\{k\}\\otimes R\_\{k\},Lk\\displaystyle L\_\{k\}=QL,k​diag⁡\(DL,k\)​QL,k⊤,\\displaystyle=Q\_\{L,k\}\\operatorname\{diag\}\(D\_\{L,k\}\)\\,Q\_\{L,k\}^\{\\top\},Rk\\displaystyle R\_\{k\}=QR,k​diag⁡\(DR,k\)​QR,k⊤,\\displaystyle=Q\_\{R,k\}\\operatorname\{diag\}\(D\_\{R,k\}\)\\,Q\_\{R,k\}^\{\\top\},whereQL,kQ\_\{L,k\}andQR,kQ\_\{R,k\}are orthogonal matrices andDL,kD\_\{L,k\},DR,kD\_\{R,k\}contain the corresponding eigenvalues\.

By rotating tensors into the eigenspaces ofLkL\_\{k\}andRkR\_\{k\}, the application ofHk−1H\_\{k\}^\{\-1\}reduces to elementwise scaling by the inverse eigenvaluesDL,k−1D\_\{L,k\}^\{\-1\}andDR,k−1D\_\{R,k\}^\{\-1\}\. Consequently, all terms required to computeτk\\tau\_\{k\}can be evaluated using matrix\-vector products and elementwise operations, without explicitly forming inverse matrices\.

Specifically, defining the rotated quantities

S~k\\displaystyle\\tilde\{S\}\_\{k\}=QL,k⊤​Sk​QR,k,Y~k=QL,k⊤​Yk​QR,k,\\displaystyle=Q\_\{L,k\}^\{\\top\}S\_\{k\}Q\_\{R,k\},\\quad\\tilde\{Y\}\_\{k\}=Q\_\{L,k\}^\{\\top\}Y\_\{k\}Q\_\{R,k\},the denominator in \([3](https://arxiv.org/html/2608.29448#S3.E3)\) can be written as

Tr⁡\(Sk⊤​Lk−1​Sk​Rk−1\)=∑i,jS~k,i​j2DL,k,i​DR,k,j,\\operatorname\{Tr\}\(S\_\{k\}^\{\\top\}L\_\{k\}^\{\-1\}S\_\{k\}R\_\{k\}^\{\-1\}\)=\\sum\_\{i,j\}\\frac\{\\tilde\{S\}\_\{k,ij\}^\{2\}\}\{D\_\{L,k,i\}\\,D\_\{R,k,j\}\},\(5\)which is computable inO⁡\(n\)O\(n\)time\. On the other hand, the numerator simplifies to

Tr⁡\(Yk⊤​Sk\)=∑i,jY~k,i​j​S~k,i​j\.\\displaystyle\\operatorname\{Tr\}\(Y\_\{k\}^\{\\top\}S\_\{k\}\)=\\sum\_\{i,j\}\\tilde\{Y\}\_\{k,ij\}\\,\\tilde\{S\}\_\{k,ij\}\.Thus, the scaling parameter is evaluated as

τk=min⁡\{1,∑i,jY~k,i​j​S~k,i​j∑i,jS~k,i​j2/\(DL,k,i​DR,k,j\)\},\\tau\_\{k\}=\\min\\left\\\{1,\\;\\frac\{\\sum\_\{i,j\}\\tilde\{Y\}\_\{k,ij\}\\,\\tilde\{S\}\_\{k,ij\}\}\{\\sum\_\{i,j\}\\tilde\{S\}\_\{k,ij\}^\{2\}/\(D\_\{L,k,i\}D\_\{R,k,j\}\)\}\\right\\\},\(6\)without ever formingHk−1H\_\{k\}^\{\-1\}\. This formulation preserves numerical stability while maintaining the computational efficiency required for large\-scale optimization\.

## Appendix CArchitectures and Additional Design Choices

### C\.1PirateNet

PirateNet\[[30](https://arxiv.org/html/2608.29448#bib.bib30)\]aims to enable stable and efficient training of deep PINN models\. It first transforms input coordinates𝐱\\mathbf\{x\}into a high\-dimensional feature space using random Fourier features

Φ⁡\(𝐱\)=\[cos⁡\(𝐁𝐱\)sin⁡\(𝐁𝐱\)\],\\Phi\(\\mathbf\{x\}\)=\\left\[\\begin\{array\}\[\]\{c\}\\cos\(\\mathbf\{Bx\}\)\\\\ \\sin\(\\mathbf\{Bx\}\)\\end\{array\}\\right\],where𝐁∈ℝm×d\\mathbf\{B\}\\in\\mathbb\{R\}^\{m\\times d\}has entries sampled i\.i\.d\. from𝒩⁡\(0,s2\)\\mathcal\{N\}\(0,s^\{2\}\)with user\-specifieds\>0s\>0\. This embedding mitigates spectral bias in PINNs by improving the eigenfunction frequency of the Neural Tangent Kernel, enabling better learning of high\-frequency components and multiscale features\. The embedded coordinates are processed through two dense layers that act as gates𝐔=σ⁡\(𝐖1​Φ​\(𝐱\)\+𝐛1\)\\mathbf\{U\}=\\sigma\\left\(\\mathbf\{W\}\_\{1\}\\Phi\(\\mathbf\{x\}\)\+\\mathbf\{b\}\_\{1\}\\right\)and𝐕=σ⁡\(𝐖2​Φ​\(𝐱\)\+𝐛2\)\\mathbf\{V\}=\\sigma\\left\(\\mathbf\{W\}\_\{2\}\\Phi\(\\mathbf\{x\}\)\+\\mathbf\{b\}\_\{2\}\\right\), whereσ\\sigmais a point\-wise activation function\. This gating mechanism is essentially the same as in modified MLP\. Let𝐱\(1\)=Φ⁡\(𝐱\)\\mathbf\{x\}^\{\(1\)\}=\\Phi\(\\mathbf\{x\}\)and𝐱\(l\)\\mathbf\{x\}^\{\(l\)\}be the input to thell\-th block\. Each block performs

𝐟\(l\)=σ\(𝐖1\(l\)𝐱\(l\)\+𝐛1\(l\)\),𝐳1\(l\)=𝐟\(l\)⊙𝐔\+\(1−𝐟\(l\)\)⊙𝐕,\\displaystyle\\mathbf\{f\}^\{\(l\)\}=\\sigma\\left\(\\mathbf\{W\}\_\{1\}^\{\(l\)\}\\mathbf\{x\}^\{\(l\)\}\+\\mathbf\{b\}\_\{1\}^\{\(l\)\}\\right\),\\quad\\mathbf\{z\}\_\{1\}^\{\(l\)\}=\\mathbf\{f\}^\{\(l\)\}\\odot\\mathbf\{U\}\+\\left\(1\-\\mathbf\{f\}^\{\(l\)\}\\right\)\\odot\\mathbf\{V\},𝐠\(l\)=σ\(𝐖2\(l\)𝐳1\(l\)\+𝐛2\(l\)\),𝐳2\(l\)=𝐠\(l\)⊙𝐔\+\(1−𝐠\(l\)\)⊙𝐕,\\displaystyle\\mathbf\{g\}^\{\(l\)\}=\\sigma\\left\(\\mathbf\{W\}\_\{2\}^\{\(l\)\}\\mathbf\{z\}\_\{1\}^\{\(l\)\}\+\\mathbf\{b\}\_\{2\}^\{\(l\)\}\\right\),\\quad\\mathbf\{z\}\_\{2\}^\{\(l\)\}=\\mathbf\{g\}^\{\(l\)\}\\odot\\mathbf\{U\}\+\\left\(1\-\\mathbf\{g\}^\{\(l\)\}\\right\)\\odot\\mathbf\{V\},𝐡\(l\)=σ\(𝐖3\(l\)𝐳2\(l\)\+𝐛3\(l\)\),𝐱\(l\+1\)=α\(l\)𝐡\(l\)\+\(1−α\(l\)\)𝐱\(l\)\.\\displaystyle\\mathbf\{h\}^\{\(l\)\}=\\sigma\\left\(\\mathbf\{W\}\_\{3\}^\{\(l\)\}\\mathbf\{z\}\_\{2\}^\{\(l\)\}\+\\mathbf\{b\}\_\{3\}^\{\(l\)\}\\right\),\\quad\\mathbf\{x\}^\{\(l\+1\)\}=\\alpha^\{\(l\)\}\\mathbf\{h\}^\{\(l\)\}\+\\left\(1\-\\alpha^\{\(l\)\}\\right\)\\mathbf\{x\}^\{\(l\)\}\.Each block comprises three dense layers with dual gating operations and an adaptive residual connection\. The trainableα\(l\)\\alpha^\{\(l\)\}parameters control block nonlinearity:α\(l\)=0\\alpha^\{\(l\)\}=0yields an identity mapping, whileα\(l\)=1\\alpha^\{\(l\)\}=1produces fully nonlinear transformation\. The final output of a PirateNet ofLLresidual blocks is given by𝐮θ=𝐖\(L\+1\)​𝐱\(L\)\\mathbf\{u\}\_\{\\theta\}=\\mathbf\{W\}^\{\(L\+1\)\}\\mathbf\{x\}^\{\(L\)\}\. Following the design choices made by\[[30](https://arxiv.org/html/2608.29448#bib.bib30)\], we initializeα\(l\)=0\\alpha^\{\(l\)\}=0, making the initial output a linear combination of first\-layer embeddings\. This initialization strategy mitigates training difficulties in deep networks by starting with effectively shallow architecture and gradually increasing depth through learnedα\\alphavalues\. Additionally, the linear structure at initialization enables direct integration of prior solution data through least squares fitting

min𝐖⁡‖𝐖​Φ−𝐘‖22,\\min\_\{\\mathbf\{W\}\}\\\|\\mathbf\{W\}\\Phi\-\\mathbf\{Y\}\\\|\_\{2\}^\{2\},where𝐘\\mathbf\{Y\}represents available measurements\. This approach provides an optimal initial guess based on various data sources, including experimental measurements, boundary conditions, or linearized PDE solutions\.

Table 3:Comparison of optimization methods showing preconditioner types, storage and computational complexity for ann×nn\\times nweight matrix, wherePPis the parameter count andNNis the batch size of collocation points, and practical compatibility with mini\-batches and scalability with large neural networks\.
### C\.2Additional Details of Standard PDE Benchmarks

Burgers equation\.The 1D Burgers equation is defined as

ut\+u​ux=ν​ux​x,u\_\{t\}\+uu\_\{x\}=\\nu u\_\{xx\},whereuurepresents the velocity field, andν\\nuis the kinematic viscosity coefficient controlling the diffusion strength\. Here we set\(x,t\)∈Ω=\[−1,1\]×\[0,1\]\(x,t\)\\in\\Omega=\[\-1,1\]\\times\[0,1\], with initial and boundary conditions

u⁡\(x,0\)\\displaystyle u\(x,0\)=−sin⁡\(π​x\)\\displaystyle=\-\\sin\(\\pi x\)u⁡\(−1,t\)\\displaystyle u\(\-1,t\)=u⁡\(1,t\)=0,\\displaystyle=u\(1,t\)=0,and viscosity parameterν=0\.01/π\\nu=0\.01/\\pi\. Following\[[33](https://arxiv.org/html/2608.29448#bib.bib26)\], we enforce a hard constraint via Taylor expansion at the origin asU=−z\+z3\+z4​U1U=\-z\+z^\{3\}\+z^\{4\}U\_\{1\}, for an odd functionU1U\_\{1\}\.

2D Boussinesq equation\.For the 2D Boussinesq equation on the half plane, in vorticity form with the self\-similar ansatz, we get the following profile equations for\(Ω,U1,U2,Φ,Ψ\)\\left\(\\Omega,U\_\{1\},U\_\{2\},\\Phi,\\Psi\\right\)as in\[[35](https://arxiv.org/html/2608.29448#bib.bib29)\]:

Ω\+\(\(1\+λ\)​\(y1,y2\)T\+\(U1,U2\)T\)⋅∇Ω\\displaystyle\\Omega\+\\left\(\(1\+\\lambda\)\\left\(y\_\{1\},y\_\{2\}\\right\)^\{T\}\+\\left\(U\_\{1\},U\_\{2\}\\right\)^\{T\}\\right\)\\cdot\\nabla\\Omega=Φ\\displaystyle=\\Phi\(2\+∂y1U1\)​Φ\+\(\(1\+λ\)​\(y1,y2\)T\+\(U1,U2\)T\)⋅∇Φ\\displaystyle\\left\(2\+\\partial\_\{y\_\{1\}\}U\_\{1\}\\right\)\\Phi\+\\left\(\(1\+\\lambda\)\\left\(y\_\{1\},y\_\{2\}\\right\)^\{T\}\+\\left\(U\_\{1\},U\_\{2\}\\right\)^\{T\}\\right\)\\cdot\\nabla\\Phi=−∂y1U2Ψ\\displaystyle=\-\\partial\_\{y\_\{1\}\}U\_\{2\}\\Psi\(2\+∂y2U2\)​Ψ\+\(\(1\+λ\)​\(y1,y2\)T\+\(U1,U2\)T\)⋅∇Ψ\\displaystyle\\left\(2\+\\partial\_\{y\_\{2\}\}U\_\{2\}\\right\)\\Psi\+\\left\(\(1\+\\lambda\)\\left\(y\_\{1\},y\_\{2\}\\right\)^\{T\}\+\\left\(U\_\{1\},U\_\{2\}\\right\)^\{T\}\\right\)\\cdot\\nabla\\Psi=−∂y2U1Φ\\displaystyle=\-\\partial\_\{y\_\{2\}\}U\_\{1\}\\Phi∂y1U1\+∂y2U2=0,Ω=∂y1U2−∂y2U1,∂y1Ψ\\displaystyle\\partial\_\{y\_\{1\}\}U\_\{1\}\+\\partial\_\{y\_\{2\}\}U\_\{2\}=0,\\quad\\Omega=\\partial\_\{y\_\{1\}\}U\_\{2\}\-\\partial\_\{y\_\{2\}\}U\_\{1\},\\quad\\partial\_\{y\_\{1\}\}\\Psi=∂y2Φ,\\displaystyle=\\partial\_\{y\_\{2\}\}\\Phi,where\(Ω,U1,Φ\)\(\\Omega,U\_\{1\},\\Phi\)are odd and\(U2,Ψ\)\(U\_\{2\},\\Psi\)are even iny1y\_\{1\}and we are in the half planey2≥0y\_\{2\}\\geq 0\. For the boundary conditions, we impose a non\-penetration boundary conditionU2​\(y1,0\)=0U\_\{2\}\\left\(y\_\{1\},0\\right\)=0along with decaying weak asymptotics at the far field, with Dirichlet boundary conditionsΦ=Ψ=0\\Phi=\\Psi=0and Neumann boundary conditions for the velocity field∇\(U1,U2\)T=0\\nabla\\left\(U\_\{1\},U\_\{2\}\\right\)^\{T\}=0\.

For the nondegeneracy condition, we adapt the approach proposed in\[[33](https://arxiv.org/html/2608.29448#bib.bib26)\]and impose∂y1Ω⁡\(0,0\)=−1\\partial\_\{y\_\{1\}\}\\Omega\(0,0\)=\-1and use Taylor expansion to enforce a hard constrain to rule out the trivial solutionU=0U=0when using weak asymptotics\. We enforce∂1Ω⁡\(0,0\)=−1\\partial\_\{1\}\\Omega\(0,0\)=\-1andΩ\\Omegais odd inz1z\_\{1\}via a Taylor expansion asΩ=−z1\+z1​z2​Ω1\+z12​Ω2\\Omega=\-z\_\{1\}\+z\_\{1\}z\_\{2\}\\Omega\_\{1\}\+z\_\{1\}^\{2\}\\Omega\_\{2\}, whereΩ1,Ω2\\Omega\_\{1\},\\Omega\_\{2\}are even and odd functions inz1z\_\{1\}respectively\.

Wave equation\.We consider a one\-dimensional wave equation in the domainΩ=\[0,1\]×\[0,1\]\\Omega=\[0,1\]\\times\[0,1\]taking the form

ut​t\(x,t\)−4ux​x\(x,t\)=0,\(x,t\)∈\(0,1\)×\(0,1\),\\displaystyle u\_\{tt\}\(x,t\)\-4u\_\{xx\}\(x,t\)=0,\\quad\(x,t\)\\in\(0,1\)\\times\(0,1\),u\(0,t\)=u\(1,t\)=0,t∈\[0,1\],\\displaystyle u\(0,t\)=u\(1,t\)=0,\\quad t\\in\[0,1\],u\(x,0\)=sin\(πx\)\+12sin\(4πx\),x∈\[0,1\],\\displaystyle u\(x,0\)=\\sin\(\\pi x\)\+\\frac\{1\}\{2\}\\sin\(4\\pi x\),\\quad x\\in\[0,1\],ut\(x,0\)=0,x∈\[0,1\]\.\\displaystyle u\_\{t\}\(x,0\)=0,\\quad x\\in\[0,1\]\.whereuurepresents the wave amplitude, andccis the wave propagation speed, determined by the medium’s physical properties\. By d’Alembert’s formula, the solutionu⁡\(x,t\)u\(x,t\)is given by

u⁡\(x,t\)=sin⁡\(π​x\)​cos⁡\(2​π​t\)\+12​sin⁡\(4​π​x\)​cos⁡\(8​π​t\)\.u\(x,t\)=\\sin\(\\pi x\)\\cos\(2\\pi t\)\+\\frac\{1\}\{2\}\\sin\(4\\pi x\)\\cos\(8\\pi t\)\.
\(a\)Training loss convergence\.![Refer to caption](https://arxiv.org/html/2608.29448v1/wave_ep200000.png)\(b\)Final prediction and absolute error field\.

Figure 3:Wave Equation Benchmark\.\(a\) PDE residual loss trajectories \(log scale\) comparingSS\-eSOAP\(red\) against Adam, Muon, SOAP, and Purifying Shampoo over 100k iterations\.SS\-eSOAPachieves 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 of10−610^\{\-6\}\.\(a\)Training loss convergence\.![Refer to caption](https://arxiv.org/html/2608.29448v1/burgers_ep200000.png)\(b\)Final prediction and absolute error field\.

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\.Allen\-Cahn equation\.We investigate the one\-dimensional Allen\-Cahn equation with periodic boundary conditions

ut−0\.0001​ux​x\+5​u3−5​u=0,t∈\[0,1\],x∈\[−1,1\]\\displaystyle u\_\{t\}\-0\.0001u\_\{xx\}\+5u^\{3\}\-5u=0,\\quad t\\in\[0,1\],x\\in\[\-1,1\]u⁡\(0,x\)=x2​cos⁡\(π​x\)\\displaystyle u\(0,x\)=x^\{2\}\\cos\(\\pi x\)u⁡\(t,−1\)=u⁡\(t,1\),ux​\(t,−1\)=ux​\(t,1\)\\displaystyle u\(t,\-1\)=u\(t,1\),\\quad u\_\{x\}\(t,\-1\)=u\_\{x\}\(t,1\)whereuurepresents the order parameter \(e\.g\., concentration difference between two phases\),ϵ\\epsiloncontrols the interfacial width,aais the reaction rate coefficient, and the term \(u−u3u\-u^\{3\}\) drives the phase separation\.

Korteweg\-de Vries equation\.The one\-dimensional KdV equation is expressed as

ut\+ηuux\+μ2ux​x​x=0,t∈\(0,1\),x∈\(−1,1\),\\displaystyle u\_\{t\}\+\\eta uu\_\{x\}\+\\mu^\{2\}u\_\{xxx\}=0,\\quad t\\in\(0,1\),\\quad x\\in\(\-1,1\),u⁡\(x,0\)=cos⁡\(π​x\),\\displaystyle u\(x,0\)=\\cos\(\\pi x\),u⁡\(t,−1\)=u⁡\(t,1\),\\displaystyle u\(t,\-1\)=u\(t,1\),whereuurepresents the wave amplitude or water surface elevation, andη\\etagoverns the strength of the nonlinearity, whileμ\\mucontrols the dispersion level\. Under the KdV dynamics, this initial wave evolves into a series of solitary\-type waves\. Like in recent PINN optimization works, we adopt the classical parameters of the KdV equation, settingη=1\\eta=1andμ=0\.022\\mu=0\.022\.

![Refer to caption](https://arxiv.org/html/2608.29448v1/comparison_kdv.png)\(a\)Training loss convergence comparison\.
![Refer to caption](https://arxiv.org/html/2608.29448v1/kdv_ep300000.png)\(b\)Final prediction and absolute error field\.

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 of10−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\.Gray\-Scott equation\.The system is described by the following coupled PDEs:

ut=ϵ1Δu\+b1\(1−u\)−c1uv2,t∈\(0,2\),\(x,y\)∈\(−1,1\)2,\\displaystyle u\_\{t\}=\\epsilon\_\{1\}\\Delta u\+b\_\{1\}\(1\-u\)\-c\_\{1\}uv^\{2\},\\quad t\\in\(0,2\),\(x,y\)\\in\(\-1,1\)^\{2\},vt=ϵ2Δv−b2v\+c2uv2,t∈\(0,2\),\(x,y\)∈\(−1,1\)2,\\displaystyle v\_\{t\}=\\epsilon\_\{2\}\\Delta v\-b\_\{2\}v\+c\_\{2\}uv^\{2\},\\quad t\\in\(0,2\),\(x,y\)\\in\(\-1,1\)^\{2\},With periodic boundary conditions, the initial conditions are

u0​\(x,y\)=1−exp⁡\(−10​\(\(x\+0\.05\)2\+\(y\+0\.02\)2\)\),\\displaystyle u\_\{0\}\(x,y\)=1\-\\exp\\left\(\-10\\left\(\(x\+0\.05\)^\{2\}\+\(y\+0\.02\)^\{2\}\\right\)\\right\),v0​\(x,y\)=1−exp⁡\(−10​\(\(x−0\.05\)2\+\(y−0\.02\)2\)\)\.\\displaystyle v\_\{0\}\(x,y\)=1\-\\exp\\left\(\-10\\left\(\(x\-0\.05\)^\{2\}\+\(y\-0\.02\)^\{2\}\\right\)\\right\)\.whereuuandvvrepresent activator and inhibitor concentrations respectively,ε1\\varepsilon\_\{1\}andε2\\varepsilon\_\{2\}are diffusion coefficients, and\(b1,b2,c1,c2\)\\left\(b\_\{1\},b\_\{2\},c\_\{1\},c\_\{2\}\\right\)control reaction kinetics\. This system generates diverse spatial patterns including spots and stripes\. We set parametersϵ1=0\.2,ϵ2=0\.1,b1=40,b2=100\\epsilon\_\{1\}=0\.2,\\epsilon\_\{2\}=0\.1,b\_\{1\}=40,b\_\{2\}=100, andc1=c2=1,000c\_\{1\}=c\_\{2\}=1,000, which generates characteristic pattern formations\.

![Refer to caption](https://arxiv.org/html/2608.29448v1/comparison_gs.png)\(a\)GS equation
![Refer to caption](https://arxiv.org/html/2608.29448v1/comparison_gl.png)\(b\)GL equation\.

Figure 6:Gray\-Scott and Ginzburg\-Landau equations\. Training loss trajectories of Adam, Muon, SOAP, andSS\-eSOAPacross different network sizes, with MLP as training backbone\.Ginzburg\-Landau equation\.The complex Ginzburg\-Landau equation in 2D takes the form

∂A∂t=ϵ​Δ​A\+μ​A−γ​A​\|A\|2,t∈\(0,1\),\(x,y\)∈\(−1,1\)2,\\frac\{\\partial A\}\{\\partial t\}=\\epsilon\\Delta A\+\\mu A\-\\gamma A\|A\|^\{2\},\\quad t\\in\(0,1\),\(x,y\)\\in\(\-1,1\)^\{2\},with periodic boundary conditions, an initial condition

A0​\(x,y\)=\(10​y\+10​i​x\)​exp⁡\(−0\.01​\(2500​x2\+2500​y2\)\),A\_\{0\}\(x,y\)=\(10y\+10ix\)\\exp\\left\(\-0\.01\\left\(2500x^\{2\}\+2500y^\{2\}\\right\)\\right\),whereAAis the complex amplitude representing the envelope of oscillations,ϵ\\epsilonrepresents the diffusion coefficient,μ\\muis the linear growth rate, andγ\\gammacontrols the nonlinear saturation\. For this example, we setϵ=0\.004,μ=10\\epsilon=0\.004,\\mu=10andγ=10\+15​i\\gamma=10\+15i\. By denotingA=u\+i​vA=u\+iv, we can decompose the equation into real and imaginary components, resulting in the following system of PDEs,

∂u∂t=ϵ​Δ​u\+μ⁡\(u−\(u−1\.5​v\)​\(u2\+v2\)\),\\displaystyle\\frac\{\\partial u\}\{\\partial t\}=\\epsilon\\Delta u\+\\mu\\left\(u\-\(u\-1\.5v\)\\left\(u^\{2\}\+v^\{2\}\\right\)\\right\),∂v∂t=ϵ​Δ​v\+μ⁡\(v−\(v\+1\.5​u\)​\(u2\+v2\)\)\.\\displaystyle\\frac\{\\partial v\}\{\\partial t\}=\\epsilon\\Delta v\+\\mu\\left\(v\-\(v\+1\.5u\)\\left\(u^\{2\}\+v^\{2\}\\right\)\\right\)\.
Lid\-driven Cavity\.We study the incompressible Navier\-Stokes equations in non\-dimensional form for a 2D domain:

𝐮⋅∇𝐮\+∇p−1R​e​Δ​𝐮=0,\(x,y\)∈\(0,1\)2∇⋅𝐮=0,\(x,y\)∈\(0,1\)2\\begin\{array\}\[\]\{rlr\}\\mathbf\{u\}\\cdot\\nabla\\mathbf\{u\}\+\\nabla p\-\\frac\{1\}\{Re\}\\Delta\\mathbf\{u\}&=0,&\(x,y\)\\in\(0,1\)^\{2\}\\\\ \\nabla\\cdot\\mathbf\{u\}&=0,&\(x,y\)\\in\(0,1\)^\{2\}\\end\{array\}where𝐮=\(u,v\)\\mathbf\{u\}=\(u,v\)represents the steady\-state velocity field,ppis the pressure field, andR​eReis the Reynolds number which characterizes the ratio of inertial to viscous forces\. This system models the equilibrium state of the flow, which is driven by the top boundary moving at a constant velocity while the other walls are stationary, leading to the formation of characteristic vortical structures whose complexity increases with the Reynolds number\. To ensure continuity at the corner boundaries, we implement a smoothed top\-lid boundary condition:

u⁡\(x,y\)=1−cosh⁡\(C0​\(x−0\.5\)\)cosh⁡\(0\.5​C0\),v⁡\(x,y\)=0u\(x,y\)=1\-\\frac\{\\cosh\\left\(C\_\{0\}\(x\-0\.5\)\\right\)\}\{\\cosh\\left\(0\.5C\_\{0\}\\right\)\},\\quad v\(x,y\)=0wherex∈\[0,1\],y=1,C0=50x\\in\[0,1\],y=1,C\_\{0\}=50\. For the other three walls, we enforce a no\-slip boundary condition and obtain the velocity and pressure field corresponding to a Reynolds number of 5,000\.

\(a\)LDC equation, width=128=128\.![Refer to caption](https://arxiv.org/html/2608.29448v1/ldc_ep488450.png)\(b\)LDC equation, width=256=256\.

Figure 7:Lid\-driven Cavity equation\. Training loss trajectories of Adam, Muon, SOAP, andSS\-eSOAPacross different network sizes\.\(a\)2D Boussinesq\(b\)Burgers![Refer to caption](https://arxiv.org/html/2608.29448v1/x2.png)\(c\)Ginzburg\-Landau
![Refer to caption](https://arxiv.org/html/2608.29448v1/x3.png)\(d\)Gray\-Scott
![Refer to caption](https://arxiv.org/html/2608.29448v1/x4.png)\(e\)Korteweg\-de Vries
\(f\)Lid\-Driven Cavity\(g\)Wave![Refer to caption](https://arxiv.org/html/2608.29448v1/ac_comparison_best.png)\(h\)Allen\-Cahn

Figure 8:Training Dynamics across PDE Benchmarks\.We compare the residual loss trajectories \(log scale\) ofSS\-eSOAP\(red\) against baselines: Adam, Muon, standard SOAP, and Purifying SOAP\. Across most stiff PDE tasks,SS\-eSOAPachieves 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\.\(a\)Matrix Quadratic Regression: Losses, residual, gradient condition numbers, and nuclear norms\.\(b\)Matrix Logistic Regression: Losses, accuracy, gradient condition numbers, and nuclear norms\.\(c\)Low\-Rank Matrix Completion: Losses, relative reconstruction errors, and nuclear norms\.Figure 9:Matrix Optimization Benchmarks\.We compareSS\-eSOAPagainst Adam, Muon, and Standard SOAP on \(a\) quadratic regression, \(b\) quadratic regression, and \(c\) low\-rank completion\.SS\-eSOAPconsistently achieves faster convergence and lower final residuals, particularly in the quadratic regime where the self\-scaling mechanism effectively captures the spectrum\.Table 4:Matrix logistic regression loss trajectory\. Mean±\\pmstandard error over random seeds\.
### C\.3Full Details of Matrix Optimization Tasks

Quadratic Regression\.We consider a matrix quadratic regression objectivef⁡\(X\):=12​‖A​X​B−C‖F2f\(X\):=\\frac\{1\}\{2\}\\\|AXB\-C\\\|\_\{\\mathrm\{F\}\}^\{2\}, whereX∈ℝm×n,A∈ℝp×m,B∈ℝn×qX\\in\\mathbb\{R\}^\{m\\times n\},A\\in\\mathbb\{R\}^\{p\\times m\},B\\in\\mathbb\{R\}^\{n\\times q\}andC∈ℝp×qC\\in\\mathbb\{R\}^\{p\\times q\}\. Then its gradient is∇f​\(X\)=A⊤​\(A​X​B−C\)​B⊤\\nabla f\(X\)=A^\{\\top\}\(AXB\-C\)B^\{\\top\}, its Hessian is∇2f​\(X\)=\(B​B⊤\)⊗\(A⊤​A\)∈ℝm​n×m​n\\nabla^\{2\}f\(X\)=\(BB^\{\\top\}\)\\otimes\(A^\{\\top\}A\)\\in\\mathbb\{R\}^\{mn\\times mn\}\. The Hessian structure implies that ideal preconditioning isolates the residual condition numberκ⁡\(E\)\\kappa\(E\)from the constant spectral skew ofAAandBB\. Unlike semi\-orthogonal projections that discard curvature information from the residualEEto enforce unit conditioning \(κ=1\\kappa=1\),SS\-eSOAPpreserves this structural information\. We set\(m,n,p,q\)=\(500,100,1000,250\)\(m,n,p,q\)=\(500,100,1000,250\)so thatffis strongly convex\. As shown in Figure[9\(c\)](https://arxiv.org/html/2608.29448#A3.F9.sf3), this allowsSS\-eSOAPto mimic a Newton\-like trajectory, resolving the low\-rank structure significantly faster than Adam and standard SOAP\. While Adam’s gradient condition numberκ​\(∇f​\(Xk\)\)\\kappa\(\\nabla f\(X\_\{k\}\)\)grows unstably,SS\-eSOAPmaintains a bounded condition number throughout training\. The residual condition numberκ⁡\(Ek\)\\kappa\(E\_\{k\}\)correlates perfectly with the loss plateaus, validating our theoretical insight that conditioning on the residual spectrum is crucial for convergence in matrix sensing tasks\. Muon’s plateau suggests it fails to fully resolve this residual structure compared to the curvature\-adaptiveSS\-eSOAP\.

Logistic regression\.We study a matrix logistic regression problem with the objectivef\(X\)=∑i=1Nlog\(1\+exp\(−ci⊙\(aiXB\)\)\)\{f\}\(X\)=\\sum\_\{i=1\}^\{N\}\\log\(1\+\\exp\(\-c\_\{i\}\\odot\(a\_\{i\}XB\)\)\), whereX∈ℝm×nX\\in\\mathbb\{R\}^\{m\\times n\},A∈ℝN×mA\\in\\mathbb\{R\}^\{N\\times m\},B∈ℝn×qB\\in\\mathbb\{R\}^\{n\\times q\}andC∈ℝN×qC\\in\\mathbb\{R\}^\{N\\times q\}, andai∈ℝ1×ma\_\{i\}\\in\\mathbb\{R\}^\{1\\times m\}andci∈ℝ1×qc\_\{i\}\\in\\mathbb\{R\}^\{1\\times q\}are the row vectors ofAAandCC, respectively\. We set\(m,n,N,q\)=\(1000,100,10000,400\)\(m,n,N,q\)=\(1000,100,10000,400\)\. From Figure[9\(c\)](https://arxiv.org/html/2608.29448#A3.F9.sf3), we observe the following key behaviors in the stochastic logistic regime:SS\-eSOAPmaintains a consistent descent rate even as predictions enter the saturation regions of the sigmoid function \(high confidence\)\. Unlike Adam, which decelerates as gradient magnitudes vanish,SS\-eSOAP’s self\-scaling curvature correction effectively rescales the step size to counteract the flattening landscape\. And while Muon exhibits strong initial convergence due to its whitening properties, it shows greater variance in the later stages compared toSS\-eSOAP\. This suggests that while spectral normalization is beneficial, the explicit variance control mechanism \(purifying trigger\) inSS\-eSOAPprovides superior stability against the noise inherent in stochastic mini\-batch gradients\.

Low\-rank matrix completion\.We also study a simple nonconvex low\-rank matrix completion problem with a mask𝒜=\(ai,j\)1⩽i⩽m,1⩽j⩽n∈ℝm×n\\mathcal\{A\}=\\left\(a\_\{i,j\}\\right\)\_\{1\\leqslant i\\leqslant m,1\\leqslant j\\leqslant n\}\\in\\mathbb\{R\}^\{m\\times n\}to mimic missing entries\. This model can be viewed as a very simplified neural network\. The objective function isf\(X,Y\)=∥𝒜⊙\(XY⊤−M⋆\)‖/2F‖𝒜∥F2\{f\}\(X,Y\)=\\\|\\mathcal\{A\}\\odot\(XY^\{\\top\}\-M\_\{\\star\}\)\\left\\\|\{\}\_\{F\}^\{2\}/\\right\\\|\\mathcal\{A\}\\\|\_\{F\}^\{2\}, whereX∈ℝm×r,Y∈ℝn×rX\\in\\mathbb\{R\}^\{m\\times r\},Y\\in\\mathbb\{R\}^\{n\\times r\}\. We choose\(m,n,r\)=\(500,250,5\)\(m,n,r\)=\(500,250,5\)\. From Figure[9\(c\)](https://arxiv.org/html/2608.29448#A3.F9.sf3), we observe thatSS\-eSOAPdemonstrates superior convergence efficiency: Compared to Adam and standard SOAP,SS\-eSOAPachieves the fastest descent rate and reaches the lowest final residual\. Muon suffers from stagnation despite rapid initial progress: While Muon’s whitening effect provides a strong initial acceleration outperforming Adam, its convergence plateaus earlier than curvature\-adaptive methods\. This suggests that orthogonalizing gradients without explicitly accounting for the magnitude of singular values limits its ability to resolve the fine\-grained structure of the low\-rank manifold in the terminal phase\.

### C\.4Optimizer Hyperparameters and Tuning Setup

To ensure a rigorous and fair evaluation, all baseline optimizers \(Adam, Muon, SOAP, and Purifying Shampoo\) were subjected to a systematic grid search over their most sensitive hyperparameters\. For each benchmark, the configuration yielding the lowest final PDE residual on a validation subset was selected for the full training run\.

For all structured preconditioned optimizers \(SOAP, Purifying Shampoo, andSS\-eSOAP\), we utilized a standard EMA decay schedule for the second\-moment estimators, withβ2=0\.95\\beta\_\{2\}=0\.95\. Momentum was uniformly set toβ1=0\.9\\beta\_\{1\}=0\.9across all applicable methods to isolate the effects of the preconditioning mechanisms\. The learning rate \(η\\eta\) was the primary axis of tuning, swept logarithmically across a broad range for all algorithms\.

#### SS\-eSOAPConfiguration\.

A key advantage ofSS\-eSOAPis its robustness to hyperparameter variance across different scales of stiffness\. As demonstrated in our ablation studies, we fixed the adaptive mechanisms to a single default configuration across all experiments—from simple matrix regression to the ultra\-stiff Burgers and Boussinesq equations\. Specifically, the eigenbasis approximation check interval was set toIc​h​e​c​k=1I\_\{check\}=1, ensuring the off\-diagonal mass ratioρ\\rhois monitored continuously\. The purifying trigger threshold was fixed atτt​r​i​g​g​e​r=0\.2\\tau\_\{trigger\}=0\.2\. When a basis update is triggered, the variance downscaling factorγ\\gammais determined dynamically based on the severity of the misalignment:γ=0\.25\\gamma=0\.25ifρ\>0\.8\\rho\>0\.8,γ=0\.5\\gamma=0\.5ifρ\>0\.5\\rho\>0\.5, andγ=0\.75\\gamma=0\.75otherwise\.

Table[5](https://arxiv.org/html/2608.29448#A3.T5)details the grid search spaces employed for the baseline tuning, while Table[6](https://arxiv.org/html/2608.29448#A3.T6)summarizes the fixed configuration used forSS\-eSOAP\.

Table 5:Hyperparameter grid search spaces for baseline optimizers\.Table 6:Fixed default hyperparameters forSS\-eSOAPused across all experiments\.

## Appendix DAdditional Experiments

### D\.1Transient behavior at basis updates

![Refer to caption](https://arxiv.org/html/2608.29448v1/plots_new/transient_gamma_ablation.png)Figure 10:Loss around triggered updates forγ=1\\gamma=1, hard reset, and adaptiveγ<1\\gamma<1\.The steady\-state bound in Theorem[4\.1](https://arxiv.org/html/2608.29448#S4.Thmtheorem1)has its smallest right\-hand side atγ=1\\gamma=1\. Figure[10](https://arxiv.org/html/2608.29448#A4.F10)instead studies transient optimizer\-state behavior near triggered basis updates\. Exact reprojection produces large immediate loss spikes\. A hard reset avoids the largest spikes but reaches a higher final loss\. Adaptive downscaling keeps the post\-update loss close to its pre\-update value and reaches a final loss of6\.54×10−56\.54\\times 10^\{\-5\}in this run\. This experiment does not establish an optimal value ofγ\\gamma\. The exact\-reprojection arm also differs from the scalarγ=1\\gamma=1transition analyzed in the theorem\.

### D\.2PINN\-specific baselines

![Refer to caption](https://arxiv.org/html/2608.29448v1/plots_new/domain_specific_pinn_baselines.png)Figure 11:Comparison with PINN\-specific loss\-balancing methods\. MultiAdam, NTK reweighting, andSS\-eSOAPtrajectories under a matched benchmark, precision, architecture, and training budget\.Figure[11](https://arxiv.org/html/2608.29448#A4.F11)comparesSS\-eSOAPwith MultiAdam and NTK\-based loss reweighting under the same architecture and training budget\. MultiAdam descends steadily before reaching a residual near3\.0×10−43\.0\\times 10^\{\-4\}\. NTK reweighting reaches a lower trend near7\.0×10−57\.0\\times 10^\{\-5\}but exhibits larger fluctuations\.SS\-eSOAPreaches4\.0×10−54\.0\\times 10^\{\-5\}in this experiment\. MultiAdam and NTK reweighting adjust the relative loss components, whileSS\-eSOAPmodifies the parameter\-update geometry\. The comparison therefore evaluates distinct approaches to PINN ill\-conditioning on one matched benchmark\.

### D\.3Dense and limited\-memory quasi\-Newton baselines

![Refer to caption](https://arxiv.org/html/2608.29448v1/plots_new/scaling_1d_burgers.png)Figure 12:Scaling on 1D Burgers\. Network width versus seconds per 1,000 steps for Adam, L\-BFGS\(m=20\)\(m=20\), SS\-BFGS, andSS\-eSOAP\.Table 7:Explicit self\-scaled baselines at the tractable network widthW=256W=256\. Results report mean±\\pmSD over three seeds\. SS\-BFGS and SS\-Broyden run out of memory atW=1024W=1024, whileSS\-eSOAPruns with approximately10×10\\timesthe Adam wall\-clock time\.Figure[12](https://arxiv.org/html/2608.29448#A4.F12)compares runtime and memory scaling on 1D Burgers\. Dense SS\-BFGS exhausts available memory before widthW=600W=600\. L\-BFGS with history sizem=20m=20avoids dense curvature storage, but its per\-step time grows with network width\.SS\-eSOAPruns atW=1024W=1024, where its cubic basis operations produce roughly10×10\\timesthe Adam wall\-clock cost\. At the tractable widthW=256W=256, Table[7](https://arxiv.org/html/2608.29448#A4.T7)shows a different trade\-off\. SS\-BFGS obtains the lowest residual, whileSS\-eSOAPobtains the lowest relativeL2L^\{2\}error\. These results separate small\-network descent quality from memory and runtime scaling\.

### D\.4Physical errors and seed variation

Table 8:RelativeL2L^\{2\}andH1H^\{1\}errors against high\-fidelity PDE reference solutions\. Values are mean±\\pmstandard deviation over three independent seeds\. Lower is better, with the best result in each row shown in bold\. NTK denotes NTK\-based loss reweighting\.Residual loss alone does not establish solution accuracy\. Table[8](https://arxiv.org/html/2608.29448#A4.T8)reports relativeL2L^\{2\}andH1H^\{1\}errors against high\-fidelity reference solutions over three independent seeds\.SS\-eSOAPhas the lowest meanL2L^\{2\}andH1H^\{1\}error in all eight displayed comparisons\. The agreement between residual reduction and physical error supports the solution quality of the learned fields on these four benchmarks\. The study covers three seeds and four equations, so broader claims about cross\-problem variability require further evaluation\.

### D\.5Diagnostic for the basis trigger

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\.Table[9](https://arxiv.org/html/2608.29448#A4.T9)tracks off\-diagonal mass and consecutive gradient cosine similarity around one basis update\. From iteration10,00010\{,\}000to iteration15,00015\{,\}000, the off\-diagonal mass rises from0\.020\.02to0\.650\.65, while the cosine similarity falls from0\.880\.88to0\.050\.05\. After the triggered update, the two statistics return to0\.030\.03and0\.850\.85, respectively\. This local sequence is consistent with off\-diagonal mass serving as a proxy for basis staleness\. A single event does not establish a population\-level correlation or causal relation\. This diagnostic usesτtrigger=0\.5\\tau\_\{\\mathrm\{trigger\}\}=0\.5, while the default experiments useτtrigger=0\.2\\tau\_\{\\mathrm\{trigger\}\}=0\.2\.

### D\.6Interaction with the inner optimizer

![Refer to caption](https://arxiv.org/html/2608.29448v1/plots_new/gs_convergence.png)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\.Figure[13](https://arxiv.org/html/2608.29448#A4.F13)isolates the interaction between the directional scaling factor and the inner optimizer on Gray–Scott\. Removingτk\\tau\_\{k\}from the Adam\-based variant lowers the final residual to5\.0×10−65\.0\\times 10^\{\-6\}\. Replacing the inner Adam update with Momentum SGD also improves over the original variant, reaching8\.0×10−58\.0\\times 10^\{\-5\}\. The effect ofτk\\tau\_\{k\}therefore depends on the inner normalization rule and the benchmark\. Under the reportedτk−1/2\\tau\_\{k\}^\{\-1/2\}convention,τk<1\\tau\_\{k\}<1enlarges the update\. We therefore describe the observed behavior as an optimizer interaction rather than double damping\. This ablation does not imply thatτk\\tau\_\{k\}should be removed on other benchmarks\.

## Appendix EComputational Complexity Analysis

We analyze the per\-iteration complexity ofSS\-eSOAPcompared to standard SOAP and Shampoo\. Letm,nm,nbe the dimensions of the layer weights𝐖∈ℝm×n\\mathbf\{W\}\\in\\mathbb\{R\}^\{m\\times n\}\. Assumem≥nm\\geq n\.

Standard iteration costs\.Gradient computation takesO⁡\(m​n\)O\(mn\), updating𝐋k\\mathbf\{L\}\_\{k\}takesO⁡\(m2​n\)O\(m^\{2\}n\),𝐑k\\mathbf\{R\}\_\{k\}takesO⁡\(m​n2\)O\(mn^\{2\}\)\. Transforming𝐆k\\mathbf\{G\}\_\{k\}to eigenspace𝐐L⊤​𝐆k​𝐐R\\mathbf\{Q\}\_\{L\}^\{\\top\}\\mathbf\{G\}\_\{k\}\\mathbf\{Q\}\_\{R\}takesO⁡\(m2​n\+m​n2\)O\(m^\{2\}n\+mn^\{2\}\)\. Element\-wise update takesO⁡\(m​n\)O\(mn\)for𝐌,𝐕\\mathbf\{M\},\\mathbf\{V\}updates\. Inverse projection takesO⁡\(m2​n\+m​n2\)O\(m^\{2\}n\+mn^\{2\}\)\. The total standard cost isO⁡\(m2​n\)O\(m^\{2\}n\)\.

Overhead ofSS\-eSOAPfeatures\.Computingτk\\tau\_\{k\}requires the following trace operations: Computing the numeratorTr⁡\(𝐘⊤​𝐒\)\\operatorname\{Tr\}\(\\mathbf\{Y\}^\{\\top\}\\mathbf\{S\}\)takesO⁡\(m​n\)O\(mn\)\. As for computing the denominatorTr⁡\(𝐒⊤​𝐋−1​𝐒𝐑−1\)\\operatorname\{Tr\}\(\\mathbf\{S\}^\{\\top\}\\mathbf\{L\}^\{\-1\}\\mathbf\{S\}\\mathbf\{R\}^\{\-1\}\), since we operate in the eigenspace where𝐋−1,𝐑−1\\mathbf\{L\}^\{\-1\},\\mathbf\{R\}^\{\-1\}are diagonal \(DL−1,DR−1D\_\{L\}^\{\-1\},D\_\{R\}^\{\-1\}\), this becomes an element\-wise product followed by a sum\. The total cost isO⁡\(m​n\)O\(mn\)\. This is negligible compared to theO⁡\(m2​n\)O\(m^\{2\}n\)projection cost\.

Checking the off\-diagonal ratio \(purifying trigger\) requires computing𝐐⊤​𝐋𝐐\\mathbf\{Q\}^\{\\top\}\\mathbf\{L\}\\mathbf\{Q\}, which leads to a total check cost ofO⁡\(m3\)O\(m^\{3\}\)\. This check happens only at intervalsIc​h​e​c​kI\_\{check\}\(default 1\)\. WhileO⁡\(m3\)O\(m^\{3\}\)seems high, for PINNs the layer dimensionsm,nm,nare typically small \(e\.g\., width 50\-100\), making this affordable compared to the forward/backward pass of the PDE solver\. Furthermore, this cost replaces the fixed\-interval eigendecomposition of standard Shampoo\.

For the eigendecomposition step, when the trigger fires \(ρ\>τt​r​i​g​g​e​r\\rho\>\\tau\_\{trigger\}\), we perform exact eigendecomposition, which costsO⁡\(m3\)O\(m^\{3\}\)\. In standard SOAP, this occurs at fixed frequencyFF\. InSS\-eSOAP, this is adaptive\. If the landscape is stable \(ρ\\rhois small\), we skip this step entirely, potentially reducing the amortized complexity compared to fixed\-schedule Shampoo\.

Standard reprojection requires𝐕n​e​w=𝐐n​e​w⊤​𝐐o​l​d​𝐕​…\\mathbf\{V\}\_\{new\}=\\mathbf\{Q\}\_\{new\}^\{\\top\}\\mathbf\{Q\}\_\{old\}\\mathbf\{V\}\\dotswhich requiresO⁡\(m3\)O\(m^\{3\}\)matrix multiplications to rotate the variance matrix\. Downscaling requires𝐕n​e​w=γ​𝐕\\mathbf\{V\}\_\{new\}=\\gamma\\mathbf\{V\}\. This is an element\-wise scalar mult, takingO⁡\(m​n\)O\(mn\)\. Therefore, the purification step inSS\-eSOAPis strictly cheaper \(O⁡\(m​n\)O\(mn\)vsO⁡\(m3\)O\(m^\{3\}\)\) than the reprojection step used in standard implementations when an update occurs\.

Summary of complexity\.SS\-eSOAPintroduces anO⁡\(m3\)O\(m^\{3\}\)check at every step\. For large models \(e\.g\., Transformers withdm​o​d​e​l=4096d\_\{model\}=4096\), this is prohibitive\. However, for PINNs where layer widths are typicallym≤256m\\leq 256, the cubic term is negligible\. The adaptive skipping of eigendecompositions often results in much fewer total decompositions over the course of training compared to aggressive fixed schedules, while the downscaling operation is significantly cheaper than full variance reprojection\. AlthoughSS\-eSOAPintroduces additional overhead per iteration due to eigenbasis management, this cost is negligible compared to the reduction in the number of iterations required to resolve stiff PDE dynamics\.

## Appendix FComputational Efficiency Analysis

Structural efficiency\.Finally, by inheriting the Kronecker structure from SOAP,SS\-eSOAPreduces the computational cost of curvature estimation fromO⁡\(N2\)=O⁡\(\(din​dout\)2\)O\(N^\{2\}\)=O\(\(d\_\{\\text\{in\}\}d\_\{\\text\{out\}\}\)^\{2\}\)toO⁡\(din2\+dout2\)O\(d\_\{\\text\{in\}\}^\{2\}\+d\_\{\\text\{out\}\}^\{2\}\)per layer\. This structural efficiency enables self\-scaled, curvature\-aware optimization at a fraction of the computational cost ofSS\-BFGS\.

Convergence speed\.Assessing wall\-clock efficiency requires distinguishing between theearly explorationphase and thehigh\-precision terminalphase\. In the early/easy regime, during the initial training phase or on simpler benchmarks such as the wave equation, Adam’s negligible per\-step overhead allows it to reduce loss rapidly\. Here,SS\-eSOAPmatches Adam’s wall\-clock time to reach the same intermediate loss\. In the stiff/high\-precision regime, the advantage ofSS\-eSOAPbecomes decisive when targeting high precision\. In this regime, Adam typically plateaus due to ill\-conditioning\. For example, on the stiff 2D Boussinesq benchmark, Adam stagnates after10\+10\+hours\.SS\-eSOAPmaintains linear convergence, crossing the10−510^\{\-5\}threshold in 4\.1 hours\. Thus, for high\-fidelity scientific discovery,SS\-eSOAPoffers a clear computational speedup by bypassing the spectral stagnation that traps first\-order methods\.

Memory overhead\.While second\-order methods typically incur high memory costs, the Kronecker factorization keepsSS\-eSOAP’s footprint competitive\. Compared to Adam \(which stores 2 states per parameter,2​N2N\),SS\-eSOAPstores the Kronecker factorsdin2\+dout2d\_\{\\text\{in\}\}^\{2\}\+d\_\{\\text\{out\}\}^\{2\}plus standard momentum buffers\. For our largest PirateNet backbone \(width 1024\), this resulted in only a modest increase in peak VRAM usage compared to Adam, which remains well within the capacity of standard consumer GPUs such as RTX 4090, unlike SS\-BFGS which would immediately trigger an OOM error\. Moreover,SS\-eSOAPmaintains the same hardware\-efficient profile as standard SOAP\. While the calculation of the self\-scaling factorτk\\tau\_\{k\}introduces a theoretical computational overhead, we note that the requisite displacement termsS~k\\tilde\{S\}\_\{k\}are cached from the previous iteration’s eigenspace update\. Consequently, computingτk\\tau\_\{k\}requires only inexpensive element\-wise operations inO⁡\(N\)O\(N\), resulting in a negligible increase in FLOPs per step compared to standard SOAP\. Thus,SS\-eSOAPdelivers the convergence benefits of self\-scaling without the hardware penalties associated with curvature\-adaptive optimizers\.

Scalability and compute trade\-off\.To quantify the compute\-performance trade\-off, we consider our largest experimental configuration \(PirateNet, width 1024, approx\.10710^\{7\}parameters\)\. Running full SS\-BFGS in this regime is impractical: storing an explicitN×NN\\times Ninverse Hessian approximation would consume over 400 TB of VRAM, far exceeding the capacity of any existing GPU cluster\. A back\-of\-the\-envelope calculation confirms that SS\-BFGS saturates the 80GB memory of an NVIDIA A100 at merelyN≈140,000N\\approx 140,000parameters, effectively restricting it to toy problems\. While limited\-memory variants like L\-BFGS mitigate thisO⁡\(N2\)O\(N^\{2\}\)storage cost, they remain computationally prohibitive due to the sequential overhead of strong Wolfe line searches, which typically require 2\-5×\\timesmore forward passes per iteration\. In contrast,SS\-eSOAPleverages the Kronecker structure to decouple memory complexity from total parameter count, requiring only megabytes of curvature storage per layer\. This efficiency enablesSS\-eSOAPto converge on 2D Boussinesq in just 4\.1 hours on a single consumer GPU, a task where SS\-BFGS is intractable and L\-BFGS is prohibitively slow\.

Table 10:Complexity Comparison per Iteration \(Assumingm=nm=n\)
## Appendix GAblation on eigenbasis update trigger threshold

Please see Table[11](https://arxiv.org/html/2608.29448#A7.T11)for the ablation results\.

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 approximately3\.5×3\.5\\times, effectively finding the optimal efficiency\-accuracy trade\-off\.
## Appendix HLimitations and trade\-offs of self\-scaling

While SS\-ESOAP demonstrates state\-of\-the\-art stability on ultra\-stiff problems, this robustness introduces specific trade\-offs in different optimization landscapes\. As observed in Table[1](https://arxiv.org/html/2608.29448#S5.T1), standard SOAP achieves marginally lower final residuals than SS\-ESOAP on the Gray\-Scott and Ginzburg\-Landau benchmarks\. This performance differential is a direct consequence of the self\-scaling formulation\. The factorτk\\tau\_\{k\}is designed as a conservative damping mechanism, bounded bymin⁡\{1,…\}\\min\\\{1,\\dots\\\}, which strictly enforces Rayleigh quotient matching to prevent divergence\.

Ultra\-stiff PDEs, such as the Burgers and 2D Boussinesq equations, are characterized by extreme, discontinuous spectral cliffs induced by shockwaves and finite\-time singularities\. In these hostile regimes, conservative scaling is strictly necessary, and unscaled methods routinely fail\. Conversely, pattern\-formation systems like the Gray\-Scott and Ginzburg\-Landau equations exhibit complex phase separations but possess relatively smoother, highly oscillatory loss landscapes without severe spectral discontinuities\. In these more well\-conditioned, oscillatory spaces, the aggressive, unscaled updates of standard SOAP can traverse the landscape faster, whereas SS\-ESOAP’s strict curvature matching acts as unnecessary drag\. Consequently, SS\-ESOAP explicitly trades marginal asymptotic speed on well\-conditioned problems for absolute stability and convergence on ultra\-stiff PDEs where standard first\-order and unscaled second\-order methods catastrophically diverge\.

## Appendix IWall\-Clock Time Across Benchmarks

Please see Table[12](https://arxiv.org/html/2608.29448#A9.T12)for the wall\-clock time details\.

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𝒪⁡\(N2\)\\mathcal\{O\}\(N^\{2\}\)memory cost of full SS\-BFGS\.
## Appendix JRelation to Natural\-Gradient Methods

Natural\-gradient and Gauss\-Newton methods precondition the gradient using function\-space geometry\. Energy natural gradients, Gauss\-Newton natural gradients, ANaGRAM, and related methods attain high accuracy on PINNs\[[21](https://arxiv.org/html/2608.29448#bib.bib36),[11](https://arxiv.org/html/2608.29448#bib.bib37),[26](https://arxiv.org/html/2608.29448#bib.bib32),[9](https://arxiv.org/html/2608.29448#bib.bib34)\]\. Dual formulations move the solve from parameter space to residual space and have been demonstrated on PINNs with up to 12\.8 million parameters\[[12](https://arxiv.org/html/2608.29448#bib.bib33)\]\. K\-FAC instead uses layerwise Kronecker blocks to approximate this geometry\[[6](https://arxiv.org/html/2608.29448#bib.bib35)\]\.

SS\-eSOAPuses a different approximation\. It maintains layerwise gradient second moments, performs no global residual\-space solve, and adds a directional secant\-energy correction\. Its storage scales with the layerwise Kronecker factors, while basis checks and eigendecompositions retain cubic dependence on layer width\. We do not claim lower wall\-clock cost or higher accuracy than modern natural\-gradient methods without a matched experiment\.

Table 13:Empirical comparison against the scalable K\-FAC, which shares our structural efficiency\.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\.
## Appendix KPotential Societal Impacts

The development of high\-precision, scalable optimizers likeSS\-eSOAPhas broad implications across several scientific and industrial domains\. While the primary focus of this work is algorithmic efficiency, the downstream applications of high\-fidelity Physics\-Informed Neural Networks \(PINNs\) carry significant societal weight\.

### K\.1Advancement of Scientific Discovery

The ability to resolve partial differential equations \(PDEs\) to near machine precision enables the use of deep learning in fields previously reserved for traditional high\-performance computing \(HPC\) methods\. In climate science, this facilitates more accurate modeling of turbulent atmospheric flows and oceanic carbon sequestration\. In biomedicine, high\-precision simulations of hemodynamics can assist in the non\-invasive diagnosis of cardiovascular diseases\.

Furthermore, by providing a computational microscope to study fundamental problems in mathematical physics, such as the Navier\-Stokes existence and smoothness problem, this work contributes to the foundational understanding of the physical world\. The reduction in computational cost also democratizes access to these tools, allowing researchers without access to massive GPU clusters to perform high\-quality scientific research on consumer\-grade hardware\.

### K\.2Potential Dual\-Use Concerns

High\-precision simulation is fundamentally a dual\-use technology\. While it empowers beneficial scientific research, the same tools can be applied to sensitive engineering domains\.

- •Aerodynamics and Defense:Precise solvers for fluid dynamics are critical in the design of advanced aerospace systems\. Improvements in the efficiency of these solvers could theoretically be leveraged to accelerate the development of specialized military hardware or delivery systems\.
- •Nuclear and Industrial Simulation:The capability to model complex, multi\-scale physical systems with high fidelity is relevant to nuclear reactor modeling and structural analysis\. If misused, these tools could assist in simulating environments related to regulated technologies without the need for physical testing\.
- •Accessibility and Security:By lowering the barrier to entry for high\-precision simulation, there is a risk that malicious actors could utilize these methods for industrial espionage or to bypass safety protocols by predicting structural failures in critical infrastructure\.

In conclusion, whileSS\-eSOAPprovides a robust framework for scientific progress, we advocate for the responsible deployment of these tools\. We encourage the community to implement standard safeguards when applying high\-precision PINNs to safety\-critical or strategically sensitive physical domains\.

Similar Articles

Soft Adaptive Policy Optimization

Papers with Code Trending

SAPO introduces a smooth, temperature-controlled gate to adaptively attenuate off-policy updates in reinforcement learning for large language models, enhancing training stability and performance compared to methods with hard clipping.

ESPO: Early-Stopping Proximal Policy Optimization

Hugging Face Daily Papers

ESPO introduces an early-stopping mechanism for reinforcement learning that detects and terminates failed reasoning trajectories in LLMs, improving mathematical reasoning performance while reducing compute by over 20%.