Gradient Alignment in Physics-informed Neural Networks: A Second-Order Optimization Perspective
Authors: Sifan Wang
Ananyae Kumar bhartari
Bowen Li
Paris Perdikaris
Venue: NeurIPS 2025 poster ()
Keywords: Scientific Machine Learning
Physics-informed Neural Networks
Partial Differential Equations
Second-order Optimization
Gradient Conflicts
Primary area: machine_learning_for_sciences
TL;DR: This work identifies gradient conflicts as a key challenge in training PINNs and shows that quasi second-order methods—especially SOAP—effectively resolve them, leading to 2-10x accuracy gains on 10 PDE benchmarks.

Abstract: Physics-informed neural networks (PINNs) have shown significant promise in computational science and engineering, yet they often face optimization challenges and limited accuracy. In this work, we identify directional gradient conflicts during PINN training as a critical bottleneck. We introduce a novel gradient alignment score to systematically diagnose this issue through both theoretical analysis and empirical experiments.
Building on these insights, we show that (quasi) second-order optimization methods inherently mitigate gradient conflicts, thereby consistently outperforming the widely used Adam optimizer. Among them, we highlight the effectiveness of SOAP \cite{vyas2024soap} by establishing its connection to Newton’s method.
Empirically, SOAP achieves state-of-the-art results on 10 challenging PDE benchmarks, including the first successful application of PINNs to turbulent flows at Reynolds numbers up to 10,000. It yields 2–10x accuracy improvements over existing methods while maintaining computational scalability, advancing the frontier of neural PDE solvers for real-world, multi-scale physical systems. All code and datasets used in this work are publicly available at:  
\url{https://github.com/PredictiveIntelligenceLab/jaxpi/tree/pirate}.
\end{abstract}

Gradient Alignment in Physics-informed Neural Networks: A Second-Order Optimization Perspective
Sifan Wang's Ananyae Kumar Bhartari Bowen Li Paris Perdikaris Institution for Foundation of Data Science, Yale University
Penn Institute for Computational Science, University of Pennsylvania 3 Department of Mathematics, City University Hong Kong DDepartment of Mechanical Engineering and Applied Mechanics, University of Pennsylvania s fan. wangwyalgaa.....mmommomo..m..m. edu, bowen. li@cityu..... edu. hk, (ananyasb, pgp) @seas upenn. edu
Abstract
Physics-informed neural networks (PINNs) have shown significant promise in computational science and engineering, yet they often face optimization challenges and limited accuracy. this work, we identify directional gradient conflicts during PINN training a critical bottleneck. We introduce novel gradient alignment score to systematically diagnose this issue through both theoretical analysis and empirical experiments. Building on these insights, we show that (quasi) second-order optimization methods inherently mitigate gradient conflicts, thereby consistently outperforming the widely used Adam optimizer. Among them, we highlight the effectiveness of SOAP II by establishing its connection to Newton's 's method. Empirically, SOAP achieves state-of-the-ard results on 10 challenging PDE benchmarks, including the first successful application of PINNs turbulent flows Reynolds numbers up to 10,000. It yields 2-10x accuracy improvements over existing methods while maintaining computational scalability, advancing the frontier of neural PDE solvers for real- world, multi-scal physical systems. All code and datasets used in this work are publicly available at: https:
Introduction
mhysics-informer neural networks (PINNs) have emerged as powerful paradigm in scientific ma- chine learning by incorporating physical principles through carefully designed loss functions. These loss functions as soft constraints, guiding neural networks to learn solutions that respect underlying physical laws while simultaneously fitting experimental data. The elegance and versatility of PINNS have led to their widespread adoption in solving both forward and inverse problems involving partial differential equations (PDEs). Their success spans numerous domains in computational science, from fluid mechanics 2
g bio-enginecring (9] 10 [T] and materials The impact of PINNs extends even further, with significant applications in
3 heat transfer 66, 7
science 12 13
14].
electromagnetics
15
16
geosciences [8
etc.
Despite their broad applications, PINNs currently face limitations convergence speed and accuracy that affect their reliability forward PDE solvers. This has motivated extensive research efforts to enhance their performance through various methodological innovations. Significant advances have emerged in neural architecture design, including novel network backbones 21 22
23
24
25 26 27].
improved activation functions
28 29 and effective coordinate embeddings
30 31
32
33 Other improvements have focused on optimizing the training process through enhanced
collocation point sampling strategies (34 35 36), more efficient optimizers [37] 38
39 40].
and
'These authors contributed equally to this work.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

advanced training strategies such as sequential training 44 42 43 and transfer learning 44
45
46
Researchers have also explored alternative formulations of the learning objective. incorporating numerical differentiation 47 variational principles inspired by Finite Element Methods 48 49
and specialized regularization terms 50 51.
A particularly active area of research has centered on addressing gradient pathologies during training [21 [2]. One prominent issue involves imbalanced backpropagated gradients across different loss terms, leading to significant discrepancies in convergence rates and reduced solution accuracy, especially in complex physical systems. This has led to the development of various adaptive weighting strategies [21] 52 53 54
55 [6, 57]. However, the equally critical issue of directional gradient conflicts where gradients from different losses point conflicting directions remains largely underexplored [58 [99. Our work aims to bridge this gap by systematically investigating, analyzing, and resolving these directional conflicts PINNs training. The key contributions of this work are summarized as follows:
We introduce a novel gradient alignment metric that extends cosine similarity to quantify directional conflicts between multiple loss terms.
We demonstrate that gradient conflicts hinder PINN training, with higher conflict scores linked to slower convergence in various PDE systems.
We show that second-order optimizers enhance gradient alignment by implicitly preconditioning the loss landscape.
We reveal that SOAP [ can be viewed as efficient approximation of the Newton preconditioner, revolving gradient conflicts.
We provide comprehensive experimental results across 10 PDE benchmarks, including the first successful PINN application to turbulent flows with Reynolds numbers up to 10,000, achieving 2-10x accuracy improvements.
Taken together, this work advances our understanding of optimization dynamics in PINNs while demonstrating how quasi second-order methods can enable more reliable neural PDE solvers for solving complex physical systems. These insights pave the way for developing next-generation optimizers for physics-informed machine learning, and beyond.
2 Overview of PINNs
Multi-task learning in deep neural networks requires simultaneously minimizing multiple competing objectives challenge that manifests acutely in physics informed neural networks (PINNs). Building upon the work of [60]. PINNs approximate solutions to partial differential equations by minimizing a composite loss function that enforces both physical constraints and data- fata-fitting objectives. Consider general PDE system:
physics-informed
+ Muu 0, t [0,2], X
(2.1)
with inital and boundary conditions
u(0,x) = g(x), x Ω
(2.2)
Blu] = 0, t E [0,T], x an
(2.3)
where [|- represents a differential operator and B[|| denotes boundary conditions. The core idea of PINNs is approximate the solution u(t, x) using a neural network un x). Through automatic differentiation [61], we can compute the PDE residual:
dus
R[ua][(, x)
(4r, x,) (( x,).
(2.4)
Øt
This leads composite loss function that encapsulates multiple competing objectives:
1 No.
N
1
L())-
uo (0. xx)
(xx))
B [ou] (lier xbr))
R [wa] (1,, x
Noc 1=1
Noc
N
())
(n)
2.
(2.5)

---

These loss functions aim to fit the initial, boundary conditions and the PDE residuals, respectively. And {xte) NN. (the: bee J No and {tz x may be selected either as fixed mesh vertices or through random sampling during each training iteration.
Gradient Alignment in PINNS
A fundamental challenge in training PINNs is that differ- ent loss terms often conflict during optimization. As illus- trated in Figure ! PINNs encounter two modes of gradient conflict during training. The first, identified by [2), involves significant imbalances in gradient magnitudes. In such cases, dominant loss terms overwhelm others, often leading model failure. Various self-adaptive weight- ing strategies have been proposed to address this issue [53 14
Type!
Type II
56 57).
Figure I: Gradient conflicts and their im- pact on PINNs optimization. The irregular green trajectory illustrates how the optimiza- tion struggles when facing two types of gra- dient conflicts: Type where gradients have similar directions but vastly different magni- tudes, and Type II, where gradients have sim- ilar magnitudes but opposing directions. The red trajectory shows how appropriate precon- ditioning through second-order information could mitigate these conflicts by aligning gra- dients both within and between optimization steps, enabling efficient convergence.
This second mode arises when gradients from different loss terms point in conflicting directions, forcing optimiza- tion into inefficient, compromised trajectories. For ex- ample, in the Navier-Stokes equations, enforcing no-slip boundary conditions demands precise control of veloc- ity gradients walls, which conflict with preserv- ing mass and momentum conservation in the bulk flow. First-order optimizers like gradient descent Adam must follow the average gradient direction, resulting in inef- ficient zigzagging between competing objectives. The severity of these conflicts increases with problem com- plexity, becoming particularly acute for turbulent flows where maintaining physical constraints across multiple scales is crucial. To better understand and address these directional gradient conflicts, we introduce the aligmment score, defined as follows.
Definition 1. Suppose that vn, 02..... are vectors, then the alignment score is defined
2i=1
A(r1, 02 ... Ux) un)= =
1.
(3.1)
71
This score ranges from [-1, 1] and naturally extends the concept of cosine similarity to multiple vectors. As illustrated in Proposition [ for the special case of 2, our exactly recovers the where 1 indicates perfect alignment, suggests [ell |
standard cosine similarity cos(u), U22 ||ez 21-22
orthogonal directions, and -1 represents complete opposition.
Proposition 1, For n=2, the alignment score A(u1, U2) equals the cosine similarity between in and U2.
21 U/
A(v1, 02) cos(U1, U2)
(3.2)
|| 22
The proof is provided in Appendix F.I The alignment score enables us to quantify both the local conflicts between individual loss terms within each gradient descent step and the global conflicts across consecutive steps. Formally:
Let L = Eim Li be a composite loss function. At the k-th step of gradient descent, let ge denote the full gradient and o ...,.. denote the gradients of individual loss terms. We define:
Definition 2.
(a) The intra-step gradient alignment score:
A A(gi ,,.),,,,,))
(3.3)
(b) The inter-step gradient alignment score:
A Alg--1,0)
(3.4)
3

---

Navier Stokes
Inter-step Alignment
Intra-step Alignment
L2 Error
1.0
1.0
0.6
0.0
0.0
0.4
-1.0
-1.0
0.0
5000
10n00
5000
10000
5000
10000
Step
Skep
Adam
SOAP
Kron
Muon
Adagrad
Figure 2: Gradient aligmment scores and test errors during PINN training for solving the Navier- Stokes equations with different optimizers. Additional benchmarks are provided in Figure where we observe the consistent phenomenon that first-order optimizers exhibit poor gradient alignment and slow convergence of test errors.
The intra-step gradient alignment score becomes especially useful for PINNs applications where the total loss typically consists of multiple terms. For example, the case of solving the 2D Navier- Stokes equations, the loss function includes separate components corresponding to momentum equations in the and y directions, the continuity equation, and boundary conditions for the velocity fields u and Since these loss terms can have different scales and properties, they should be treated separately rather than grouped together. Our intra-step score effectively quantifies gradient conflicts across all these terms simultaneously.
The impact of aligmment can be formalized through the following result, which shows that the rate of loss decay under preconditioned gradient descent depends on the alignment scores.
Proposition 2. Let L Ei= Li IRd + be L-smooth (w.r.. the Euclidean norm). Consider preconditioned gradient descent
O++1 -0---11 1j11s
he Pist.
9t Σ
9e
where each Pe > O satisfies S (P)) < Amax (P.) S M. Let A. = L(0,) - [(0++1) and
assume & 1/M. Then:
(i) Single-step drop. For all if [hgll X for all i,
A. 2
n M
Lnd
n
[ intra + 1) Σ |1//122
(ii) Two-step cumulative drop. Let = || ||, b = |hell-- Then
7 M (1 Lxnh a + M (1 LnM) ab A Laj b².
A + A.
2
This result shows that higher intra- and inter-step aligmment scores directly accelerate loss reduction. The assumption |h,] = A for all is not restrictive in the PINNs setting, where balancing the scales of different losses is common practice. In fact, adopt the weighting scheme proposed by Wang al. [21] (Appendix G.2), which ensures that all weighted gradients have the same norm.
In the following, we will demonstrate that gradient direction conflicts widely exist in training PINNs, especially in the carly stages training. To this end, we conduct experiments on five representative PDEs spanning from linear wave propagation to reaction-diffusion systems like the Sinzburg-Landee equation and fluid dynamics governed by the Navier-tookes equations. The detailed results are presented in Figure 5 and corresponding experimental setup is provided in Appendix G
Figure presents gradient alignment scores and test errors during PINN training for solving the Navier-Stokes equations with different optimizers. Importantly, we compute these scores from the gradients after applying each optimizer's gradient transformations, which captures the actual update directions and their degree of alignment conflict. We observe that the scores from first order optimizer oscillate significantly near below zero at early stages. provides strong evidence for persistent directional conflicts between gradients throughout the training process. Intuitively, these

---

conflicting gradients force the network parameters to follow an inefficient zigzag trajectory in the loss landscape. significantly impeding convergence speed.
In contrast, the quasi second-order optimizers (e.g., Muon 62] SOAP 0. Kron
68 consistently
maintains the higher positive values for both inter-step and intra-step gradient alignment scores throughout training. This effective resolution of gradient direction conflicts directly corresponds to significantly faster convergence in error. is worth nothing that SOAP achieves the highest alignment scores among all optimizers, demonstrating its superior effectiveness. In the following sections, we provide theoretical explanation for why PINNs inherently suffer from directional gradient conflicts and establish a formal connection between SOAP and Newton's 's method, offering deeper insight into its empirical success.
4 Intra-step Gradient Alignment of PINNs at Initialization
In this section, we show that gradient conflicts are intrinsic to PINNs near initialization, regardless of the optimizer used. For simplicity, we analyze the intra-step gradient alignment for two-layer neural network solving the one-dimensional Laplace equation under small weight initialization. This setup allows for tractable analysis while capturing key phenomena, and the results naturally extend to more general PDEs, Following the setup in Section 2 and without loss of generality, we consider the ID Laplace equation:
Au= = =0= =0 on [-1,
(4.1)
u(+1) 9+1-
We approximate the solution u(x) by a two-layer network with width N:
u(x, Σ a: (w;x) o(wx)
(4.2)
where a = (ai..... a. = (wi...., wx) E IR and 0 = (a. w) E R2N Moreover, we limit
ourselves the activation function (I) tanh(x:). ). In this case, the loss 2.5
reduces to
N N
min ((0) 8=(a.w)
Iu" (xp, 0)1 + Σ Ju(s, 0) gel
(4.3)
C.(0)
(oo10)
Proposition 3. At initialization, we assume that the weights a, wA are initialized by i.i.d. Gaussian (0 with small E = o(1). Then the alignment score converges to binary random variable in the infinite width limit:
sgn (0-1 91) sgn (g-1 -9)
with prob. 3' (4.4)
lim ((CC)) O(=3) -( On
with prob. 2-
where D= GD, Adam,, Shampoo, or Soap denotes the corresponding optimizer update rule, and Co is constant depending the optimizer.
The proof is provided in Appendix As a direct implication, the intra-step alignment score exhibits no preferred direction. Consequently, all optimizers fail to induce consistent intra-step alignment between V. and VL, near initialization, as empirically validated in Figure In contrast, SOAP maintains a perfect inter-step step alignment score of throughout the entire training process. In the following section, we provide theoretical explanation for this distinctive behavior.
Gradient Alignment in Quasi Second-Order Optimization
In this section, we aim to reveal why quasi Second-order optimizers can naturally promote gradient alignment via the following lemma.
Proposition 4. Let L(0) be twice differentiable loss function with Hessian ((),, and let P(0) be a preconditioner with uniformly bounded inverse P-) Consider the preconditioned gradient descent update with exponent 0 S < 1 and learning rate 0:
6 8,41 02 - P 5 (0, V L(0.)
(5.1)
5

---

be interpreted as Let = (())) then the alignment score Algi, 9t-1) satisfies:
Algr.9t+11 = 17²
ННР* " + O
(5.2)
2
1
To mainlain alignment A (t+1))+)) 2 ffr > 0. the learning rate v must satisfy:
2e gel
(5.3)
[HP-"gel
The proof is presented in Appendix gradient descent (P I), 1/max
C This bound specializes to the following cases: For vanilla
2r (Hr and for Newton's method (P H, = 1), the
V
maximum learning rate can be relaxed to fJmax = 2e. These results imply that preconditioners approximating the Hessian effectively relax learning rate constraints. enabling the use of larger learning rates while maintaining optimization trajectory consistency. Comparing these cases reveals that Newton's method eliminates the dependency on the condition number of the Hessian, allowing for a constant maximum learning rate regardless of problem conditioning. This finding aligns with our theoretical understanding that accurate second-order information mitigates the effects of ill-conditioning. allowing for stable optimization even with aggressive learning rates.
Next, establish that under assumptions, the SOAP optimizer approximating to using Hessian as the preconditioner:
Wh+11 Wty 7) Soap(g:) R w. H 9c-
(5.4)
Formal assumptions and proof are provided in Appendix
Table 1: Comparison of optimization methods showing preconditioner types, storage and compu- tational complexity for n x weight matrix, and practical compatibility with mini-batches and scalability with large neural networks. The theoretical connections between Shampoo, Muon, and quasi-second-order methods are detailed in Appendix D and
Method
(Approx.) Precond.
Storage
Computation
Mini-batch
DNN
Natural Gradient 664 37
F-
O(n²)
O(z))
X
x
BFGS/L-BFGS 65
11-1
O(n²)
(
X
SOAP
H-1
((n²)
((z²)
Kron
63
F-1/2
O(n²)
(
Shampoo 66
H-1/2
O(n²)
(n
Muon
Ha
O(n²)
(n
ConFig
N/A
O(n²)
(
DCGD
59
N/A
O(n²)
O(n?)
In Table we compare various popular methods to tackle the identified directional gradient conflicts in the context of PINNs. While all these methods are theoretically promising, some face practical limitations in the context of complex PDE systems. L-BFGS unsuitable for large-scale or stochastic training, as gradient noise disrupts its Hessian updates and line search procedures. First-order variants such as ConFiG (58) and DCGD (59] attempt to alleviate conflicts via gradient surgery or projection. These methods yield incremental improvements but remain constrained by the inherent limitations of first-order updates.
Natural gradient descent (NGD) requires computing and inverting the Fisher information matrix at every iteration and is restricted to float64 precision, which is inefficient on GPUs and consumes 2x more memory while being 2-4x slower than float32. As result, NGD has only been demonstrated on relatively simple benchmarks with smooth solutions, where very small networks suffice and convergence issues rarely appear. On more challenging PDEs with sharp transitions complex dynamics, NGD fails to scale: it is highly sensitive to hyperparameters, lacks hani-batching support, and often diverges.
To illustrate the limitations of NGD, we revisited the 2D Poisson and heat benchmarks of (37 using their official implementation and comparing against SOAP under both float32 and float64

---

precision. We tested MLPs of varying depth and width, reporting the best results across ten random seeds for each method. While SOAP reliably converged, NGD exhibited some failures for some random seeds, which is also acknowledged in [37] Even under successful runs, NGD underperforms SOAP both accuracy and stability, as summarized in 2
Table 2: Comparison of relative L errors for different MLP architectures and optimization methods (NGD vs. SOAP) on Poisson and Heat 2D PDE benchmarks, evaluated in both float32 and float64 precision.
PDE
Architecture (Method)
Float32
Float64
[2. 32, 1J (NGD)
4.87 x 10 3.61
10
3.84 10 : 2.92 10-7
[2, 32, 32, 1] (NGD)
1.65 x 10 +.993 x 10
1.27 x 10 2.86 x 10
[2. 32, 32, (NGD)
2,48 10-000111.00 10
3.14 100000---
10-7
Poisson
[2, 256, 1] (NGD)
1.56 x 10-1
10-2
6.10 x 10-7
10-7
[2, 256, 11 (SOAP)
3.06 10 + 10
6.08
7 2.13
[2. 256, 256, 1J (SOAP)
1.87 x 1,111 10-7
4.06 x 10- 1 10-7
[2. 256, 256, 256. (SOAP)
1.35 x 105 x 10-7
2.99 x 10-1111000
10
12, 32, 11 (NGD)
5.98 x 10 10 2 x 10
7.68 x + 10
[2. 32, 32, 1] (NGD)
x.95 10x10022..20
10-1
2.32 x 10-6 10-66
10-6
[2, 32, 32, 32, 1] (NGD)
5.95 x 10 + x
5.13 x 10 6.2. 10
Heat 2D
[2. 256, 1J (NGD)
5.95 x 10-1 1.05 x 10-3
8.69 x 10 1..4 10-6
[2, 256, 1J (SOAP)
4.61 x 1 x 10-7
8.122
3.03 x 10 6.06
10-7
[2, 256, 256, 11 (SOAP)
2.74 10 +.552 x 10-7
2.04 x 10 6 4.10
10
[2. 256, 256, 256. 1] (SOAP)
2.65 10 5.00
10
1.33 10 2.61 10
In contrast, quasi-second second-order methods (e.g., SOAP M. Kron 633 Muon (62 naturally promote directional gradient conflicts via preconditioning and mitigate ill-conditioning in the loss landscape [67] while maintaining computational tractability. As previously illustrated in Figurd3 SOAP consistently achieves the highest gradient alignment scores, likely due its closer approximation to Newton's 's method an observation further supported by our theoretical analysis in Appendix D
Building these insights, now present comprehensive numerical experiments that demonstrate the superior accuracy and convergence behavior of quasi-second-order methods across a diverse set of PDE benchmarks.
Experiments
To rigorously evaluate the performance of the aforementioned quasi second-order methods, examine a diverse set of 10 representative and challenging PDEs that govern fundamental physical phenomena. These equations span wave propagation, shock formulation, chaotic systems, reaction- diffusion processes, fluid dynamics, and heat transfer. The detailed description of the problem setup, including the PDE parameters, initial and boundary conditions, numerical implementations, and supplementary visualizations, are presented in Appendix G
Baselines. We focus comparisons PINN approaches that scalable to large neural networks, as scalability is essential for solving realistic, large-scale physical problems. Consequently, we exclude methods based natural gradients (37 38) and L-BFGS variants [40] these typically rely on full-batch training, require high-preedt gighhprecision (e.g., float64) computation, and are limited to small
network sizes sizes-making them impractical for the complex PDE benchmarks considered in this work.
Our baseline setup is based on the current state-of-ther thet--t training pipeline proposed by [27] Specif- ically, we adopt PirateNet [27] the backbone architecture, which is known for its stability and scalability to deeper networks. All weight matrices are initialized using random weight factorization
(RWF) [69]. Exact periodic boundary conditions strictly enforced when applicable [70].
For model training, we use mini-batch gradient descent with the Adam optimizer [71), which has become the de facto standard for training PINNs due to its robust performance and computational
efficiency. To improve training efficiency and robustness, we use learning rate annealing 21 72
for loss balancing. In addition, we employ causal training 77 72 to address causality violations

---

Spectral Energy Distribution
10
10
E(k)
10
128
128
256
512
10
Wavenumber (k)
010.5.V
0.5
VT.. 0.5)
C-
[-4
[-6
Em8
Temperature
0.4
02
00
0.5
10
Rebeexa
Policd
Ret 2?? Eree
10
Temperaturs
16
0.6
Figure Simulating complex fluid dynamics using PINNs with SOAP optimization. (a) Kolmogorov flow at Re=10,000; comparison between reference solution and PINN predictions demonstrates accurate capture of turbulent structures across multiple time steps. (b) Spectral energy distribution showing PINN's superior resolution of fine-scale dynamics compared to traditional numerical solu- tions various grid resolutions. (c) Lid-driven cavity flow at Re=5,000: streamlines and centerline velocity profiles show excellent agreement with benchmark data from [68]. (d) Kuramoto-Sivashinsky equation: PINNs accurately predicts complex spatiotemporal patterns and chaotic dynamics. (e) Soslligh-Tayloon instability (Pr=0.71, Ra=10%): evolution of temperature field shows precise capture of interface dynamics and mushroom-shan shaped structures characteristic of this flow.
when solving time-dependent PDEs. For challenging benchmarks, we implement time-marching and curriculum learning strategies +p
Importantly, this baseline represents one of the most accurate and scalable PINN pipelines currently available, achieving state-of-the-att performance across diverse set of PDE benchmarks [72] 74). For instance, on the standard Allen-Cahn benchmark, which has been extensively evaluated by various PINN methods, proposed baseline outperforms most existing approaches, demonstrated in Table g A comprehensive description of the techniques employed and the hyperparameter configurations provided in Appendices G.T and G.2 respectively.
State-of-the-art results. Table highlights the consistent performance improvements of quasi- second-orde optimizers across a wide range of PDE benchmarks. Among them, SOAP achieves the best overall results, likely due to its closer alignment with Newton's method (Table I making it particularly well-suited for the multi-objective nature of PINN optimization.
Compared to our baselines, SOAP reduces relative error by 6.4x on the wave equation. For nonlinear ID problems, it achieves 6.9x improvement on Allen-Cahn equation, and approximately 2x improvements both the Korteweg-de Vries and Kuramoto-Sive saanooo-Sivashinsky equations. The performance gains become particularly pronounced for coupled diffusion-reaction systems. The Grey-Scott and -inzburg-Landau
Ginzburg- Landau equations exhibit order of magnitude reduction in error. On challenging Navier-Stokes benchmarks, including the lid-driven cavity and Rayleigh-Taye instability problems, SOAP demonstrates more than 10x error reduction. We highlight and discuss these substantial improvements in detail below.
Complex fluid dynamics. Our most significant achievement is successfully applying PINNs to complex fluid dynamics problems that were previously considered beyond their capabilities. In

---

Table Comparison of optimizer performance obtained by training PINNs with Adam, Adam+L.- BFGS, and SOAP, respectively, across various PDEs, following the training pipeline described in The evaluation metric is the relative L error over the entire spatial-temporal domain.
Section
Benchmark
Adam
Adam+I.BFGS
Muon
SOAP
Wave
5.15 10-5
5.08 10-5
8.62 x 10-6
9.34 10
8.05 x 10 -
Burgers
8.20 10-5
8.20 10-5
4.85 10
4.52 10
4.03 10-
Allen-Cahn
2.24 10
2.25 10
3.63 10
4.95 10
3.48 x 10
Korteweg-Do De Vries
7.04 10-4
7.33 10-4
5.48 x 10-4
4.19 10-4
3.40 x 10-4
Kuramoto-Svvahhinsky
7.48 10-²
5.49 x 10
3.51 x 10
3.86 x 10-2
Grey-Scoll
3.61 x 10
1.89 x 10
1.95 10
1.88 x 10
Ginzburg-Landee
1.49 10-2
7.53 x 10-3
4.58 10-5
4.78 x 10-3
Lid-driven cavity (Re 5 10")
3.24 10-1
7,05 10
6.70 10-
3.99 x 10
Kolmogorov flow (Rc 10*)
2.04 10
8.62 x 10
6.89 10
3.20 x 10
Rayleigh- Taylor instability (Ra 10%)
7.32 10
5.74 x 10
1.80 10 5.22 x 10
particular, we demonstrate breakthrough results in three challenging cases that combine multiple physical constraints and have historically proven difficult for PINNs, see Figure
For the lid-driven cavity flow at Reynolds number 5,000, SOAP enables a dramatic improvement in accuracy, reducing the relative L error from 42.4% to 3.99%. As shown in Figure B our model successfully captures intricate flow features including secondary and tertiary corner vortices, showing excellent agreement with the benchmark results of [68].
The Rayleigh-Tayy Taylor instability presents an even more challenging test, requiring simultaneous handling of interface dynamics and coupled velocity denssty evolution. SOAP enables accurate prediction of the characteristic mushroom-shaper structures that develop as heavier fluid penetrates into lighter fluid, achieving relative L error of 0.52% nearly an order of magnitude improvement over the best baseline's is 7.32%. Figure 3 demonstrates excellent agreement with reference solutions across multiple time steps, capturing both the initial linear growth phase and subsequent nonlinear development.
Our most impressive result comes from the turbulent Kolmogorov flow at Reynolds number 10,000 marking the first time PINNs have successfully captured turbulent dynamies at such high Reynolds numbers. Our model achieves relative 2 eтror of 3.20%, compared to 20.4% with our baseline. Figure 33 shows that our predictions accurately reproduce both the scage flow structures and the complex cascade of smaller eddies characteristic of turbulent flows. Moreover, spectral analysis reveals that our PINN solution maintains higher spectral energy at high wavenumbers compared to traditional numerical solvers, even those using a 1024 1024 grid resolution. This demonstrates PINNs' potential advantage in resolving fine- scale dynamics without requiring explicit grid discretization key capacity for turbulence modeling.
Ablation studies. We conduct systematic experiments to evaluate SOAP's performance across different architectures and hyperparameter settings, establishing the robustness of our approach. Our first investigation examines SOAP's effectiveness across three representative architectures: standard MLP. modified MIP (21, and PirateNet (27). As shown in the top panel of Figure testing each architecture on four benchmark PDEs (Wave, Burgers, Allen-Cahn, and KdV equations), we find that SOAP consistently improves accuracy compared to Adam regardless of the underlying network backbones. In particular, PirateNet seems to be the most effective architecture across all test cases, leading to its selection for our main experiments.
As illustrated in the bottom panel of Figure 4 our results of SOAP's hyperparameters reveal two critical factors affecting performance. The preconditioner update frequency presents a clear trade-off between accuracy and computational cost. While more frequent updates yield better results, the improvements diminish beyond an update frequency of 2, which we selected as the optimal balance for our experiments. The momentum parameter 31 proved especially crucial: high momentum (3) = 0.99) consistently achieves the best results, while low momentum (3,=0.01) significantly degrades performance across all test cases.

---

Wave
Burgets
Allea-Cuha
Kurteweg Vries
10
ConFIG
10
Rel L22 Error
SOAP
10
10
10
10"
10
10
MLP
Modillal MLP
PinsteNot
MLP
Modified MLP
PisteNer
10
MLP
Modifial MLP
Pimtteeer
MLP
Modifial MUPP
PasteNe
10
10
Wave
Rel L3 Error
10
Burgers
Allen Allen-Cahn
10"
Kdv
10"
10"
10
10
10
0.01
0.90
0.99
Frequency
Momnentum B1
Figure 4: Optimizer performance comparison and ablation studies. Top: Relative L error across PDE benchmarks using different optimizers. Bottom left: Relative I error for varying preconditioner update frequencies in SOAP optimizer. Bottom right: Relative L error with different momentum values in SOAP optimizer.
For completeness, we also compared SOAP against ConFiG [58], a recently proposed method for addressing gradient conflicts in PINNs that has demonstrated relative good performance compared to established baselines such as PCGirad [7] and IMTL-G (76) in multi-task learning. Despite its promising theoretical foundations, ConFiG showed some sensitivity to hyperparameter settings in experiments, resulting in inconsistent performance. These results highlight the practical advantages of SOAP's more robust optimization approach.
Computational costs. While SOAP requires approximately 2x longer training time compared to baselines (Table 7) our focus is exploring the performance frontier of PINNs through extended training to full convergence. Impressively, error and loss convergence curves (Appendix G.6) indicate that SOAP typically achieves rapid initial convergence, reaching a reasonable accuracy (approximately 10-+) within the first 10,000 iterations, followed by gradual error reduction in subsequent iterations. This suggests the potential for reducing training time by up 10x while maintaining competitive performance. These findings motivate future research into designing computationally efficient optimization algorithms and training strategies for PINNs, paving the way for practical and scalable applications in complex physics simulations.
7 Conclusion
This work advances our understanding of gradient conflicts in training of PINNs. We proposed gradient alignment scores as quantitative measure of such conflicts, and provide both theoretical analysis and empirical evidence demonstrating their prevalence during PINN optimization. Fur- thermore, we show that second-order optimization methods implicitly promote gradient alignment, offering a principled approach to mitigating these conflicts. Particularly, we uncover a connection between SOAP and Newton's 's method. Extensive experiments across 10 PDE benchmarks confirm the effectiveness of quasi second-order optimizers, with SOAP achieving state-of-the-e performance.
Building on these insights, several promising research directions emerge. While SOAP demonstrates the power of gradient alignment in handling coupled physical constraints, opportunities exist for more efficient preconditioned algorithms that maintain their effectiveness with reduced computational cost. More broadly, our work suggests that the principles of gradient alignment and second-order preconditioning could benefit many deep learning applications involving competing objectives, though challenges remain in scaling these approaches to larger systems. Success in these directions could transform both scientific computing and multi-task optimization.
10

---

Acknowledgment
B.L. would like to acknowledge support from National Key R&D Program of China Grant No. 2024YFA10160M0 P.P. and S.W. acknowledge support from the US Department of Energy under the Advanced Scientific Computing Research program (grant DE-SC0024563), the Nvidia Academic Grant Program, and the Institute for Foundations of Data Science at Yale University. We also thank the developers of the software that enabled our research, including JAX [77], Matplotlib 788. and NumPy [79].
References
[1] Nikhil Vyas, Depen Morwani, Rosie Zhao, Itai Shapira, David Brandfonbrener, Lucas Janson, and Sham Kakade. Soap: Improving and stabilizing shampoo using adam. arXiv preprini arXiv:2409.1 11321, 2024.
[2] Maziar Raissi, Alireza Yazdani, and George Em Karniadakis. Hidden fluid mechanics: Learning velocity and pressure fields from flow visualizations. Science, 367(6481 (1:6448111066610333 2020.
[3] Muhammad M Almajid and Moataz O Abu-Al-Saud. Prediction of porous media fluid flow using physies informed neural networks. Journal of Petroleun Science and Engineering, 109205 2022.
[4] Hamidreza Eivazi, Mojtaba Tahani, Philiip Schlatter, and Ricardo Vinuesa. Physics-informed neural networks for solving reynolds-averaged navier-stokes equations. Physics of Fluids, 34(7), 2022.
[5] Zhen Cao, Kai Liu, Kun Luo, Sifan Wang. Liang Jiang, and Jianren Fan. Surrogate modeling of multi-dimensional premixed and non-premixed combustion using pseudo-time stepping physics-informed neural networks. Physics of Fluids, 36(11), 2024.
[6] Jiaxuan Xu, Han Wei, and Hua Bao. Physics-informet neural networks for studying heat transfer in porous media. International Journal of Heat and Mass Transfer, 217:124671, 2023.
[7] Hassan Bararnia and Mehdi Esmaeilpour. On the application of physics informed neural networks (pinn) to solve boundary layer thermal-fluid problems. International Communications Heat and Mass Transfer, 132:105890, 2022.
[8] Gargya Gokhale, Bert Claessens, and Chris Develder. Physics informed neural networks for control oriented thermal modeling of buildings. Applied Energy, 118552, 2022.
[9] Georgios Kissas. Yibo Yang, Eileen Hwuang, Walter Witschey, John A Detre, and Paris Perdikaris. Machine learning in cardiovascular flows modeling: Predicting arterial blood pres- sure from non-invasive 4D flow MRI data using physics-informed neural networks. Computer Methods in Applied Mechanics and Engineering, 358: 12623, 2020.
[10] Xuelan Zhang, Baoyan Mao, Yue Che, Jiaheng Kang, Mingyao Luo, Aike Qiao, Youjun Liu, Hitomi Anzai, Makoto Ohta, Yuting Guo, et al. Physics-informed neural networks (pinns) for 4d hemodynamics prediction: investigation of optimal framework based on vascular morphology. Computers in Biology and Medicine, 164: 164:107287, 2023.
(11) Federica Caforio, Francesco Regazzoni, Stefano Pagani, Elias Karabelas, Christoph Augustin, Gundolf Haase, Gernot Plank, and Alfio Quarteroni. Physics-informed neural network esti- mation of material properties in soft tissue nonlinear biomechanical models. Computational Mechanics, pages 1-27, 2024.
[12] Enrui Zhang, Ming Dao, George Em Karniadakis, and Subra Suresh. Analyses of internal structures and defects in materials using physics-informed neural networks. Science advances, 8(7)xcabk0644, 2022.
[13] Hyogu Jeong, Jinshuai Bai, Chanaka Prabuddha Satuwatta-Gamaee Charith Rathnayaka, Ying Zhou, and YuanTong Gu. physics-informed neural network-based topology optimization (pinnto) framework for structural optimization. Engineering Structures, 278: 15484, 2023.
11

---

[14] Haoteng Hu, Lehua Qi, and Xujiang Chao. informed neural networks (pinn) for com- putational solid mechanics: Numerical frameworks and applications. Thin walled Structures. page 112495, 2024.
[15] Alexander Kovacs, Lukas Exl, Alexander Kornell, Johann Fischbacher, Markus Hovorka, Markus Gusenbauer, Leoni Breth, Harald Oezelt, Masao Yano, Noritsugu Sakuma, et al. Conditional physics informed neural networks. Communications in Nonlineur Science and Numerical Simulation, 104:106041, 2022.
[16] Arbaaz Khan and David Lowther. Physics informed neural networks for electromagnetic analysis. IEEE Transactions Magnetics, 58(9):1 4, 2022.
[17] Marco Baldan, Paolo Di Barba, and David A Lowther. Physics-informet neural networks for inverse electromagnetic problems. IEEE Transactions on Magnetics, 19,5):1-5, 2023.
[18] Jonthan D Smith, Zachary E Ross, Kamyar Azizzadenesheli, and Jack B Muir. Hyposvi: Hypocentre inversion with stein variational inference and physics informed neural networks. Geophysical Journal International, 228(1):698-710) 2022.
[19] Chao Song and Yanghua Wang. Simulating seismic multifrequeney wavefields with the fourier feature physics-informet neural network. Geophysical Journal International, 2303 1514, 2023.
[20] Pu Ren, Chengping Rao, Su Chen, Jian-Xun Wang. Hao Sun, and Yang Liu. Seismicnet: Physics- informed neural networks for seismic wave modeling in semi-infinite domain. Computer Physics Communications, 295: 109010, 2024.
[21] Sifan Wang, Yujun Teng, and Paris Perdikaris. Understanding and mitigating gradient flow pathologies in obysics-informed neural networks. SIAM Journal on Scientific Computing, 43(5):A3055-3308) 2021.
[22] Vincent Sitzmann, Julien Martel, Alexander Bergman, David Lindell, and Gordon Wetzstein. Im- plicit neural representations with periodic activation functions. Advances in Neural Information Processing Systems, 33:7462-7473, 2020.
[23] Rizal Fathony, Anit Kumar Sahu, Devin Willmott, and J Zico Kolter. Multiplicative filter networks. In International Conference on Learning Representations, 2021.
[24] Ben Moseley, Andrew Markham, and Tarje Nissen-Meyer. Finite basis physies-iformmed neural networks (fbpinns): scalable domain decomposition approach for solving differential equations. arXiv preprint arXiv:2 107.0777, 2021.
[25] Namgyu Kang, Byeonghyeon Lee, Youngjoon Hong, Seok-Bae Yun, and Eunbyung Park. Pixel: Physics-informed cell representations for fast and accurate pde solvers, arXiv preprint arXiv:2207. /00../28000 2800, 2022.
[26] Junwoo Cho, Seungtae Nam, Hyunmo Yang, Seok-Bae Yun, Youngjoon Hong, and Eunbyung Park. Separable physics-informed neural networks. Advances in Neural Information Processing Systems, 36, 2024.
[27] Sifan Wang, Bowen Li, Yuhan Chen, and Paris Perdikaris. Piratenets: Physics-informat deep learning with residual adaptive networks. arXiv preprint arXiv:2402.00322. 2024.
[28] Ameya D Jagtap, Kenji Kawaguchi, and George Em Karniadakis. Adaptive activation functions accelerate convergence in deep and physics-infomme neural networks. Journal of Computational Physics, 404:109113 2020.
[29] Jassem Abbasi and Pål Østebø Andersen. Physical activation functions (pafs): An approach for more efficient induction of physics into physics-informed neural networks (pinns). Neurocom- puting, 608: 608:12835,, 2024.
[30] Sifan Wang, Hanwen Wang, and Paris Perdikaris. On the eigenvector bias of fourier feature networks: From regression to solving multi-scale PDEs with physics-informed neural networks. Computer Methods in Applied Mechanics and Engineering, 384:113938, 2021.
12

---

[31] Francisco Sahli Costabal, Simone Pezzuto, and Paris Perdikaris. 6-pinns: physics-informes neural networks on complex geometries. Engineering Applications of Artificial Intelligence. 127:107324, 2024.
[32] Chengxi Zeng, Tilo Burghardt, and Alberto M Gambaruto. Rbf-pinn: Non-fourier positional embedding in physics-innieee informed neural networks. arXiv preprint arXiv:2402.08367, 2024.
[33] Xinquan Huang and Tariq Alkhalifah. Efficient physics-informed neural networks using hash encoding, Journal of Computational Physics, page 112760, 2024.
[34] Mohammad Amin Nabian, Rini Jasmine Gladstone, and Hadi Meidani. Efficient training of physies-informed neural networks via importance sampling. Computer-Aided Civil and Infrastructure Engineering. 2021.
[35] Arka Daw, Jie Bu, Sifan Wang, Paris Perdikaris, and Anuj Karpatne. Rethinking the importance of sampling in physics-informe neural networks. arXiv preprint arXiv:2207.023338 2022.
[36] Chenxi Wu, Min Zhu, Qinyang Tan, Yadhu Kartha, and Lu Lu. A comprehensive study of non- adaptive and residual- based adaptive sampling for physics- informed neural networks. Computer Methods in Applied Mechanics and Engineering, 403:1 15671, 2023.
[37] Johannes Müller and Marius Zeinhofer. Achieving high accuracy with pinns via energy natural gradient descent. In International Conference on Machine Learning, pages 25471-25488 PMLR, 2023.
[38] Anas Jnini, Flavio Vella, and Marius Zeinhofer. Gauss- natural gradient descent for physics-informed computational fluid dynamics. arXiv preprint arXiv:2402. 10680, 2024.
[39] Yongcun Song, Xiaoming Yuan, and Hangrui Yue. The admm-pin algorithmic framework for nonsmooth pde-constrained optimization: a deep learning approach. SIAM Journal on Scientific Computing, 46(6): 2024.
[40] Jorge F Urbán, Petros Stefanou, and José A Pons. Unveiling the optimization process of physics informed neural networks: How accurate and competitive can pinns be? Journal of Computational Physics,
2025.
[41] Colby Wight and Jia Zhao. Solving Allen-Cahn and Cahn-Hilliard equations using the adaptive physics informed neural networks. arXiv preprint arXiv:2007.04542, 2020.
[42] Aditi S Krishnapriyan, Amir Gholami, Shandian Zhe, Robert M Kirby, and Michael W Mahoney. Characterizing possible failure modes in chysics-informed neural networks. arXiv preprint arXiv:2109001050, 2021.
[43] Wenbo Cao and Weiwei Zhang. Tsonn: Time-stepping-orrennee neural network for solving partial differential equations. arXiv preprint arXiv:2310. 2023.
[44] Shaan Desai, Marios Mattheakis, Hayden Joy, Pavlos Protopapas, and Stephen Roberts. One- shot transfer learning of physics-informed neural networks. arXiv preprint 11X.1:2110.112866 2021.
[45] Somdatta Goswami, Cosmin Anitescu, Souvik Chakraborty, and Timon Rabczuk. Transfer learn- ing enhanced physics informed neural network for phase-field modeling of fracture. Theoretical and Applied Fracture Mechanics, 106:17 2020.
[46] Souvik Chakraborty. Transfer learning based multi-fidelity physics informed deep neural network. Journal of Computational Physics, 426: 109942, 2021.
[47] Pao-Hsiung Chiu, Jian Cheng Wong, Chinchun Ooi, My Ha Dao, and Yew-Soon Ong. Can-pinn differentiation
fast physics-informed neural network based on ponpllddautmmati-nnnmeee method. Computer Methods in Applied Mechanics and Engineering, 395: 14909, 2022.
[48] Ehsan Kharazmi, Zhongqiang Zhang, and George Em Kariaddkis.. hp-vpinns: Variational physies-informed neural networks with domain decomposition. Computer Methods Applied Mechanics and Engineering.
2021.
13

---

[49] Ravi G Patel, Indu Manickam, Nathamiel A Trask, Mitchell A Wood, Myoungkyu Lee, Ignacio Tomas, and Eric C Cyr. namiicall consistent inffsiee neural networks for hyperbolic systems. Journal of Computational Physics, 449:1 10754, 2022.
[50] Jeremy Yu, Lu Lu, Xuhui Meng, and George Em Karniadakis. Gradieet-enhaneee physics- informed neural networks for forward and inverse pde problems. Computer Methods Applied Mechanics and Engineering, 393:1 14823, 2022.
[51] Hwijae Son, Jin Woo Jang, Woo Jin Han, and Hyung Ju Hwang. Sobolev training for physics informed neural networks. arXiv preprint arXiv:2101 088932, 2021.
[52] Sifan Wang, Xinling Yu, and Paris Perdikaris. When and why PINNS fail to train: A neural tangent kernel perspective. Journal of Computational Physics, 449: 1 10768, 2022.
[53] Wensheng Li, Chao Zhang, Chuncheng Wang, Hanting Guan, and Dacheng Tao. Revisiting pinns: Generative adversarial bhysics-informed neural networks and point-weighting method. arXiv preprini arXiv:2205.0877 2022.
[54] Wenqian Chen, Amanda A Howard, and Panos Stinis. Self-adaptive weights based on balanced residual decay rate for physics-informed neural networks and deep operator networks. arXiv preprint arXiv:2407.01613, 2024.
[55] Sokratis J Anagnostopoulos, Juan Diego Toscano, Nikolaos Stergiopulos, and George Em Karniadakis. Residual-basee attention in physies-informed neural networks. Computer Methods in Applied Mechanics and Engineering, 421:1 16805, 2024.
(56) Li Liu, Shengping Liu, Hui Xie, Fansheng Xiong, Tengchao Yu, Mengjuan Xiao, Lufeng Liu, and Heng Yong. Discontinuity computing using physics-informed neural networks. Journal of Scientific Computing, 98(1):22, 2024.
[57] Jiahao Song, Wenbo Cao, Fei Liao, and Weiwei Zhang. Vw-pinns: A volume weighting method for pde residuals in physics-informed neural networks. Acta Mechanica Sinica, 41(3):224148 40, 2025.
[58] Qiang Liu, Mengyu Chu, and Nils Thuerey. Config: Towards conflict-free training of physics informed neural networks. arXiv preprini arXiv:2408. /114, 2024.
[59] Youngsik Hwang and Dongyoung Lim. Dual cone gradient descent for training physics-informed neural networks. Advances in Neural Information Processing Systems, 37:98563-98595, 2024.
[60] Maziar Raissi, Paris Perdikaris, and George E Karniadakis. Physics- infoomed neural networks: deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational Physics, 378:686-707, 2019.
[61] Andreas Griewank and Andrea Walther. Evaluating derivatives: principles and techniques of algorithmic differentiation. SIAM, 2008
[62] Keller Jordan, Yuchen Jin, Vlado Boza, Jiacheng You, Franz Cesista, Laker Newhouse, and Jeremy Bernstein. Muon: An optimizer for hidden layers in neural networks, 2024.
[63] Xi-Lin Li. Preconditioned stochastic gradient descent. IEEE transactions on neural networks and learning systems, 2955): 1454-146, 2017.
[64] Shun-Ichi Amari. Natural gradient works efficiently in learning. Neural computation, 10(2):251-
276, 1998.
[65] Dong C Liu and Jorge Nocedal. On the limited memory bígs method for large scale optimization. Mathematical progrumming, 45(1):503-528, 1989.
[66] Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor optimization. In International Conference on Machine Learning, pages 1842-1850. PMLR, 2018
[67] Pratik Rathore, Weimu Lei, Zachary Frangella, Lu Lu, and Madeleine Udell. Challenges in training pinns: A loss landscape perspective. arXiv preprint arXiv:2402.01868 2024.
14

---

[68] U Ghia, K.N Ghia, and C.T Shin. High-re solutions for incompressible flow using the navier- stokes equations and a multigrid method. Journal of Computational Physies, 48(3):387-411, 1982.
[69] Sifan Wang, Hanwen Wang, Jacob H Seidman, and Paris Perdikaris. Random weight factorization improves the training of continuous neural representations. arXiv preprint arXiv:2210.0127.. 2022.
[70] Suchuan Dong and Naxian Ni. A method for representing periodic functions and enforcing exactly periodic boundary conditions with deep neural networks. Journal of Computational Physics, 435:1 10242, 2021.
[71] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
[72] Sifan Wang, Shyam Sankaran, Hanwen Wang, and Paris Perdikaris. An expert's guide to training physics-informed neural networks. arXiv preprint arXiv:2308.08400 2023.
[73] Sifan Wang, Shyam Sankaran, and Paris Perdikaris. Respecting causality is all you need for training physics-informed neural networks. arXiv preprint arXiv:2203.07404, 2022.
[74] Zhongkai Hao, Jiachen Yao, Chang Su, Hang Su, Ziao Wang, Fanzhi Lu, Zeyu Xia, Yichi Zhang. Songming Liu, Lu Lu, et al. Pinnacle: A comprehensive benchmark of physics-informed neural networks for solving pdes. arXiv preprint arXiv:2306. 08827, 2023.
[75] Tianhe Yu, Saurabh Kumar, Abhishek Gupta, Sergey Levine, Karol Hausman, and Chelsea Finn. Gradient surgery for multi-task learning. Advances in Neural Information Processing Systems, 33:5824-588 5836, 2020.
[76] Liyang Liu, Yi Li, Zhanghui Kuang, J Xue, Yimin Chen, Wenming Yang, Qingmin Liao, and Wayne Zhang. Towards impartial multi-task learning. iclr, 2021.
[77] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018.
[78] John D Hunter. Matplotlib: A 2D graphics environment. IEEE Annals of the History of Computing, 903))900-95 95, 2007.
[79] Charles R Harris, K Jarrod Millman, Stéfan J van der Walt, Ralf Gommers, Pauli Virtanen, David Cournapeau, Eric Wieser, Julian Taylor, Sebastian Berg, Nathaniel J Smith, et al. Array programming with numpy. Nature, 585(7825):357-362 2020.
[80] Hanxu Zhou, Zhou Qixuan, Tao Luo, Yaoyu Zhang, and Zhi-Qin Xu. Towards understanding the condensation of neural networks at initial training. Advances in Neural Information Processing Systems, 55:2184-2196, 2022.
[81] Zheng-an Chen and Tao Luo. On the dynamics of three-layer neural networks: initial condensa- tion. arXiv preprint arXiv:2402. 15958, 2024.
[82] Anas Barakat and Pascal Bianchi. Convergence and dynamical behavior of the adam algorithm for nonconvex stochastic optimization. SIAM Journal on Optimization, 31(1):244-277 2021.
[83] Lukas Balles and Philipp Hennig. Dissecting adam: The sign, magnitude and variance of stochastic gradients. In International Conference Machine Learning, pages 404-413. PMLR, 2018.
[84] Adepu Ravi Sankar, Yash Khashage, Rahul Vigneswaran, and Vineeth N Balasubramanian. A deeper look at the hessian cigenspectrum of deep neural networks and its applications to regularization. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pages 9481-9488, 2021.
[85] Igor Molybog, Peter Albert, Moya Chen, Zachary DeVito, David Esiobu, Naman Goyal, Punit Singh Koura, Sharan Narang, Andrew Poulton, Ruan Silva, et al. theory on adam instability in large-scale machine learning. arXiv preprint XXiv:2304.099777 2023.
15

---

[86] Depen Morwani, Itai Shapira, Nikhil Vyas, Eran Malach, Sham Kakade, and Lucas Janson. A new perspective on shampoo's preconditioner. arXiv preprint arXiv:2406.17748, 2024.
[87] Rohan Anil, Vineet Gupta, Tomer Koren, Kevin Regan, and Yoram Singer. Scalable second order optimization for deep learning. arXiv preprint arXiv:2002.09011 2020.
[88] Sokratis J Anagnostopoulos, Juan Diego Toscano, Nikolaos Stergiopulos, and George Em Karniadakis. Residual- attention and connection to information bottleneck theory in pinns. arXiv preprint arXiv:2307.00377 2023.
[89] Xenofon Karakonstantis, Diego Caviedes-Nozal, Antoine Richard, and Efren Fermandez-Gramee Room impulse response reconstruction with physics-informed deep learning. The Journal of the Acoustical Society of America, 155(2):0048-105) 2024.
[90] Jiaming Zhang, David Dalton, Hao Gao, and Dirk Husmeier. Physics-informed deep learning based on the finite difference method for efficient and accurate numerical solution of partial differential equations.
[91] Matthew Tancik, Pratul P Srinivasan, Ben Mildenhall, Sara Fridvvic-Keil, Nithin Ragha- van, Utkarsh Singhal, Ravi Ramamoorthi, Jonathan T Barron, and Ren Ng. Fourier features let networks learn high frequency functions in low dimensional domains. arXiv preprint arXiv:2006. 10739, 2020.
[92] Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedfor- ward neural networks. In Proceedings of the thirteenth international conference on artificial intelligence and statistics, pages 249-256, 2010.
[93] Tim Salimans and Durk P Kingma. Weight normalization: A simple reparameterization to accelerate training of deep neural networks. Advances in neural information processing systems, 29, 2016.
[94] Tobin A Driscoll, Nicholas Hale, and Lloyd N Trefethen. Chebfun guide, 2014.
[95] Syver Døving Agdestein, Simone Ciarella, and Benjamin Sanderse. IncompressibleNavier Stokes.jl, November 2024.
[96] Levi McClenny and Ulisses Braga-Neto. Self-adaptive physics-informed neural networks using a soft attention mechanism. arXiv preprint arXiv:2009.04544 2020.
[97] Revanth Mattey and Susanta Ghosh. A novel sequential method to train physics informed neural networks for allen cahn and cahn hilliard equations. Computer Methods in Applied Mechanics and Engineering, 390: 14474, 2022.
[98] Vasiliy A Es' kin, Danil V Davydov, Ekaterina D Egorova, Alexey O Malkhanov, Mikhail A Akhukov, and Mikhail Smorkalov. About optimal loss function for training thisics-iformed neural networks under respecting causality. arXiv preprint 2022:2230..02282, 2023.
[99] Norman Zabusky and Martin D Kruskal. Interaction of" solitons" in collisionless plasma and the recurrence of initial states. Physical review letters, 15(6):240, 1965.
16

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The abstract and introduction accurately reflect the paper's contributions, including: introducing gradient alignment score to quantify directional conflicts (Sec. 1-2), demonstrating how second-order optimizers enhance gradient alignment (Sec. 3), establishing SOAP's connection to Newton's method (Sec. 4), and showing improved performance across 10 PDE benchmarks (Sec. 5).
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: The paper discusses computational costs Section 5 (paragraph "Computa- tional costs"), noting that SOAP requires approximately 2x longer training time compared baselines, though achieves better convergence. The conclusion also acknowledges the need for "more efficient preconditioned algorithms that maintain their effectiveness with reduced computational cost.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best
17

---

judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: The paper includes theoretical results with clear assumptions and proofs. Proposition Section 3 establishes the connection between alignment score and cosine similarity. Proposition 2 analyzes gradient alignment at initialization. Proposition 3 connects preconditioned gradient descent to Newton's method. Complete proofs are provided in the appendix (referenced throughout Section 3-4).
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions the paper (regardless whether the code and data are provided or not)?
Question:
Answer: [Yes]
The paper provides sufficient information to reproduce the main experimental results. Section 5 describes the baseline setup, and the appendices (referenced in Sec- tion 5) contain complete information on experimental settings, hyperparameters, and data generation.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., s., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm.
18

---

(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g.. a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [NA]
Justification:
We will release the code once the paper is accepted.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips. cc/
gablic/guides
pondeSubmisssionollcco for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: The paper specifies training and test details in Section 5,
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
19

---

Answer: [Yes]
The paper includes multiple performance comparisons across different PDEs and optimizers, with ablation studies in Figure 4 showing the impact of different hyper- parameter settings. The consistent performance improvements across diverse benchmarks (Table 2) and multiple experimental configurations support the statistical significance of the results.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report I-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative eтror rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: The paper provides information about computational resources in Appendix G.5.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics httpps:/neuurpp.
Answer: [Yes]
Justification: The research conforms to the NeurIPS Code of Ethics. focuses on funda- mental advances in scientific machine learning with no apparent ethical concerns. The work acknowledges all related research properly and presents results transparently.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
20

---

If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive socictal impacts and negative socictal impacts of the work performed?
Answer: [Yes]
Justification: The conclusion discusses broader positive impacts, noting that the principles of gradient alignment could benefit applications involving competing objectives beyond scientific computing. The research enables more accurate simulation complex physical systems, which has positive societal impact for scientific and engineering applications. As with any tool that furthers our understanding and ability to predict the outcomes complex systems, there may be ill- il-innenniooed cases, but do not expect any specific negative impact from this work. Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied partieular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (c.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (c.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: The paper focuses optimization techniques for solving PDEs and does not involve models or data with high risk for misuse. The methods presented are used for scientific computing applications without foreseeable harmful applications.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
21

---

We recognize that providing effective safeguards is challenging, and many papers do require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
The paper properly cites existing work and acknowledges software used (JAX, Matplotlib, Chebfun, and NumPy) in the Acknowledgments section. The baseline implementation builds on prior work that is appropriately cited.
Justification:
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms use in package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification:
We will release code implementing the methods described this paper upon acceptance.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/model as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Question:
Answer: [NA]
Justification: The paper does not involve crowdsourcing or research with human subjects.
focuses on algorithmic improvements and numerical experiments for solving PDEs.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
22

---

Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The research does not involve human subjects, IRB approval is not applica- ble.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods this research? Note that if the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, or originality the research, declaration is not required.
Answer: [NA]
Justification: The paper does not indicate the use of L.LMs as part of the core methods or research.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non non-standard components.
Please refer to our L.LM policy (https: Buttep///eearrpsscccc///eeeeeeep for what should or should not be described.
23

---

A Nomenclature
Table Notation used throughout the paper. We use uppercase letters for matrices and lowercase letters for their vectorized forms. All gradients and Hessians are with respect to the loss function L.
Symbol
Description
C
Loss function
Lios
Initial condition losses
Cres
PDE residual losses
0
Neural network parameters
W
Weight matrix for a given layer Vectorized weight matrix, t Vec(W)
G
Gradient matrix, G = W
Vectorized gradient, 9 = Vec(G)
F
Fisher information matrix
Hessian matrix, H V
Add
Full Adagrad preconditioner matrix
HGN
Gauss-Newtar approximation of the Hessian
Gradient alignment score between loss components
B Analysis of Intra-step Gradient Alignment
We present some preliminary analysis to understand intra-step gradient conflicts in training PINNs via standard gradient descent, Adam (7]). and Shampoo algorithms [66]. and how SOAP can effectively resolve them during training. For simplicity, consider the simplest case of using PINNs with two-layer NN to solve one-dimensional Laplace equation and focus on the analysis of intra-step gradient alignment (3.3) with small initialization. The analysis can be easily extended to without loss generality, we consider
other types PDEs. Following the general setup Section 2 ID Laplace equation as follows
Au= d" =0 on [-1,
(B.1)
u(+ = J+1-
We approximate the solution u(x) by two-layer network with width N:
u(x, 0) Σ a.a (w,.) a o(wz),
(B.2)
where a = aN) = (w), ........... wx) E IR N and 0 = (a. w) E 22N
Moreover, we limit reduces to
ourselves the activation function (Tx) tanh (z). In this case, the loss 2.5
min @=(nn (n,()
((0)
(Epp 0)1 +
[u(8, 0) gal
(B.3)
N
(.)
(se(0)
To analyze the gradient conflict phenomenon in training PINNs, consider the small initialization regime.
Assumption 1.
The weights 0g, w; initialized by i.i.d. Gaussian (0,ce²) with small o(1).
This allows us to introduce the normalized parameters:
a, =
initialized standard Gaussian.

---

Burgers
Inter-step Alignment Score
Intra-step Alignment Score
Relative L Error
0.0
10"
Allea-Caha
0.0
1.0
Korteweg- Vries
Ginzburg-Lande
0.5
10"
-0.5
10"
Navier Stokes
1.00
0.75
0.50
0.25
0.00
Step
Step
Step
Adam
SOAP
Kron
Muon
Adagrd
RMSprop
Figure 5: Gradient alignment and test obtained by training PINNs with different optimizers different PDEs. From left to right: ground truth PDE solution, intra-step gradient alignment scores (Eq. (3.3] inter-step gradient alignment scores (Eq. (3.4)), and test error convergence during training.
Lemma 1. Under small initialization, the gradients the residual and boundary loss terms can be approximated as:
Vo),(0) = GT (a, w) + Oss)),
(B.4)
v0Lsc(0) = 5 w) + O(e²),
(B.5)
where
G (a, w) c,a 0 03 3a C w
(B.6)
bhe (a, w) = (9-1 91) a).
(B.7)
Here à = laa, w = - ww are the normalized parameters, and G", T are the effective gradiens terms.
We remark that these elementary computations also provide insights into the gradient magnitude imbalance discussed in Section noting / = O(s?) while [/VoLc(())] = O(s).
Proof. We recall the Taylor expansions the activation function (xx) tanh(x) and its derivatives for later use:
(<)) = 3
+ +z(2),
=1=))=1-²+ +0(zz)))
(B.8)
d"
2x
+0(r), d"(x) 8z O
3
25

---

The gradient of loss function [()) consists of two parts computed as follows:
Nr Uzr (p. 0)V () 0),
2
Lz)
N
and
V Cbc = Σ (u(s, 0) - (0)))))))VV))))))))))))))))))))) o=
with, thanks to
4.2 and B.8
Vau(r, 0) o [wxx) = swr + O(s²),
V u. u(2 0) a,o'(w,x)xz = Ea,e E² ,'a,w;2'" -(E"))
and
V a, Uax (z, 0) o" o"(wx) w -2 + O
V wy Mrx (x, () a,o (w;xr) ) [ 2-x 20,0 (w;z)) 20;o"(w;x))ir
-65 a,,, 2x O(e²).
We also compute
u(I, 0) ==²a- wer ..a.w O(Et),
and
Maz (x,, 0) Σ 0:00 (w:"(wx (w;z)| 25
T. x + O(5))
For convenience, we define the componentwise power xO = (af,...., xx and product x y = (218/1... ....,,INNN.. NNN for x, E RN By the above computation, it follows that at the initialization, there hold
N,
2
V==(a,)CCr(0
(-25
a.
w w03xp)
(-2e"w03mm
- -6=a w +
222z))+ O(5))
N
p=1
N,
8
03
03 3a 22++(())
a
w
w
N
p=1
and
V (t)
Σ
(c²a Ws O((") - -g(s)) (ews (sws 0(=)), cās caa o wo2s3 + +(()))
s=#1
8g(s) (w,a) + 0
(
We then define constant C. = 8N, Σn 2 and the effective gradients
Ga = (G% aa w), G (a w)) ca w (w 3a ④ 02
(B.9)
and
(a, = (G a, w),
(a) Ca = - 11)1-a - (w,
(B.10)
enabling us to write
o +0(6)) lc
(B.11)
We are now ready to understand the gradient conflict for various optimizers applied to the residual and boundary loss terms separately.
26

---

Proposition 3 Ar initialization, the alignment score converges to a binary random variable in the infinite width limit:
sgn (9-1 91)
with prob. 27
lim A(O(V2)), = O(6²) Cu N N
(B.12)
-ggn
sgn (0-1 - 91) with prob. 2
where 0 GD,, Adam, Shampoo, or Soap denotes the optimizer update rule, and Co is a constant depending the optimizer.
We can see that these optimizers fail to resolve intra-step gradient conflicts in the initialization. aligning with the near-zero initial intra-step gradient scores shown in Figure
Proof. Gradient descent. We start with the standard continuous-time gradient descent:
d0
dt vof(o).
(B.13)
Motivated by 80 81 under small initialization Assumption 1 in the initial stage of training dynamics where the leading-order expansion (B.11) holds for the weights a, w, the gradients VeC, (0) and (s) can be effectively described by G" (a. w) and (* w), respectively, up some scaling factors, then the gradient flow (B.13) can be approximated by the effective dynamics for the normalized parameter 8 (a, w):
dl
d - (s' -- (0) + =Goo
(B.14)
Recalling Definition
1 under our assumptions, we have the intra alignment score:
(Vool, voCooc)) A (G" (0), ()) (0) +
where
Σ w
C. 2 w. W + a₃
A (G), (0), (0), (0) sgn(a w w sgn(g-1 91) v Σ, W. 9a, W,
Then we find that at the initialization, by a,, w. ~N 1) from Assumption I and the law of large numbers,
E W. + + E; a w
6
3
almost surely.
W, 9a,, w.
V15+ + 27
21
Also, by the symmetry of Gaussian, there holds POE, a,w; > 0) = _.. follows that the alignment score A(G (0), the (0))== converges to binary random variable with expectation zero in the infinite width limit:
sgn(g-1 -o)) / sgn 91) /21
with prob. 2 with prob.
A (G" (0), th (0)) A -=
Adam. We now consider the deterministic version of the Adam optimizer 71. recalled below for completeness. Let (xx be differentiable objective function Rd. The Adam iteration is defined by 2n Ty 2n- for on (In, Mns Un) Rd x gd x Rd with (xo, O, 0), where
7(1-a") (am+(1)))mmm))) -1- -10- f(x))
T Tn,
am (1 f())
30 + (1 (1- fV
We still consider the gradient conflict at the initialization, since the Adam dynamics is more compli- cated than the gradient flow one. From (82). we have that starting from (æo, 0, E 123d the Adam dynamics at = 0 satisfies (0) x fzoo). (() x vf(ro) 62 and
V((x))
2(0)
+
fx)) ON e=o(1
V 02
27

---

indicating that at early iterations of Adam, the algorithm performance would be similar to the sign gradient descent [83]. Back our problem (4.3), by the above discussion, if we apply Adam the loss functions L, (0) and (0), respectively, at the initialization, the normalized weights 0 = e 0
will be updated along the directions:
VC,
("
COO(O)
(0) (r +c=)),
+ 06
(0)0
(0)
Then, by (B.6 and B.7 one compute
(7 = sgn(a (a- w 03 (ign(w), sgg(())
Adamn (0)
V (0)
and
G
Ghe Adamm
(0)) sgn(g-1 91) (sgn(w) (w), (sgg())())
follows that the alignment score is
A (d'm (0), Ghe Adam (0) = sgn(ā 03 sgn(g-1 - 91) [ggn(g-1 - 91)
with prob.
sgn(gg1 91)
with prob.
which holds for any two-layer NN with width N.
Shampoo and SOAP. We proceed to consider Shampoo [66], which second-order optimizer with Kronecker product preconditioners. For the reader's convenience, we recall the Shampoo iterations for training neural networks, Following the notation Section W,, G. E R=xx be the weight matrix and gradient matrix for a layer at time step I, respectively. Shampoo generates left and right preconditioners:
-- L1-1 + G.G.
RG = R G Gr,
and then updates the weight matrix by
W = We nT, 1/4 1-11//GGr G.
with step size v) > O. we disable the accumulation the preconditioners and set L, = G,G.T and R, = G. G., then the Shampoo optimizer is simplified to
W++1 = W. 7) Shampoo(Gs), Shampoo(Gt) =
with Shampoo Gt) = U,V,", where Ug and Vú are from the reduced singular value decomposition of G. = U,2.VV V/²".
clear that if we apply Shampoo to (0) or (0) with the two-layer NN (.2 then under small initialization Assumption I at the initialization, the updates of the normalized weights 0 = (a, w) = 00 would be a 17 Shampoo (G be w))
1) Shampoo (G be (a. (e w)).
Here Gr or N are given in [B6 and B.7 Moreover, note that for any vector E Rd Shampoo(z) is simply x/||=r| Therefore, we can compute the (effective) initial Shampoo gradient directions for (0) and (0):
Shampoo = (Shampoo (G,(() Shampoo (G)()) sgn(a woi) (
02
[oB
a
W
and
Shampoo Gee (0) (Shampoo (G (8) Shampoo (Gh(())) (f))) sgn(g-1 91)
w
a
|w|
|al
It follows that the alignment score is, as N + 0o,
sgn(g-1 - 91)
with prob.
A Shampoo (0), The Shampoo (0)
Cshampoo
sgn(gg11 9i) with prob.
We finally consider SOAP. Following the notations in Appendix D if G2 is a vector, then G. [1, 0, .... T and A = G. We transform this gradient back and obtain G./[Ge] || means that at initialization.
Shampoo (G be (a, w)) SOAP or (a, w)) Therefore, the initial gradient conflict of SOAP follows from the case of Shampoo
28

---

C Inter-step gradient alignment of preconditioned gradient descent
Proposition 4 Let C(0) be a twice differentiable loss function with Hessian H()), and let P()) be positive definite preconditioner with uniformly bounded inverse p (0). Consider the preconditioned gradient descent update with exponent
and learning rate 1 >
B1t1 = 8g rpP~~(Pp)) V
(C.1)
For conseculive gradient vectors = V vl(,,) and g++1 vL(0++1), the alignment score
A(ge: (601))t+)++-- = list 11 satisfies:
2c+1
Algt: 9++1) =1 1/2 "HHH|HP-- "gell
0
(C.2)
|1
Furthermore, to ensure gradient alignment Alge, 9(+1) 2 1 e for some e > O, the learning rate must satisfy:
2e 2e
1 S
(C.3)
|HP "gell
This bound specializes to the following cases:
1. Vanilla Gradient Descent (P² I): Tmma
Xuna (H)
2. Newton's 's Method (P = H, 1): Tmux 22c
Proof. Let P(0) be a positive definite preconditioner with uniformly bounded inverse P-) For learning rate v) >> 0 and exponent 0 < < 1, the parameter update is:
0c+1 = - (0.) gt
(C.4)
where 9. = VL (6,).
We interested in the behavior of the alignment between successive gradients under this update. To study this, we expand g++1 VL(0++1) via second-order Taylor series around Or, using the Hessian (0,)) V2L(0,)
g++1 = 9: + H (0,) A0: + Dd D (0))0))A0, 20,] + + (140,11")
(C.5)
Substituting s
112 D (p P-*9.] + 1r
Gr+1 = 9r -HHH-*
(C.6)
Assuming 7) is small, we retain only the first-order term:
g+++ g: - - gt (12||
(C.7)
We compute the inner product between the current and next gradient:
g g++1 = ||g: 1| 2 n9.
HPPHHP- "g. + (1² p gi (1llll))
(C.8)
We also expand the norm of the new gradient:
||g++1|| || ||
2ng. HP 91 + n HHP* "9.l|" 0 (17²)
(C.9)
Taking the square root and using the approximation
vI
for small x, get:
9. HP
1 HHP* " gell
|1 = |19.11 [1- (1-
+0 (vi²)
(C.10)
2
19: 2
29

---

The cosine similarity (alignment score) between and G++1 is defined as:
A 9c+1)= = 9 96+1
(C.11)
| |
Substituting the approximations above:
1196 79.
HP 9t +0 (1))
A 9c+1) = 19:11 (1 (1-
(C.12)
9. HP- MP-*9 St
+ 42
[HP-
+ 0 +11))
|1 11
11 11
Factor |9::l2 in denominator:
92 HHP* 91 + (1²²
A (c+1))+))
(C.13)
1 1)
H HP
(1)))
11-11"
9ell
Using (1 - a a b)=1 =1 b O (a² + b²) for small a;
9i HPP 9t
9r HP 9t
n72
H 9t
A 1
1- 1- |19e|1
1 11
+0
11
2
19 11
(C.14)
Multiplying and simplifying to O (17²)
1/2 |HP 9t
A (gt. 9++1)
0 O17))
(C.15)
2
|1
Thus, the alignment between gradients remains close to 1 the second-order term is small. To ensure the gradients stay well-aligned, i.e., Algr, 9c+
1) 2 1 for given tolerance c 0, we require:
1/2 P ge S
(C.16)
2 Iseel
Solving for the maximum permissible learning rate 1) yields:
2c | l
5
(C.17)
P- Il
We now examine this upper bound in few special cases of interest:
1. Vanilla GD (I=I):
2e
Mmax
V
X2 X (H)
(C.18)
where Amax(()) is the largest eigenvalue of H.
2. Newton's Method P H, 1)):
HP-1 mmxx
(HP-1)- =
7max V2e
(C.19)
D Connection between SOAP and Newton's method
SOAP [ enhances Shampoo's efficiency by performing optimization in a transformed space aligned with the preconditioner's principal directions. For each layer's weight matrix and gradient G, E Rm SOAP maintains two covariance matrices using exponential moving averages:
L, 8 + -1- G.G
(D.1)
R, = R, G.
(D.2)
These matrices are then eigendecomposed as L; = Q. and R, = Q RARQT. where A, and An n contain the eigenvalues that capture the principal curvature directions of the loss landscape.
At each iteration I, SOAP updates each layer's weight matrix W, using its corresponding gradient G, as follows:
30

---

1. Project the weight and gradient into the eigenspace:
W. = Q G. Q.
2. Apply the Adam update in the rotated space:
W++1 W 1) (G)).
3. Transform back the original parameter space:
W1+1 = Qr Q
To reduce computational overhead, the preconditioners L, and RG are updated with frequency f in practice. We will analyze the impact of update frequency and momentum parameters through ablation studies Section
Before diving into the formal analysis, let us build some intuition for why SOAP is particularly effective at resolving gradient conflicts. The key insight comes from understanding how second-order information captures interactions between different loss terms. When gradients conflict, it typically indicates that improving one objective requires coordinated changes across multiple parameters information that is encoded in the Hessian matrix's ' off-diagonal elements.
SOAP approximates this second-order information in two complementary ways: (i) Its block- diagonal structure naturally captures parameter interactions within each network layer; (ii) Its adaptive preconditioner accumulates information about gradient correlations across training steps. This allows SOAP to implicitly identify and exploit parameter update directions that simultancously improve multiple objectives. Rather than simply following the average gradient, SOAP can utilize the local loss landscape geometry to find more direct paths to good solutions. The following sections make this intuition precise through formal analysis of SOAP's convergence properties and gradient alignment characteristics.
To establish SOAP's connection Newton's method, begin by examining how the Hessian matrix can be approximated in neural networks. We limit our analysis with networks trained with cross-entropy loss. Here the Gauss-Newton approximation takes the form:
Of 8 Of dw 0 Ow
HGN = EE
= = [99"]
(D.3)
where C denotes the loss function, f represents network outputs, and G aW. with vectorization g vec(G). Empirical evidence from [84] supports key simplifying assumption: Assumption 2. The Gauss-Newton component provides a good approximation to the true Hessian: Hax H.
8C is the gradient matrix
For our purpose, we begin by noting that there exists a one-to-one correspondence between the original parameter space and the rotated space that preserves the nectox multiplication.
Lemma 2. Let QL E Rmxm and QR E Raxn be two orthogonal matrices. For any matrix A E R and vector u E Rmm, define T = (Q1 & QR)U and Ā = (QL & QR)A(QQ & QT).
Then there holds
Ae = (Qu & QR) = Āö.
The proof follows directly by applying the transformation QL @ QR to Av and the definitions of À and U. Building on the above lemma, one can easily transform the preconditioned gradient descent in the original space to the rotated one and vice versa.
can establish the equivalence between the preconditioned gradient descent in the original and rotated spaces.
Corollary 1. Let W,, G, E RRxx be the weight matrix and gradient matrix for a layer iteration I, respectively, with vectorizations = vec(W.) and gt vec(Gr). The preconditioned gradient descent update:
Wi+1 Wt - 7p 9t.
(D.4)
equivalent to performing preconditioning in the rotated space:
W w. nP-1z gt
(D.5)
where E Rmmaxmn is the preconditioner, and w. g. and P are the rotated weight, gradient, and preconditioner defined by the transformations in Lemma
31

---

Proposition 5. Let L, E [ and R. - E have eigendecompooittons L, QLArQ and R₁ = Q Under the assumption of Lemma 4 the equivalent preconditioner in the rotated space diagonal, i.e. HaN diag
Proof. The proof follows from the combination of Lemma 2 and Lemma 4 using the Kronecker approximation:
First, we express HGN
HON - 1/22 & / / (E [GG"])
(D.6)
Then, derive the rotated preconditioner:
HGN (Qz @ QR) HON (0f @ Q1i
(Qr @ Qn) (L. & R (0% & ( / Tr
(QI.) (Qn / (B
=A @ A 1/2 R (B(A)))
The final expression shows that HaN is diagonal, as it is the Kronecker product of diagonal matrices scaled by a scalar factor.
Finally, we connect our analysis Adam's update rule by adapting the following result from Molybog et al [85]:
Proposition 6 (Adapt from 85
Suppose that 0r is a local minimum and assume that 0 8" ~
(1 IT). For Adam update rule denoted by 0.+1 = 02 17 Ad)m(gr), we have
Adam(g)) 2 diag (H) gr.
(D.7)
Proof. The Adam optimizer follows the update rule:
m = 3, + (1 - Bi)gt,
1/2 82 + (1 -2) @ 9t.
m m1/(1-88) -1-1))
ût /t/(1 (1-35),
0. At 1 - --mii/( (v/ + c).
Taking a first-order Taylor expansion of the gradient around a local minimum 8":
go go- + He- (0 8*) H (0 8*).
This yields
gage ~ Hg- (0 - 0*) (0 8T) H
Under our assumption that 0 - 0" ~ N1000 0²1),
E [oogo] H&lE (0 (0- 0r) (0 0-]_| HJ. o² Ho Hj.
By construction, v approximates the diagonal of Egwer (geehl) }. where 0, represents the distribution of model weights over the past O(1/(1 B2)) steps:
diag (Ea-o, [gogh] o² ddagg (H
Finally, assuming m, x gr:
Adam (gt) R
m
2
my
diag( 1 gr.
+
32

---

Theorem 1. Under assumption of Proposition 6 SOAP's update approximates Newton's 's method:
= Wt - -1 Soap(g.) s w, - 1
(D.8)
Proof. Combining Propositions Band 6 we obtain
Ad, diag(H) H))
" ge diag(HoN) "g HaN se H --gh
't.
(D.9)
By Corollary I this is equivalent to the Newton update original space:
W++1 = wi
H H-'g. ge-
(D.10)
As direction implication, the Hessian matrix approximately diagonal in the rotated space.
The key insight that SOAP effectively approximates the block-diagonal Gaussian Newton com- ponent of Hessian in rotated space, with each block corresponding to layer- wise Kronecker factorization. This structure naturally promotes gradient alignment across optimization steps, as we demonstrated theoretically in Proposition and observed empirically in Figure [5]
Remark 1. While SOAP effectively approximates Newton's method through its block-diagona struc- ture, other optimizers make different compromises in their approximations. Adam can approximate Newton's 's method, but requires the highly restrictive assumption that the Hessian matrix is diago- nal. Similarly, Shampoo takes different approach by using the square root of the Gauss-Newton component preconditioner 1861
Shampoo(ge) = H 1/2 = H-1 9t-
(D.11)
These structural differences help explain why SOAP achieves better gradient alignment than both Adam and Shampoo.
E
Connection of Shampoo and Muon to Quasi-second-order Methods
We review the connections between Shampoo, Muon, and quasi-second-order optimization methods, building upon results from 87 and 86 This section aims to provide a self-contained exposition of these relationships.
Adagrad. Adagrad is preconditioned online learning algorithm that leverages the accumulated covariance of gradients as a preconditioner. Let E Rº denote the parameters at time and Sr R
denote the corresponding gradient. Adagrad maintains a preconditioner HAds Σt= 9:9. The parameter update with learning rate 1) is given by:
1/2 0r+1 = 0r HH Ada gr
(E.1)
Shampoo. Shampoo tracks two statistical matrices throughout training, La E Rmxm and R. Raxx defined as:
* =I= Σ G.G.; R. = cla G.
(E.2)
where G3 E Rmxx is the gradient matrix at step 8, and stability.
0 is small constant for numerical
The full matrix Adagrad preconditioner H, can be approximated as (L, @ R.) This approximation transforms the Adagrad update rule w4+1 w: ,HH, 1/2 gs into the Shampoo update rule for parameter matrix W:
W W, nL, 1/4 G.R," 1/4
(E.3)
The theoretical foundation for this approximation is provided by the following lemma:
33

---

Lemma 3 (177). Lemma 1). Let G1...., Gt m be matrices of rank at most r. Let 9s vec (G,) and define Hi Elmn E._ 9ag. Define Lt, Rs above: Lt =elm + Σ G,G.g R₁ eIn E'= GTG. G.. Then for any p, > 0 such that 1/p + 1/g 1, we have H, S
rL, R
1/p
follows from this lemma that for any p.q q with 1/p- 1, the full Adagrad preconditioned 1/2)), The 1/2 11 can be approximated by (L,") @ R R//9)-1/2g gt vec(L, 1/2p R 1/2r).
gradient H
Ada
case where p q 2 yields the standard Shampoo update.
Moreover, 86 explored the Hessian approximation perspective of Shampoo, demonstrating that the preconditioner in Shampoo is Kronecker product approximation of the Causs-Newtor component of the layerwise Hessian (see Remark
Muon. The Muon optimizer 662 was recently proposed for optimizing neural network weights representable as matrices. At iteration t. given current weight Wi-1. momentum μ, learning rate T).- and objective Lt, the update rule of the Muon optimizer is:
M, ,M, V&, (W,-1),
(E.4)
O, = Newton-Schulz (Mr)
(E.5)
W. W,-1 0
(E.6)
Here, M, is the momentum of gradient at iteration t, initialized as zero matrix when t 0. The Newton-Schull iteration process is adopted to approximately compute (M,,))-
When preconditioner accumulation removed, can observe that the update simplifies
62
W+1 = W, - (GG,G.)))AAG(G G,)
(E.7)
W, "UUSSTT)""UUSS""VVVVV""
(E.8)
W. ((US 1(2-1
(E.9)
= W, - TT--UUS-VSSS-133
(E.10)
W. nUV
(E.11)
From this derivation, Muon be viewed as approximating the use of H 1/2 with additional orthogonalization benefits.
Ada
preconditioner,
Additional Lemma and Proof
Lemma (186), Corollary 2). Under the assumption that the reshaping of the Hessian tensor HGN rank-1,
Hon (E [GG]] & E [G Tr ())
F11
Proof of Proposition
Proof. When 2, we note
A(V1, 12) =
122
1
C1
1/
+ 2
|1
|12
||0
2 + cos(U1, U2) + 1) - = cos(v11 U2)
F.2 Proof of Proposition 2
Proof. (i) Single-step drop. By L-smoothness (descent lemma),
L(0r+1) S L(O,) + (ge, Oc+1 (,) + $ 110.41 0,112.
34

---

With 0141 - 0, = -n
A. = ((0.) - [(0++1) 2 n(g,,
La2
2 11/e11??
Since h4 Pige and P. = 21 I, we have
(ge. ha) h1 Ph- 2 ir lhelll??
hence
A, 2
4
2 |1bell?.
Under the equal-norm assumption |||||| X, let u/ h2/A. By the chosen definition of Ajutra
n² (Aimtra + 1).
i=1
Since h, M and the refined bound.
h = > ut and >; |||1| n2? obtain the displayed equal- formulas
(ii) Two-step cumulative drop. Apply L-smoothness with 0 = =0--1, Ø = =0++1:
[(6(41) - [()--1) + 101-1.041 - 61-1) + 1/2 [|0 - 11-1112.
Since 0.41 - 01-1 = -n + h4), we get
A D. 7/0t-1, h-- h1)
| 1111?. Ly 2
For the linear term, using Pi-1 M 1,
(9 h4-1) h Pille-
M a1lhe-1l? 1l?
and
h1) - 2 1o (t)-
Let : 1|111111, := |hel]. and Aiuter = (
Then
ab
n/ h--1 + h1) 2 31 (a² + ab Altter).
For the quadratic term, ||h,-1 + h4 12 = a² + 2ab Aintor + b². Putting the bounds together,
A A. z + (a² + ab Altter) La (a² + 2ab A inter + b?),
which rearranges to the stated inequality. The coefficient of Ainter equals M (1 - LyM) 2 0 for
n S 1/LL hence monotonicity.
G Experimental Details
G.1 Architectures
This section outlines the network architectures employed in our work, along with the enhancements introduced to improve their performance.
Modified MLP. The modified MLP architecture proposed by 21 which has been extensively used in the literature [35 73 88 [4 88 90] due to its improved capability in learning complex PDE solutions. The network processes input coordinates through two parallel encoders:
U (Wix + b1), = (W2x + b2).
(G.I)
Then, for / - 1,2, L, the forward pass is defined as:
p(((x)) = w g (1-1)) + (()
(G.2)
( (x) (r." @ U + (1 - (1.) (x))))) @V
(G.3)
35

---

The final network output is given by
fx(x) = ((++)) (x) + b(4+1)
(G.4)
where a is nonlinear activation function, @ denotes element-wise multiplication, and the trainable parameters are:
0 = W1. b1, W2, b2, (WO, b L+1
(G.5)
This architecture extends the standard MLP by incorporating dual input encoders and merging their features through point-wise multiplication at each hidden layer. While computationally more demanding, this modification demonstrates superior performance minimizing PDE residuals compared to standard MLPs.
PirateNet. PirateNet is proposed by (27], which aims to enable stable and efficient training deep PINN models. The architecture first transforms input coordinates into nigh-dimensional feature space using random Fourier features [9I]
coos Bx) sin [xx)
((x)
where B RRd has entries sampled i.i.d. from (O,,N(0, s²) with user-specified & 0. This embedding mitigates spectral bias in PINNs by improving the eigenfunction frequency of the Neural Tangent Kernel, enabling better learning of high-frequenge components and multiscale features (30..
The embedded coordinates are processed through two dense layers that act as gates:
U ( +(x) + b1), V (W2((x) + b2);
where is woint-wise activation function. This gating mechanism is essentially the same in modified MLP.
Let (!) = (xx and (4) be the input to the I-th block (1 < 535 L). Each block performs:
p() a (W ( ("),
(G.6)
(0) p() о UU (1- (1- e V,
21
(G.7)
g ( a (w 21 +b²)) (1)
(G.8)
22 = g @U -1- (1-g")) o
(G.9)
h a (W (1) (1)
22 + +b(1)
(G.10)
x ((+1) = (h) + (1
(G.11)
Each block comprises three dense layers with dual gating operations and an adaptive residual connection. The trainable (1)) parameters control block nonlinearity: (!)) 0 yields an identity mapping, while (!)) produces fully nonlinear transformation.
The final output of a PirateNet of I. residual blocks is given by
Ug = (1++1)
(G.12)
Importantly, we initialize ( = O, 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 a values. Addi- tionally, the linear structure at initialization enables direct integration of prior solution data through least squares fitting:
min |W -
(G.13)
where Y represents available measurements. This approach provides optimal initial guess based on various data sources, including experimental measurements, boundary conditions, or linearized PDE solutions.
36

---

Exact imposition of periodic boundary conditions. We adopt the approach of [70] to enforce periodic boundary conditions as hard constraints, improving both training convergence and accuracy. Consider one-dimensional periodic function with period P satisfying:
(1) (a) = w (a P);
= 0,
(G.14)
We construct Fourier feature embedding:
v(z) (cos(wx)) sin(wr)),
(G.15)
where 2x, Any network uo (v(x)) using this embedding inherently satisfies the periodic boundary condition.
The same idea can directly extended to higher-dimennional domains. For two-dimensional domains, the periodicity constraints are:
0 Pa.3), E 16,6+ Pyl,
dx ().
(a
(G.16)
dr
(x., a)
Oy u (r,bb
E (a, + Pa], Pa),
(G.17)
oy
for - 1, where P2 and Py are the periods in the x and y directions, respectivlly. Similarly, these constraints are encoded using the embedding:
v(z v) [cos (wax) sin (Wzx), cos (wyj)) sin (wyy))
(G.18)
with Uig- =
2x p. p.
2x
For time-dependent problems, we concatenate time coordinates t with spatial embeddings: uo([t, v(r)]) or uo((t, v(x.
Random weight factorization. We implement random weight factorization (RWF) 66 to enhance PINN performance. RWF decomposes each neuron's weight vector as:
w (R.1) = (A.1). (R.1).
(G.19)
where k = 1,....... d. 1, 1,..., L + (((.)) IRch- is the k-th row of weight matrix W(), g(k.l) E IR is a trainable scale factor, and ((e()) E IRdi-1, This factorization can be expressed in matrix form as:
w diag (sW) v(O, 1= 111,2 1.,, L+1
(G.20)
with (0) E Idd.
Implementation involves: (1) initializing MLP parameters using the Glorot scheme (2) initializing scale vectors cxp(s) where N(,, aI), (3) factorizing each weight matrix as W diag(exp(s)) V, and (4) optimizing parameters s, V directly. We employ exponential parameterization following Weight Normalization (93 to ensure non-zero scale factors across varied magnitudes. We recommend μ 0.5 1 and a 0.1, as these values consistently improve convergence and accuracy while avoiding the instability of larger values or the diminished effect of smaller ones.
G.2 Training pipeline
This section details the methodologies and strategies used to train PINN models.
Causal training. Recent work by [73] shows that PINNs may violate temporal causality when solving time-dependent PDEs, as they tend to minimize residuals at later times before correctly
solving earlier times. To address this, introduce causality- -ware training approach. We partition the temporal domain into M equal segments and denote the PDE residual loss within the i-th segment as L,. The modified residual loss becomes:
M
(+) = M
Σ i=1
wi Ct (0)-
(G.21)
37

---

We compute the temporal weights as
i-1 1
W/ = exp
=xp
-E
for i = 2.3....,
(G.22)
h=1
Then,
1 M (, M
Σ exp
c Σ A=1
c)) (0).
(G.23)
The weight wy decreases exponentially with the cumulative residual loss from previous time steps. This ensures that (6) is minimized only after previous residuals (0) become sufficiently small, enforcing temporal causality in the optimization process.
The causality parameter requires careful tuning: small values may insufficiently enforce causality, while large values can create optimization difficulties by requiring extremely small early-time residuals before later times are considered. We recommend selecting a moderate that allows all temporal weights to converge to by training completion, reducing it necessary.
Learning rate annealing, Another key challenge in training PINNs is balancing loss components, as they often exhibit multi-scale behaviors, resulting in unbalanced gradients and training failures.
We implement self-adaptive learning rate annealing algorithm 21 that automatically balances the weighted loss:
(())= = >ic (0) Abc (0) C,),
(G.24)
The global weights are dynamically computed to equalize the gradient norms of each loss component:
dic [V) + [Vo.s. + (0)]
(G.25)
(0)]
Aac Vo (0) + | Lbe (1V4++
[Vo.,
(G.26)
X, (oLic + VoL + (0)
(G.27)
(0)]
where || denotes the 2 norm. Then we obtain
|| VoLin( = |e VoLia (0)l = || Vo) = [ IV (,
(G.28)
This formulation equalizes the gradient norms of weighted losses, preventing bias toward any particular term during training. The weights are updated as running averages of their previous values, stabilizing stochastic gradient descent. These updates occur at user-specified intervals (typically every 100-1000 iterations), incurring minimal computational overhead.
Curriculum training and time-marching. Despite the improvements described above, PINNs still face challenges in complex domains requiring high accuracy, such chaotic systems like high Reynolds number Navier-Stokes equations where error accumulation can cause trajectory divergence. We address these challenges using curriculum training 422 which decomposes the optimization into more manageable sub-tasks.
The core idea involves decomposing the entire optimization task for PINNs into a sequence of more manageable sub-tasks. In this work, we mainly focus integrating this strategy into our training pipeline for solving time denendepndeent PDEs and singular perturbation problems.
An effective approach employ is the curriculum training strategy introduced by 422
For time-dependent PDEs, we implement temporal domain decomposition: the time domain is divided into smaller intervals. After the first window, initial conditions for subsequent windows are set using predictions from the final step of the previous window. This approach reduces the difficulty of the optimization task of learning full system dynamics, though at an increased computational cost due to wer-wwndow model retraining.
38

---

Table Parameter settings and numerical configurations for generating the reference solution across PDE benchmarks.
PDE
Parameter
Package
Resolution
Wave
4
N/A
200 128
Burgers
0.01
Chebfun
200 512
AC
10 100$,00 1,0+,
Chebfun
200 x 512
Kdv
1,y= 0.022
Chebfun
200 512
KS
a 100/16,9 100/1007
Chebfun
250 x 512
GS
0.2, 0.1,b
40,
100,
cı 1,000
Chebfun
100 200 200
M. 10,7 15k
Chebfun
LDC
Re=5x109
IncompeessibleNNvvirStoee.
128 x 128
KF
Re=10²
IncompressiheeNavvieSSukee
50 x 512 x 512
Ra-10", Pr 0.71
IncompeesibbeNNvvirsStoess
x 100 x 200
While we also partition the temporal domain to compute causal weights within each window, this differs from the ime-marching strategy. Both techniques promote learning solutions sequentially along the time axis to respect causality, but causal weighting complements rather than replaces time-marching, as causality violations may still occur within individual time windows.
G.3 Data Generation
We generate our reference dataset using two numerical packages: Chebfun [94] in MATLAB and IncompresibleNNwiesStoee [95] in Julia. The data generation process employs time step of dt 10-1, followed by temporal downsampling to construct the final dataset. Table summarizes the PDE parameters and dataset details.
G.4 Hyper-parameters
Unless otherwise specified, we adopt the following hyperparameter configuration for our experiments.
Architecture. We employ PirateNet [27 as our backbone architecture, configured with three residual blocks (9 layers in total), hidden dimension of 256, and Tanh activation functions. All weight matrices are initialized using random weight factorization (RWF) [69] with parameters
1.0 and 0.1. When applicable, strictly enforce exact periodic boundary conditions following 770.
Training Protocol. We train our models using mini-batch gradient descent with 8,192 randomly sampled collocation points per iteration. The learning rate schedule comprises an initial linear warm-up phase over the first 5,000 steps (from 0 to 10 10- followed by exponential decay with factor of 0.9. To enhance training stability and convergence, we implement learning rate annealing for loss balancing 21 72], updating the loss weights every 1,000 iterations with moving average. For time-dependent PDEs, we apply causal training [73] 72) with a causal tolerance of 1.0. Additionally, we leverage time-marehing and curriculum learning strategies 42 for particularly challenging benchmarks.
Optimizers. We evaluate several optimizers with the following configurations:
Adam: We use the Adam optimizer 7 with standard hyperparameters B1 = 0.9 and 3r = 0.999, which has become the de facto standard for training PINNs due to its robust performance and computational efficiency.
SOAP [ Based on our ablation studies presented in Figure 4 select 31 = 0.99 and
3g 0.999 for SOAP, which yield optimal performance across our experimental tasks.
Muon 622 : For Muon, we adopt momentum hyperparameters matching those of Adam:
31 0.9 and B3 0.999. Notably, we employ more accurate Newton-Schull coefficients (2, 15.-1,,0..3 0.5) and implement 10 Newton-Schulz matrix iterations to ensure convergence of the orthogonalization procedure.
39

---

Table 6: Hyperparameter configurations for benchmark PDEs. Hyperparameter settings used to reproduce our experimental results. The backbone architecture is PirateNet, where Depth indicates the number of adaptive residual blocks, and Width denotes the number of neurons per hidden layer. RFF and RWF represent Random Fourier Features and Random Weight Factorization, respectively.
Parameter
Wave
Burgers
AC
Kdv
KS
GS
GL
LDC
KF
RT
Architecture
Depth
Width
256
256
256
256
256
256
256
256
384
384
Activation
Tanh
Tanh
Tanh
Tanh
Tanh
Swish
Swish
Tanh
Tanh
Tanh
RFF scale
10.0
2.0
2.0
2.0
2.0
2.0
2.0
10.0
2.0
2.0
RWF
==1 0, 0=0011
Learning rate schedule
Initial learning
10-2
Decay
0,9
Decay steps
10a
2 10a
5 10a
x 109
10
10a
2 10a
x 10"
x 10
2 10
Warmup steps
10"
Training
Iters (per time window)
10
10a
3x 10
10
102
10a
10P
2x 10
2x 10'
10a
Batch size
8,192
# Time windows
10
N/A
25
Weighting Scheme
Grad Norm
Causal weighting
Tolerance
1.0
1.0
1.0
N/A N/A
1.0
1.0
# Chunks
16
Kron [63]: For Kronecker-factored optimization, we use the same momentum hyperparam- eters Adam: 31 0.9 and 33 0.999.
The complete set of hyperparameters is detailed in Table 6 largely following the configurations established in [72] and [27]. The decay step is tailored for each benchmark to ensure the learning rate reaches a sufficiently small value (approximately 10-7) by the end of training. The number of time windows is determined empirically based problem complexity, with fine-tuning guided by the loss convergence behavior observed during preliminary experiments.
G.5 Computational Cost
Our implementation is based on JAX-PI 72 and we conducted all experiments on single NVIDIA A6000 GPU, with detailed runtime benchmarks reported Table 7
G.6
Benchmarks
Wave equation. taking the form
We consider one-dimensional wave equation in the domain Ω [0. 1] x [0,
un((,, 1) 4uza (v. 0, (x, (0, (0,1),
u(0, f) u(1,t) 00 LE [0, 1].
1 u(2, 0) sin(Tx) sin(4 x E [0.1].
ue(z, 0) 0, x [0, 1].
where represents the wave amplitude, and is the wave propagation speed, determined by the medium's physical properties.
By d' Alembert's formula, the solution u(z, t) is given by
u(x, t) = sin(TT) cos((2t) + sin(AAz) cos (8T)).
40

---

Table 7: Computational runtime (in hours) comparison of different methods across various PDEs. All experiments are performed on an Nvidia A6000 GPU, reporting the total training time needed to achieve convergence using PINNs with Adam and SOAP, respectively
Benchmark
Adam
SOAP
Wave
2.80
4.35
Burgers
1.18
4.05
Allen-Cahn
1.48
5.83
Korteweg-DD De Vries
1.61
3.90
Suramoto-Sivashinsky
19.51
34.16
Grey-Scott
19.52
40.01
Ginzburg-Landau
15.98
23.75
Lid-driven cavity (Re x 102)
5.67
8.25
Kolmogorov flow (Re 10a)
9.56
11.00
Rayleigh-Tayle Taylor instability (Pr 0.71, Ra 10%)
20.23
21.73
Reference
Predicted
Absolute Error
x10
0.5
0.0
IC Loss
PDE Loss
Rel. L Frmo
10
Adam
SOAP
10
Wwd www. www.mmmmmmm mmm.o
10
10
11
www.mmp
0.0
0.5 Step
1.0
0.0
0.5
1.0
10
0.0
0.5
1.0
x10*
Step
x10
Step
X10
Figure 6: Wave equation. Top: Comparison between the reference solution and the model predictions. Bottom:
Training loss and test error trajectories for the Adam and SOAP optimizers.
Burgers equation.
The ID Burgers equation is defined as:
+ utl 1/2lga
where represents the velocity field, and the kinematic viscosity coefficient controlling the diffusion strength. Here we set (a, t) E l = 11 1] x [0, 1]. with initial and boundary conditions:
"(x, sin sin)),))
((-1,1) ((, = 0.
and viscosity parameter 0.01/x.
Allen- Allen-Cahn equation. boundary conditions:
We investigate the one-dimensional Allen-Cahn equation with periodic
dt 0.00010ez + 5u 5u =0, te [0, 1], x [-1,11
u(0, x) x² coo((TT),
u(t,-1) -1) u(t, 1), (t, -t. = U2 (t,
41

---

Reference
Predicted
Absolute Error
x10
1.00
0.30
0.25
TC Loss
PDE Loss
Rel. L2 Frror
10
10"
Adam
SOAP
10"
"WW
10-
0.0
0.5 Step
1.0
10
0.0
0.5
1.0
0.0
0.5
1.0
x10
Step
x10
Step
x10'
Figure 7: Burgers' equation. Top: Comparison between the reference solution and model predictions. Bottom: Training loss and test trajectories for the Adam and SOAP optimizers.
Reference
Predicted
Absolute Error
x10
1.0
0.5
IC Losss
PDE Loss
Rel. L Rrmr
10
SOAP
10.2
10
10-
10
Step
x10
Step
x10
Step
X10
Figure 8: Allen-Cale equation. Top: Comparison between the reference solution and model predictions. Bottom: Training loss and test error trajectories for the Adam and SOAP optimizers.
where u represents the order parameter (e.g., concentration difference between two phases), controls the interfacial width, is the reaction rate coefficient, and the term (u w²) drives the phase separation.
It is worth noting that this benehmark has been extensively used to validate the effectiveness of PINNs methodologies. In Ta we compare the test errors across different PINNs advancements, demonstrating that our approach achieves state-of-the-art performance with improvement of up one order of magnitude accuracy.
Korteweg-D Vries equation.
The one-dimensional KdV equation is expressed as follows:
+ nuu, + yLurer= =0, t (0, x E (-1,1)
u(x; 0) COS(TTT),
u(t, -1) u(t, 1),
where represents the wave amplitude or water surface elevation, and governs the strength of the nonlinearity, while μ controls the dispersion level. Under the KdV dynamics, this initial wave evolves into a series of solitary-type waves.
42

---

Table Allen-Cahn equation: Relative L test errors obtained by different PINNs variants.
Method
Relative L error
Original formulation of Raissi al. 660
4.98 10
Adaptive time sampling 41
2.33 X 10
Self-attention 96
2.10 X 10
Time marching 97
1.68 X 10
Causal training
1.39 X 10
Dirac delta function causal training 98
6.29 X 10-6
JAX-PI 72
5.37 X 10
RBA-PINNs 88
4.55 X 10
PirateNet 27
2.24 x 10-6
BRDR-PINNS
1.45 x 10-6
Ours
3.48 x 10 -
Reference
Predicted
Absolute Error
x10
1.00
0.50
0.25
IC Loss
PDE Loss
Rel. L2 Error
10
10"
Moo W Myd
Adam SOAP
10
10
Wmm wh
0.0
0.5
1.0
1.0
0.0
0.5
1.0
Step
x10"
Step
x10
Step
X10'
Figure 9: Korteweg-De Vries equation. Top: Comparison between the reference solution and model predictions. Bottom: Training loss and test error trajectories for the Adam and SOAP optimizers.
For our study, we adopt the classical parameters of the KdV equation, setting r) = and μ = 0.022
999
Kuramoto-Sivashins. equation.
The one-dimensional equation takes the form:
Wt + auu + Buge yuazzz O, E [0, [0, x e [0, 2m),
u(0, x) = up(x))
where represents the height of thin film flame front. This equation arises in various physical contexts, including flame front propagation, thin film flows, and plasma instabilities.
In this example, we take T = 0.8, 100/16, B 100/162, 7 100/16* and wo(x)
c1 + si)).
Grey-Scott equation.
The system is described by the following coupled PDEs:
e Au + (1 u) cuue, EE (0, 2), (I, E (-1,1)² Ar by + czuv², t E (0.2). (a.v) (-1,1)²
43

---

L
10.
10
Adam SOAP
Initital Coadition
10
PDE Loss
10"
Man
0.0
0.2
0.4
1.0
Step
x10*
Figure 10: Kurammto-SSvvasiisky equation. Training loss and test error trajectories for the Adam and SOAP optimizers.
tw0.4
2
075
Rel.
0.55
Prest
Als Earaa
Figure 11: Grey-Scott equation. Comparison between reference solution and model predictions.
With periodic boundary conditions, the initial conditions are:
uo( y) = 1 exp( 10p 0.05) + (y + 0.02))) up (x, y) = 1ep(-10((x=0.00)))) + (y - 0.02)3))
where u and v represent activator and inhibitor concentrations respectively, E1 and F2 are diffusion coefficients, and (b1, bz. C1, 02) control reaction kinetics. This system generates diverse spatial patterns including spots and stripes.
We set parameters characteristic pattern formations. Due to the similar behavior of and v, we report only the relative L error of in Table 3
0.2, 0.1, b 40, by 100, and c = 000, which generates
44

---

Rei e Bire
10
Adie SOAP
10
Rei 1 Ens
10
10
10
110
10.4
0.6
03
Step
Xis'
Figure 12: Grey-Scott equation. Test error trajectories for the Adam and SOAP optimizers.
10"
Inttal Conditoo Lone
10
10"
16-1
107
Insial Cenalioe Loo
107
10"
10
PDE Lewa
10"
Vill
10"
PDE
10
Step
Figure 13: Grey-Scotl equation. Training loss and test error trajectories for the Adam and SOAP optimizers.
Ginzburg-Landan equation.
The complex Ginzburg-Land Landau equation in 2D takes the form cAA + μA 7A/4/2, e (0,1), (x, v) E (-1,1)²,
0A
Öt
with periodic boundary conditions, an initial condition
Ao( y) (10y + 10kr) exp (-0.0122500.² 2500y
where A is the complex amplitude representing the envelope of oscillations, represents the diffusion coefficient, is the linear growth rate, and controls the nonlinear saturation. For this example, we set 0.004, μ 10 and 7 = 10 15i.
By denoting A = w + ic, we can decompose the equation into real and imaginary components, resulting in the following system of PDEs.
Ou
eAu uuu (u (w5 + oz)),
Øt
0v
eAv u(u (u (u v²))
Øt
Given the coupled dynamics of and present the relative L error of in Table
45

---

t= 0.2
t=0
0=0.8
1=10
Ref.
11.
Phott U
60
-05
Als. Frior
Figure 14: Ginzburg- Landau equation. Comparison between the reference solution and model predictions.
10
10"
Figure 15: Ginzburg-Lane Landau equation. Test error trajectories for the Adam and SOAP optimizers.
90J2
16
10"
Soutted Mornddahh Hora
107
10
10
10"
Mrn Jaann
16*
10"
18a
PDE Loos
10
10**
107
Step
Figure 16: Gindburg-Landau Landau equation. Training loss trajectories for the Adam and SOAP optimizers.
46

---

Boundary Condition
Adam
Adaz
SUAP
Nawderh A MMM
10"
Boundary
10"
Leos
Adam
Adaz
WWM AAM AMAM
M
10-
Mapee Ma MMMM
10
PDE Momentum Lova V
10
PDE Continuity Loss
Adam
Adam SOAP
10
10
May
10
0.75
0.75
Step
X10'
Step
X10
Figure 17: Lid-driven Caviry. Training loss and test error trajectories for the Adam and SOAP optimizers.
Lid-driven Cavity. We study the incompressible Navier-Stokes equations in non-dimensional form for a two-dimensional domain:
u Vu Vp
Au = 0, (z,y) E (0,
Re
uu= 0, (z. E (0,
where u = (u, v) represents the steady-state velocity field, is the pressure field, and Re is 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 constant velocity while the other walls are stationary, leading the formation of characteristic vortical structures whose complexity increases with the Reynolds number.
To ensure continuity the corner boundaries, implement smoothed top-lid boundary condition:
u( c,, y) cosh (Co(z - 0.5))
v(2, v) 0,
(G.29)
cosh (0.5Co)
where æ [0,1].y 1, 50. For the other three walls, we enforce no-slip boundary condition. Our goal is to obtain the velocity and pressure field corresponding to Reynolds number of 5, 000. The accuracy of our method is evaluated using the velocity magnitude vu² + vZ, with results presented in Table
Kolmogorov flow. We study the two-dimensional Kolmogorov flow governed by the incompressible Navier-Stokes equations:
u. +u Vu -p+
Re Au +
v. 0,
on the unit square domain (x, E [0, 1)2.
Here = (u, represents the time-varying velocity field, and denotes the external forcing term that maintains the flow structure. The system evolves from random initial state and develops
47

---

320
Figure 18: Kolmogorov flow. Comparison between reference solution and model predictions.
Rel. L Error a
10
SOAP
10
Rel. Enory
10
10
10"
Step
X10
Figure 19: Kolmogorov flow. Test error trajectories for the Adam and SOAP optimizers.
characteristic flow patterns, where energy transfers between different spatial scales through nonlinear interactions and viscous dissipation.
For our study, the system is driven by sinusoidal forcing f = sin (Arry),)) 0). The numerical experiment initializes with a random initial condition and evolves until T 2. The model's performance is quantified by the relative 1,2 error of vorticity (Table
Cayleigh-Taylon instability. We investigate coupled flow-temperature system that models buoyancy-driven instability rectangular domain (x.y) E [0, 1] 2]:
Pr
u. + u. Vu = Vp V Ra
Au Tey,
(G.30)
V u
(G.31)
T uT
Tu
(G.32)
PPrRa
where T' the temperature field (acting as a density proxy through the Boussinesq approximation), Pr is the Prandtl number (ratio of momentum to thermal diffusivity). and Ra is the Rayleigh number (measuring buoyancy-driven flow strength). This system captures the characteristic mushroom-shaped plumes that develop the heavier fluid penetrates into the lighter fluid below.
48

---

Inttal Condition Leek
SOAP
10
Inital Cundition
10
10"
PDE Momesturn Loss
10
10"1
10
PDE Momentum Loss V
10
10"
Continiiy Loss
10
Step
X10
Figure 20: Kolmogorov flow. Training loss trajectories for the Adam and SOAP optimizers.
We set the Prandtl number Pr = 0.71 and Rayleigh number Ra 10 The boundary conditions periodic in the horizontal direction for both u and T. with Dirichlet conditions u = T = 0 imposed on the top and bottom boundaries. The accuracy of our method is evaluated using the temperature field, with results presented in Table
49

---

C=2
t=4
C-
t=
t=2
t=4
t-6
1-8
10
05
Ref. u
Nef.
-.0
10
Pred
Phod
-0.
1.0
Abr. Emnor
Abs.
Figure 21: Soyleigh-Taylon instabiliry. Comparison between reference solution and model predictions.
Rel. L Error u
Adam SOAP
10
mmmm
10
Rel. L2 Error
10
10
10
m
10
Rel. L Error T
10
10
10
0.5
2.0
Step
x10
Figure 22: Raaleigh-Tayyy Taylor instability. Test error trajectories for the Adam and SOAP optimizers.
50

---

Fnriee Corditton
Iritial Condition Less
Adam:
10"7
10"
Muyl
W
Mi
Mali Manda Wymm
1677
10. 11
16.
Boundary Conditoon
Boundary Condition
10a
10"
MAd
My
Malled Aumenal
10
10"
Whimmae
10
10.
PDC Momassttan Loss
PDE Mortennam Loss
Migh
107
Will
1674.
M
107
10"
PDE Cortinitty Loss
PDE Eaergy Los
MA
1071
10
MM
10a
0.0
1.5
2.0
10"
0.0
1.
210
SteP
X10
Step
x1n
Figure 23: Rayleigh-Toone Taylor instability. Training loss trajectories for the Adam and SOAP optimizers.
51


Reviews:
[{"review_id": "nAc5Tx0Sg9", "reviewer": "Reviewer_wvKq", "summary": "This paper investigates a key optimization challenge in PINNs, where conflicting gradients from different loss terms hinder effective training. The authors propose a novel gradient alignment score, distinguishing between intra-step (within a step) and inter-step (across steps) misalignments. They reveal that first-order optimizers like Adam suffer from poor alignment, especially during early training.\n\nTo address this, they propose a quasi-second-order optimizers, especially their proposed method SOAP, which leverages preconditioning to mitigate gradient conflicts. Extensive experiments across 10 PDE benchmarks and achieving better accuracy and faster convergence.", "questions": "None.", "limitations": "None.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 4, "originality": 4, "strengths_and_weaknesses": "**Strengths**\n1. The idea of formalizing and quantifying gradient conflicts via alignment scores in PINNs is novel. \n2. The paper provides solid theoretical definitions (alignment score), along with analysis on why quasi-second-order methods mitigate conflicts. The proposed SOAP optimizer is well-motivated and rigorously tested.\n3..Writing is very clear.\n\n**Weakness**\n1. The theoretical analysis focuses mainly on 1D Laplace equations at initialization.\n2. What's the results of Adam+ L-BFGS as it is a very useful technique in training PINNs.\n3. The sensitivity of performance to learning rates in SOAP vs. Adam is not deeply discussed.", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer's thoughtful and constructive feedback. Below, we address each of the raised concerns in detail, incorporating additional experiments and clarifications to strengthen our submission.\n\n> ### W1. The theoretical analysis focuses mainly on 1D Laplace equations at initialization.\n\nWe agree that extending the analysis to higher-dimensional or nonlinear PDEs would strengthen the theoretical contribution. However, the 1D Laplace equation provides the most tractable setting for establishing rigorous analytical results, allowing us to clearly demonstrate that gradient conflicts are intrinsic to PINNs at initialization. Importantly, our key theoretical insights are consistently validated through extensive experiments across a broad range of benchmarks (Figure 2 and Figure 5 of the manuscript), supporting their generality beyond the 1D setting.\n\n\n> ###  W2. What's the results of Adam+ L-BFGS as it is a very useful technique in training PINNs.\n\nThe Adam + L-BFGS optimization procedure was introduced in the original PINNs paper by Raissi et al [1]. At the time, Adam alone often struggled to reduce the PINN loss to sufficiently small values, motivating the use of L-BFGS as a second-stage optimizer for further refinement.\n\nSince then, numerous techniques—such as adaptive loss weighting, Fourier features, PirateNet, and causal training—have substantially improved the loss convergence of PINNs. With these advances, the training loss can now reach very small values (e.g., below $10^{-8}$ , as demonstrated in Figures 7–9 of our paper. In this regime, the additional benefit of L-BFGS becomes negligible. This is because the networks are typically trained with `float32` precision, making L-BFGS prone to numerical instability when the loss is already very small.\n\nFor benchmarks without time-marching, we conducted additional experiments by fine-tuning the trained models with Adam followed by L-BFGS. The results are summarized below. As can be seen, the effect of adding L-BFGS is minimal across all cases.\n\n| Benchmark         | Adam                        | Adam + L-BFGS               |\n|-------------------|-----------------------------|-----------------------------|\n| Wave              | $5.15 \\times 10^{-5}$       | $5.08 \\times 10^{-5}$       |\n| Burgers           | $8.20 \\times 10^{-5}$       | $8.20 \\times 10^{-5}$       |\n| Allen–Cahn        | $2.24 \\times 10^{-5}$       | $2.25 \\times 10^{-5}$       |\n| Korteweg–de Vries | $7.04 \\times 10^{-4}$       | $7.33 \\times 10^{-4}$       |\n\n> ###  W3. The sensitivity of performance to learning rates in SOAP vs. Adam is not deeply discussed.\n\nTo investigate the sensitivity of performance to the learning rate, we conducted an ablation study using the Allen–Cahn benchmark. The table below summarizes the final relative  $L^2$ errors obtained by Adam and SOAP across a range of learning rates.\n\nBoth optimizers demonstrate reasonable stability over learning rates ranging from 0.01 to 0.0001. However, Adam exhibits slightly higher sensitivity. In contrast, SOAP maintains consistently low error across all tested learning rates, highlighting its robustness to learning rate selection.\n\n| Learning Rate | Adam                      | SOAP                      |\n|---------------|---------------------------|---------------------------|\n| $0.01$        | $8.70 \\times 10^{-5}$     | $5.45 \\times 10^{-6}$     |\n| $0.005$       | $1.31 \\times 10^{-4}$     | $4.40 \\times 10^{-6}$     |\n| $0.001$       | $2.68 \\times 10^{-5}$     | $\\mathbf{2.84 \\times 10^{-6}}$ |\n| $0.0005$      | $\\mathbf{2.14 \\times 10^{-5}}$ | $3.87 \\times 10^{-6}$     |\n| $0.0001$      | $3.68 \\times 10^{-4}$     | $6.34 \\times 10^{-6}$     |\n\n[1] Raissi, Maziar, Paris Perdikaris, and George E. Karniadakis. \"Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations.\" Journal of Computational physics 378 (2019): 686-707."}, {"review_id": "UN0ja4o77E", "reviewer": "Reviewer_aauB", "summary": "This work focuses on the gradient alignment problem of PINNs, which belongs to an important topic that has been researched extensively. The authors locate the issue both theoretically and empirically. They propose to use quasi second-order methods, especially SOAP, to resolve gradient conflicts, which shows superior accuracy.", "questions": "1. The authors only focus on type II of the gradient conflicts, \"where gradients have similar magnitudes but opposing directions.\" It would be interesting to see type I of such during PINN's training, particularly while second-order methods are applied, and whether the two types are orthogonal. By orthogonality I mean whether different learning rate annealing algorithms can work with SOAP.\n2. As the authors listed, there are many quasi second-order approaches. Why does specifically SOAP (and sometimes Muon) perform better than others? Insights or analyses are appreciated.", "limitations": "They are mentioned above.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n\n- The decomposition of two types of gradient conflicts and inter-step/intra-step gradient alignment.\n- Theoretical explanation of the model's dynamics.\n- Good performance on a range of datasets\n\nWeaknesses:\n\n1. Some datasets are not very challenging, on which Adam has already achieved sub 1% error rates. I recommend testing SOAP on some hard datasets such as Poisson2d-MS and Wave2d-CG from PINNacle, where the boundary conditions are more difficult.\n2. Didn't show why second-order methods can help with inter-step gradient alignment.\n3. The method-wise contribution is limited by borrowing SOAP, which affects the novelty of this work.\n4. One key contribution as claimed is \"We reveal that SOAP [1] can be viewed as an efficient approximation of the Newton preconditioner.\" But this is revealed by SOAP, not this work.", "was_revised": true, "final_justification": "In discussions with authors, W2, W3, and Q1 are solved; Q2 is partially solved, and I hold concerns about W1 and W4. However, I think this paper is a good and insightful paper overall.", "rebuttal": "We sincerely thank the reviewers for their careful reading, thoughtful comments, and constructive suggestions. We appreciate the opportunity to clarify our contributions and address the raised concerns. Below, we respond to each point in detail.\n\n> ###  W1. Some datasets are not very challenging ...\n\n**Target Accuracy:** While sub-1% relative error may be acceptable in some machine learning tasks, it remains insufficient for many scientific computing applications, where high-fidelity simulations often require accuracy approaching machine precision. Our goal is therefore not merely to outperform Adam in moderately challenging cases, but to reduce errors **as much as possible** across a broad spectrum of PDEs. As demonstrated in our results, SOAP consistently outperforms our strong baseline with Adam by up to an order of magnitude.  This improvement is highly significant in scientific computing contexts.\n\n**Benchmarks:** We believe our current experimental suite is already **comprehensive and representative**. Our paper includes 10 diverse benchmarks, ranging from relatively simple problems (e.g., 1D wave equation) to significantly more complex ones, such as coupled reaction-diffusion systems and Navier–Stokes equations at high Reynolds numbers. This coverage exceeds that of most related PINNs papers published at top conferences [1–7], which typically evaluate on only 3–5 relatively simple tasks.\n\nImportantly, across all benchmarks, SOAP consistently demonstrates superior performance and robustness compared to other optimizers. This provides strong empirical support for our claim that (quasi) second-order preconditioning helps resolve directional gradient conflicts, thereby improving the accuracy of PINNs.\n\n> ###  W2. Didn't show why second-order methods can help with inter-step gradient alignment.\n\nIn Proposition 3, we prove that preconditioned gradient descent with preconditioner $P$ maintains high inter-step alignment, specifically  \n$\nA(g_t, g_{t+1}) \\geq 1 - \\epsilon\n$\n\nprovided the learning rate satisfies  \n$\n\\eta \\leq \\sqrt{\\frac{2\\epsilon \\|g_t\\|^2}{\\|H P^{-s} g_t\\|^2}}.\n$\n\nFor Newton’s method, where $P = H$ and s = 1, this upper bound simplifies to \n\n$ \\eta_{\\max} = \\sqrt{2\\epsilon} $, \n\nremoving dependence on the Hessian’s condition number. This enables larger and more stable learning rates while preserving alignment. In fact, we show that the closer the preconditioner is to the true Hessian, the greater the inter-step gradient alignment and the more aggressive the learning rate that can be safely used.\n\n\n>  ### W3. The method-wise contribution is limited by borrowing SOAP, which affects the novelty of this work.\n\nWhile SOAP is used as the primary optimizer, our contributions extend significantly beyond the application of an existing algorithm.\n\n- Our core novelty lies in identifying and formalizing **directional gradient conflicts** as a fundamental bottleneck in PINN optimization. To this end, we introduce *gradient alignment score* to quantify such conflicts, which generalizes cosine similarity to multiple loss components. \n\n- We further provide rigorous theoretical analysis demonstrating that such conflicts are intrinsic to PINNs at initialization (Proposition 2). Additionally, we establish a novel connection between SOAP and Newton’s method (Theorem 1), which offers a principled explanation for why second-order preconditioning can effectively resolve these conflicts.\n\n- Our empirical results—including the first successful PINN applications to high-Reynolds-number Kolmogorov flows  $\\text{Re} = 10^4$ . \n\nSOAP is one *instantiation* of a broader principle: that quasi second-order preconditioning mitigates directional conflicts in scientific ML tasks. Thus, the main methodological contribution lies in conceptually understanding, diagnosing, and overcoming optimization barriers in PINNs and related scientific ML models, **not** in the application of a new optimizer per se.\n\n\n> ###  W4. One key contribution as claimed is \"We reveal that SOAP can be viewed as an efficient approximation of the Newton preconditioner.\" But this is revealed by SOAP, not this work.\n\nWe respectfully disagree with this assessment and kindly ask the reviewer to revisit the original SOAP paper by Vyas el al. [10], which provides key theoretical insights:\n\n- Vyas el al. establishes that Shampoo with power 1/2 is equivalent to Adafactor operating in the eigenbasis of Shampoo’s preconditioner (Claim 1, Section 4.1).\n- Building on prior work by Morwani et al. (2024), Vyas el al. note that Shampoo with power 1/2 approximates the optimal Kronecker factorization of the Adagrad preconditioner.\n- Leveraging these insights, SOAP was designed to apply **Adam (rather than Adafactor)** in the eigenbasis of the Shampoo preconditioner, aiming for improved stability and efficiency.\n\nImportantly, the SOAP paper does **not** provide any theoretical analysis connecting SOAP to Newton's method, **nor** does it claim that SOAP approximates Newton steps. In contrast, our work is the first to formally establish this connection.\n\n\n> ###  Q1. The authors only focus on type II of the gradient conflicts...\n\nOur decision to study Type II gradient conflicts was deliberate and grounded in the following rationale:\n\n- **Type I conflicts** (i.e., imbalanced gradient magnitudes) have been extensively studied in the PINN literature:\n    - Prior work, including Wang et al. [8,9] and many follow-ups, has shown that such conflicts are prevalent.\n    - A wide range of adaptive loss weighting schemes (e.g., GradNorm, NTK-based weighting) has been developed to address them.\n\n- In contrast, **Type II conflicts** (i.e., directional gradient misalignment) remain largely underexplored:\n    - To the best of our knowledge, only a few works ([5,6]) have explored these issues.\n    - Our work is among the first to systematically demonstrate that directional conflicts are widespread in PINN training, see Figures 2 and 5 for empirical evidence using our proposed gradient alignment score.\n    - We further provide theoretical analysis showing that Type II conflicts are intrinsic to PINNs  even for simple equations like the 1D Poisson problem (Proposition 2).\n\nImportantly, the two types of conflicts often **coexist** and require **complementary solutions**. In our training procedure (described in Appendix G.2), we apply learning rate annealing (as in [7]) to help mitigate Type I conflicts. SOAP specifically targets Type II conflicts.  All ten benchmarks reported in our paper use this unified training protocol, demonstrating the robustness and compatibility of SOAP with standard loss-weighting techniques.\n\nIn summary, both conflict types are present in PINN training. Our work complements prior research on Type I conflicts by introducing new tools and theoretical insights to understand and resolve the less-addressed but equally important Type II conflicts.\n\n> ###  Q2. Why does specifically SOAP (and sometimes Muon) perform better than others?\n\nFormally, we prove that SOAP approximates the Newton preconditioner $H^{-1}$ whereas **Muon** and **Kron** approximate only $H^{-1/2}$. From **Proposition 3**, we show that full inverse preconditioning (i.e., $H^{-1} $) enables higher stable learning rates and better inter-step gradient alignment.\nWhile Muon and Kron provide some alignment benefits, they are inherently limited by the weaker curvature information encoded in their preconditioners. Consequently, their performance falls short of SOAP.\n\nIn summary, SOAP performs better because it more closely approximates Newton’s method, achieving superior alignment and optimization efficiency.\n\n-----\n\n[1] Characterizing possible failure modes in physics-informed neural networks, Neurips 2021\n\n[2] Separable Physics-Informed Neural Networks, Neurips 2023\n\n[3] Mitigating Propagation Failures in Physics-informed Neural Networksusing Retain-Resample-Release (R3) Sampling, ICML 2023\n\n[4] PINNsFormer: A Transformer-Based Framework For Physics-Informed Neural Networks, ICLR 2024\n\n[5] ConFIG: Towards Conflict-free Training of Physics Informed Neural Networks, ICLR 2025\n\n[6] Dual cone gradient descent for training physics-informed neural networks. Neurips 2024\n\n[7] PIG: Physics-Informed Gaussians as Adaptive Parametric Mesh Representations, ICLR 2025\n\n[8] Understanding and Mitigating Gradient Flow Pathologies in Physics-Informed Neural Networks, SISC\n\n[9] When and why PINNs fail to train: A neural tangent kernel perspective, JCP\n\n[10] Vyas, Nikhil, et al. \"Soap: Improving and stabilizing shampoo using adam.\" arXiv preprint arXiv:2409.11321 (2024)."}, {"review_id": "3bPkXEvvM8", "reviewer": "Reviewer_SoJw", "summary": "This paper studies the gradient alignment of multiple objectives in Physics-informed neural networks (PINNs). Training of PINNs faces optimization challenges due to the composite nature of the loss function, which consists of competing objectives (e.g. the initial and boundary conditions and a PDE residual). \nBy introducing a new alignment score, defined as \n\n$$\n\\mathcal{A}(v_1,v_2,\\dots,v_n) := 2 \\left|\\left| \\frac{\\sum_{i=1}^n \\frac{v_i}{\\||v_i\\||}}{n} \\right|\\right|^2 -1,\n$$\nwhich generalizes the cosine similarity measure to multiple vectors, the authors identify the misalignment of gradient directions of competing objectives as a critical obstacle. \nBuilding on this result, the authors show that (quasi) second-order optimization methods can mitigate these gradient conflicts resulting in superior performance over classically used methods, including Adam, and that a recently proposed optimizer SOAP performs best among the tested methods.\nFinally, this paper shows empirically that PINNs trained with SOAP are able to learn challenging PDEs, including turbulent flows, which was not possible before with PINNs.", "questions": "- I am a little bit puzzled regarding the inter-step alignment score of SOAP. The plots show that the alignment score is almost or exactly 1.0, which would imply that the trajectory of the iterates is almost straight. This seems a bit unlikely to me, given that all the methods use mini-batches. Could the authors perhaps clarify on this? \n\n- Can the authors also elaborate on why a very high inter-step alignment score close to 1 is desired? (which SOAP seems to achieve, see my question above.) Clearly having a negative or close-to zero inter-step alignment score is undesirable, but given the highly non-convex landscape, I would not expect the trajectory to be very straight. For instance, [1] also studies the directionality of trajectories in neural networks (although not in PINNs), but argue that a cosine similarity score of consecutive gradients close to 1 would be undesired. \n\n- I had a closer look at Figure 5 and noticed that the relative L2-error of SOAP dominates all other methods except for the Korteweg-De Vries problem, which aligns with the higher intra-step alignment score. Here Muon has the lowest relative L2-error for three quarter of the training trajectory, while having comparably low alignment scores, even lower than Adam and Kron (after initial 20% of steps). Do the authors have any explanation for this? \n\n- Given that all these methods have different computational costs, can the authors also provide figures for relative L2-error vs. wall-clock time? \n\n[1] Singh, Sidak Pal, et al. \"The directionality of optimization trajectories in neural networks.\" The Thirteenth International Conference on Learning Representations. 2025.", "limitations": "yes.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "Strengths\n- This paper is overall clearly written and easy to follow. \n- This paper introduces a novel alignment score, which shows that the update directions of multiple objectives (they are not strictly gradient directions anymore due to the diagonal preconditioning) in Adam are misaligned, while this is less the case in quasi second-order methods, such as Kron, Muon and SOAP. This alignment score also correlates strongly with the relative L2-error achieved by each optimizer across multiple benchmark problems.\n- I am not an expert in PINNs, but the empirical results that the authors achieve in this work with the SOAP optimizer appear to be quite impressive, which advances the state-of-the-art of what PINNs can achieve.\n- The benchmarks that the authors provide can serve as a valuable reference for future work on PINNs. \n\nWeaknesses\n- The novelty of this paper seems a bit limited, given that SOAP has been proposed already and that gradient conflicts appear to be a known problem in training PINNs given that there is previous work which propose specifically designed optimizers (e.g. ConFiG) to address this problem (even though the overall performance is below SOAP).\n- The provided explanation of SOAP's superior performance is not very convincing to me. Proposition 3 only addresses the inter-step alignment (which is equivalent to cosine similarity), not the intra-step alignment, which seems to be the more critical point regarding gradient conflicts. Furthermore, Proposition 3 only provides a statement regarding the maximal (local) step size that can be taken to maintain a certain gradient alignment between consecutive steps, but does neither explain why higher alignment scores between consecutive steps is necessarily better nor why a larger step size is more desired.", "was_revised": true, "final_justification": "After a thorough discussion with the authors the contributions of this submission were further clarified to me. \n\nConcretely, the benefit of the gradient alignment scores is clearer as it is not only a cheaper metric to track than the eigenspectrum of the Hessian (which can be used to study the ill-conditioning of the loss landscape). Moreover, when an explicit preconditioning matrix is not available (for instance because one does some form of implicit preconditioning such as in SOAP) or only coarse approximations are available, the gradient alignment scores appear to be a more reliable metric to study optimization behavior. \n\nFurthermore, during the discussion phase it became more clear that SOAP (while not the novelty of this work) is not \"just another quasi-Newton method\", which not only performs better, but also more robustly than more \"classic\" quasi-Newton and second order methods, such as NGD, which require high precision formats, such as float64 to converge and perform well, while SOAP also works robustly with float32. As the authors elaborate, this can be intuitively understood, as previous methods try to approximate the inverse of the Hessian, which is numerically sensitive, while SOAP performs implicit preconditioning. SOAP has been developed in the context of LLMs, where this problem of numerical sensitivity is not present. Thus, identifying this and transferring it to the domain of PINNs is valuable by itself.\n\nFinally, a further discussion provided by the authors putting this contribution in the broader context of other quasi-Newton methods will provide a good reference point for the community to train PINNs more reliably and robustly in the future. \n\nTherefore, after this discussion phase, I have decided to adjust my score from an initial 3 (\"Borderline reject\") to a 5 (\"Accept\").", "rebuttal": "We sincerely thank the reviewer for their careful reading and thoughtful feedback. We appreciate the opportunity to clarify our work and address the concerns raised. Below, we provide detailed responses and additional insights that we hope will strengthen the overall understanding of our contributions.\n\n> ###  W1. The novelty of this paper seems a bit limited...\n\nWhile SOAP is used as the primary optimizer, our contributions extend significantly beyond the application of an existing algorithm.\n\n- Our core novelty lies in identifying and formalizing *directional gradient conflicts* as a fundamental bottleneck in PINN optimization. To this end, we introduce *gradient alignment score* to quantify such conflicts, which generalizes cosine similarity to multiple loss components. \n\n- We further provide rigorous theoretical analysis demonstrating that such conflicts are intrinsic to PINNs at initialization (Proposition 2). Additionally, we establish a novel connection between SOAP and Newton’s method (Theorem 1), which offers a principled explanation for why second-order preconditioning can effectively resolve these conflicts.\n\n- Our empirical results—including the first successful PINN applications to high-Reynolds-number Kolmogorov flows  $\\text{Re} = 10^4$ . \n\nSOAP is one *instantiation* of a broader principle: that quasi second-order preconditioning mitigates directional conflicts in scientific ML tasks. Thus, the main methodological contribution lies in conceptually understanding, diagnosing, and overcoming optimization barriers in PINNs and related scientific ML models, **not** in the application of a new optimizer per se.\n\n> ###  W2. The provided explanation of SOAP's superior performance is not very convincing...\n\nOur rationale for SOAP’s superior performance is as follows: \n\n- Figures 2 and 5 show a strong positive correlation between inter-step gradient alignment and the predictive accuracy of PINNs, suggesting that higher alignment scores empirically lead to better performance.\n\n- In Theorem 1, we show that SOAP approximates Newton’s method under certain assumptions.\n\n- Proposition 3 further implies that more accurate curvature approximations—such as using the Hessian—lead to higher gradient alignment, enabling more stable updates and faster convergence.\n\nIn summary, these theoretical results and empirical observations together explain why SOAP’s approximate second-order nature yields improved gradient alignment and, consequently, superior PINN performance.\n\n> ###  Q1.The plots show that the alignment score is almost or exactly 1.0, which would imply that the trajectory of the iterates is almost straight.\n\nWe compute the inter-step alignment score between consecutive gradient steps and report these scores at regular intervals (e.g., every 100 iterations). To reduce the effect of stochastic noise, alignment scores are computed using the loss evaluated on the **full set of test collocation points** (i.e., without mini-batching), while training continues to use mini-batches. Note that our training batch size is large (e.g. 8192), so switching to mini-batches for computing alignment scores would not significantly affect the results.\n\nIt is worth noting that SOAP shows some small oscillations in alignment scores, these fluctuations are much less pronounced than those of other methods. To improve clarity and readability, we apply a running average to smooth all curves, which explains why the reviewers observe SOAP’s alignment scores as nearly constant.\n\nWe emphasize that a high inter-step alignment score close to 1 does not imply that the entire optimization path is nearly straight. Since the alignment is measured only between consecutive steps, small deviations can accumulate over many iterations, resulting in a significantly curved overall trajectory. Therefore, a high inter-step alignment score does not guarantee a perfectly straight optimization path throughout training.\n\n> ###  Q2. Why a very high inter-step alignment score close to 1 is possible and desired?\n\n**On possibility:** It is important to note that achieving an alignment score close to 1 can be **trivial** if the learning rate is made sufficiently small. In this case, parameter updates become negligible, and consecutive gradients naturally align since they hardly change. What truly matters is maintaining a nearly linear trajectory while using a reasonable step size.\n\n**On desirability:**\nIntuitively, high inter-step alignment indicates that the optimization trajectory is more linear, which is often a sign of stable and efficient convergence. This is empirically supported in our results (Figures 2 and 5), where higher alignment correlates with faster error convergence.\n\nTo build more intuition, consider the simplest case of a quartic objective function. If the Hessian is used as the preconditioner, the optimization trajectory becomes a **straight line**, resulting in an inter-step alignment score exactly equal to 1. In this idealized scenario, the optimizer reaches the local minimum in a single step.\n\nAlthough training deep networks involves complex, non-convex objectives, good preconditioners still help locally **whiten** the curvature, making the loss landscape more isotropic. This allows updates to approximate those in well-conditioned **quadratic** problems. As long as the curvature (i.e., Hessian) does not change drastically between steps, the optimization path remains stable, and high inter-step alignment naturally emerges—even in non-convex settings. Thus, high alignment is not only possible, but also reflects desirable optimization behavior when achieved under a meaningful learning rate.\n\n\n> ###  Q3. Alignment Scores in Korteweg-De Vries\n\nWhile our theoretical insights and metrics such as alignment scores explain general trends across problems, they are not sufficient to capture all **problem-specific dynamics**.\n\nThere are several potential reasons for this discrepancy:\n\n- **Alignment is a proxy**, not a guarantee of performance. It captures how well gradients are aligned across time (inter-step) or loss terms (intra-step), but does not account for other factors like curvature mismatch, optimizer noise characteristics, or problem-specific regularities.\n\n- Deep learning optimization inherently involves non-convex landscapes and complex interactions; thus, **our theoretical results—while broadly predictive—may not hold uniformly across all settings**. This instance could be one where assumptions behind our alignment-based interpretation do not fully apply.\n\n- Muon’s update rule may incidentally match the structure of the KdV problem better, even without high alignment scores. This may point to implicit regularization or favorable dynamics specific to this PDE.\n\nOverall, we believe this deviation is indeed intriguing and motivates future work toward a more general and dedicated theory that explains not just average trends but also exceptions.\n\n\n> ### Q4. provide figures for relative L2-error vs. wall-clock time.\n\nThank you for the suggestion — this is an excellent point. We will include such figures in the revised version. In our experiments, for the same amount of training time, SOAP consistently achieves the lowest relative L2-error compared to other methods."}, {"review_id": "Z7RkQNFVfF", "reviewer": "Reviewer_aDSZ", "summary": "The paper thoroughly investigates the issue of gradient alignment that arises during training of PINNs, which arise when combining the physics-informed loss (residuals of the PDE) and the supervised one (boundary conditions/ initial condition/ data). The authors propose a novel gradient alignment score, inspired from the cosine similarity, to quantify these conflicts and demonstrate how and why quasi-second-order optimization alleviate them. The paper combines the theoretical analysis with extensive numerical validation across benchmark PDEs.", "questions": "I have the following questions:\n* Second-order training for PINNs can be very effective even with small (approx. 2000 parameters) and vanilla NN architectures. How does the SOAP optimizer perform in this scenario? Why not including such results to highlight the effect of SOAP alone?\n* Based on the theoretical findings, I would have expected SOAP to reach much higher accuracy. Quasi-Newton methods as L-BFGS can at times face struggle when the hessian matrix they approximate is ill-conditioned (as often the case with PINNs). Can it be that the SOAP pre-conditioner is sensible to the approximation of singular hessians?", "limitations": "The authors clearly discuss the limitations of their method.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Overall, the paper is in a very good format. In particular:\n* The paper is very well written and provides a clean and complete understanding of the problem studied, both numerically and theoretically.\n* The underlying problem is a critical issue in PINNs for years, and the analysis proposed by the authors gives a clear and easy understanding of the problem\n* The alignment score introduced gives a very intuitive understanding of the empirical results, while carrying interesting theoretical advancements\n\nHowever, there are some factors that are concerning.\n* The problem of loss balancing in PINNs has been already studied and solved in several literature.\n* The authors strongly recommends the usage of SOAP, but the theoretical backbone suggests that a pure second-order method would be the optimal choice. Comparison with methods as Natural Gradient Descent, Levemberg-Marquardt or other quasi-Newton approaches would bring value to the choice of SOAP.\n* The improvements in the relative $L^2$ error shown in Table 2 are in general unimpressive. The usage of second-order methods can boost the accuracy of PINNs of several orders of magnitude, even when the underlying network is small (making it feasible to adopt such methods).\n* According to the experimental setup described in the appendix, PINNs trained in the numerical evaluation use random weight factorization, random Fourier features, a PirateNet architecture and gradient norm-based loss weighting. Some of these methods are implicitly addressing the gradient alignment issue, making it harder to understand the standalone impact of SOAP.", "was_revised": true, "final_justification": "The thorough discussion with the authors, especially the additional experiments and discussion on second-order optimizer gives a different cut to the paper, which I recommend acceptance for.", "rebuttal": "We appreciate the reviewer's careful reading and insightful questions, particularly regarding the comparison with second-order methods and the theoretical expectations for SOAP's performance. These points have prompted valuable clarifications. We will address each point in detail below.\n\n> ###  W1. The problem of loss balancing in PINNs has been already studied and solved in several literature.\n\nWhile the issue of loss balancing in PINNs has indeed received significant attention, we respectfully disagree with the claim that it has been **solved.** As summarized in our paper, gradient conflicts in PINNs can be broadly categorized into two types:  \n- **Type I**: Conflicts due to imbalanced gradient magnitudes across loss terms.  \n- **Type II**: Conflicts due to gradients having opposing directions.\n\nType I conflicts have been extensively explored, beginning with Wang et al. [1] and followed by numerous loss weighting schemes proposed in subsequent works. In contrast, **Type II conflicts remain largely underexplored**. To the best of our knowledge, only two recent works [2, 3] explicitly discuss this issue. \n\nMoreover, even methods specifically designed to handle Type II conflicts—such as ConFiG [2]—do not outperform our baseline on standard benchmarks (see Figure 4). This suggests that directional gradient conflict in PINNs remains an **open challenge**.\n\nOur work makes three key contributions in this regard:\n\n- We show that directional gradient conflicts (Type II) are pervasive during PINNs training, as evidenced by the proposed gradient alignment score (Figure 5).\n\n- We prove that Type II conflicts are intrinsic to PINNs optimization—even for simple problems like the Poisson equation (Proposition 2).\n\n- We demonstrate that quasi-second-order preconditioners, such as SOAP, can effectively alleviate Type II conflicts and lead to improved optimization and accuracy (Theorem 1 and Proposition 3).\n\nTherefore, while prior work has made important progress on loss balancing, especially in mitigating Type I conflicts, our results suggest that Type II conflicts are far from resolved. We believe our contributions offer new insights and practical tools for tackling this underexplored yet critical challenge in PINNs.\n\n\n\n> ###  W2, W3, Q1. Unimpressive Results from SOAP and Comparison with Second-Order Methods\n\nWe would like to clarify that while second-order methods (e.g., [4]) can achieve near machine-precision accuracy, this has only been demonstrated on relatively simple benchmarks—where solutions are smooth and exhibit no sharp transitions or discontinuities. In such cases, very small networks are often sufficient to approximate the solution, and convergence and memory issues typically do not arise.\n\nIndeed, **none** of the existing second-order methods have shown strong performance on the more challenging benchmarks we consider, which involve sharper features and more complex dynamics. This is not merely due to missing evaluations but second-order methods face fundamental scalability challenges: they do not support mini-batching, are sensitive to hyperparameters, and require high-precision `float64` arithmetic for stable convergence. \nHowever, `float64` is less efficiently supported on GPUs, consumes **2x** much memory, and is typically **2–4×** slower than `float32`, which is the precision used in **most neural network training—including all experiments reported in this manuscript.**\n\nTo better illustrate these issues, we revisited the 2D Poisson and heat benchmarks from [4], which uses natural gradient descent (NGD) to train PINNs for solving PDEs. Using the official implementation from [4], we varied only the hyperparameters such as random seed and depth or width of the network and evaluated their performance under both `float32` and `float64` precision, comparing against SOAP method.  \n\nWe tested MLPs with one to three hidden layers, such as [2, 32, 1], [2, 32, 32, 1], and [2, 32, 32, 32, 1], and used wider networks with the same depth for SOAP, such as [2, 256, 1], [2, 256, 256, 1], and [2, 256, 256, 256, 1]. All models were trained to convergence over ten random seeds.  It is worth noting that a perfectly fair comparison is inherently difficult due to differing architecture preferences (e.g., SOAP benefits from overparameterization). Therefore, we report the best result for each method under its respective optimal hyperparameters.\n\nWe note that for some seeds, NGD failed to converge and produced highly inaccurate results; we excluded these runs from the mean/standard deviation. However, such failures are also acknowledged in [4], as reflected in the large maximum errors reported in Table 3 and Table 5 of that work. The mean and standard deviation of the resulting relative $L_2$ errors are summarized in the table below.\n\n### Poisson 2D\n\n| Architecture     (Method)               | Float32                | Float64                |\n|----|------|---|\n| [2, 32, 1] (NGD)                    | 4.87e-2 ± 3.61e-2       | 3.84e-7 ± 2.92e-7       |\n| [2, 32, 32, 1]   (NGD)               | 1.65e-1 ± 3.93e-2       | 1.27e-6 ± 2.86e-7       |\n| [2, 32, 32, 32, 1]  (NGD)            | 2.48e-1 ± 1.75e-2       | 3.14e-6 ± 5.10e-7       |\n| [2, 256, 1]  (SOAP)   | 3.06e-6 ± 7.12e-7   | 6.08e-7 ± 2.13e-7   |\n| [2, 256, 256, 1]  (SOAP)     | 1.87e-6 ± 5.19e-7   | 4.06e-7 ± 1.81e-7  |\n| [2, 256, 256, 256, 1]  (SOAP)     | **1.35e-6 ± 3.45e-7**   | **2.99e-7 ± 1.05e-7**   |\n\n### Heat 2D\n\n| Architecture    (Method)              | Float32                | Float64                |\n|----|---|--|\n| [2, 32, 1]        (NGD)              | 5.98e-2 ± 5.46e-2       | 7.68e-6 ± 1.85e-6       |\n| [2, 32, 32, 1]   (NGD)               | 5.95e-1 ± 2.29e-4       | 2.32e-6 ± 1.17e-6       |\n| [2, 32, 32, 32, 1]    (NGD)         | 5.95e-1 ± 4.58e-4       | 5.13e-6 ± 5.25e-7       |\n| [2, 256, 1]   (SOAP)             | 4.61e-6 ± 8.12e-7       | 3.03e-6 ± 6.06e-7       |\n| [2, 256, 256, 1]    (SOAP)       | 2.74e-6 ± 9.52e-7       | 2.04e-6 ± 4.10e-7       |\n| [2, 256, 256, 256, 1]   (SOAP)   | **2.65e-6 ± 5.00e-7**   | **1.33e-6 ± 2.61e-7**   |\n\nThese results highlight several key limitations of NGD:\n\n- **Extreme sensitivity to precision**: NGD's performance deteriorates significantly under `float32`, often leading to divergence.\n- **Inconsistent scalability**: Increasing network depth/width does not always improve accuracy for NGD and often leads to worse performance or instability.\n\nIn contrast, SOAP achieves consistently strong performance across architectures and numerical precision, demonstrating both robustness and scalability. It can match or outperform NGD in accuracy, while avoiding the computational cost and stability issues associated with second-order methods.\n\n> ### W4. Combined Techniques in PINN Training May Confound SOAP’s Effect\n\nTo clarify, the alignment scores reported in Figures 2 and 5 are computed using the full training procedure described in the manuscript (Appendix G.2), which includes random feature embeddings, the PirateNet architecture, and gradient norm-based loss weighting, etc.\n\nWhile it is true that these components can improve optimization to some extent, the effect of SOAP is both **distinct and consistent**. Specifically, under identical training configurations, we observe that SOAP significantly enhances both intra-step and inter-step gradient alignment compared to baseline methods (Figures 2 and 5). This suggests that SOAP directly addresses gradient misalignment and provides benefits beyond what is achievable by architecture or loss-balancing heuristics alone.\n\n> ### Q2. Concerns About SOAP’s Accuracy and Sensitivity to Ill-Conditioned Hessians\n\nWe thank the reviewer for this insightful comment. Our theoretical analysis (see Theorem 1) draws a connection between SOAP and Newton’s method under certain assumptions. However, as is often the case with theoretical guarantees in deep learning, these assumptions may not hold precisely in practice—particularly in the highly nonlinear and nonconvex optimization landscape of PINNs. Consequently, it is not unexpected that SOAP may not consistently match the idealized performance of a true second-order method.\n\nRegarding the comparison to L-BFGS: while L-BFGS approximates the inverse Hessian using historical gradients and updates, it is known to be sensitive to noises, especially in mini-batch or stochastic training regimes common in PINNs. From our experience, we find that L-BFGS often underperforms our baselines  likely due to these instabilities.\n\nIn contrast, SOAP demonstrates greater robustness to mini-batch noise and appears to yield a more stable approximation of the inverse Hessian, even in ill-conditioned settings. This empirical robustness suggests that SOAP may better handle singular or near-singular Hessians. Nonetheless, a formal theoretical analysis of SOAP’s behavior in such regimes remains an important direction for future work.\n\n\n[1] Understanding and mitigating gradient flow pathologies in physics-informed neural networks, SIAM SISC\n\n[2] ConFIG: Towards Conflict-free Training of Physics Informed Neural Networks, ICLR 2025\n\n[3] Dual cone gradient descent for training physics-informed neural networks. Neurips 2024\n\n[4] Achieving High Accuracy with PINNs via Energy Natural Gradient Descent, ICML 2023"}]

Author rebuttal:

Meta-review: This paper identifies directional gradient conflicts as a key challenge in training PINNs and introduces a gradient alignment score to diagnose them. It shows that quasi-second-order methods mitigate these conflicts, with the SOAP optimizer—closely tied to Newton’s method—achieving 2–10× accuracy gains on 10 PDE benchmarks, including turbulent flows up to Reynolds 10,000.
The work is praised for its clear formalization, intuitive alignment score, strong theory–experiment balance, and effective optimizer design. Concerns include limited novelty, narrow theory, modest gains relative to known second-order methods, missing baselines, and unclear standalone impact of SOAP given prior conflict-mitigation techniques.

During the rebuttal period, the authors offer comprehensive comparisons with NGD, a new theoretical link between gradient alignment and faster loss decay, detailed ablations on hyperparameters and costs, and an extended discussion of the benefits and limitations of second-order methods, which solve most of the concerns of reviewers. During the AC-reviewer discussion period, Reviewer SoJw and Reviewer aDSZ express a positive attitude toward acceptance. 

Overall, the pros outweigh the cons, and the AC recommends for acceptance, while strongly encourages the authors to incorporate the rebuttal into the final version of the paper.
Decision: Accept (poster)