Gradient Alignment in Physics-informed Neural Networks: A Second-Order Optimization Perspective

1 Introduction

• Directional conflicts between gradients from different losses remain underexplored, despite prior efforts to address imbalanced gradient magnitudes in PINN training.The introduction identifies these conflicts as the gap the work aims to investigate and resolve.Source: §1 · p. 2
• The work introduces a gradient alignment metric that extends cosine similarity to quantify directional conflicts among multiple loss terms.Source: §1 · p. 2
• In various PDE systems, gradient conflicts hinder PINN training, and higher conflict scores are linked to slower convergence.Source: §1 · p. 2
• Second-order optimizers enhance gradient alignment through implicit preconditioning, and SOAP [1] is presented as an efficient approximation of the Newton preconditioner for addressing gradient conflicts.The introduction connects this optimization mechanism to more reliable neural PDE solvers for complex physical systems.Source: §1 · p. 2
• Across 10 PDE benchmarks, the work reports 2-10x accuracy improvements and the first successful PINN application to turbulent flows with Reynolds numbers up to 10,000.These results are presented as comprehensive experimental evidence for the approach.Source: §1 · p. 2

2 Overview of PINNs

• PINNs address competing objectives by minimizing a composite loss that enforces physical constraints and data fitting, building on.The loss targets initial conditions, boundary conditions, and PDE residuals.Source: §2 · pp. 2–3
• A neural network uθ(t, x) approximates the PDE solution u(t, x), and automatic differentiation computes the PDE residual.Source: §2 · p. 2
• The composite loss encapsulates multiple competing objectives, including fitting the initial and boundary conditions and PDE residuals.Source: §2 · pp. 2–3
• Training points may be fixed mesh vertices or randomly sampled at each iteration.Source: §2 · p. 3

• PINN loss terms can conflict through unequal gradient magnitudes or opposing directions; for Navier–Stokes, enforcing no-slip wall conditions can conflict with bulk mass and momentum conservation.First-order methods follow an average gradient, producing inefficient zigzagging that becomes especially acute in turbulent flows.Source: §3 · p. 3
• The alignment score extends cosine similarity to multiple vectors, ranges from -1 to 1, and supports intra-step and inter-step measures of gradient alignment.For two vectors, it equals cosine similarity: 1 indicates perfect alignment, 0 orthogonality, and -1 complete opposition.Source: §3 · p. 3
• Under the stated preconditioned-gradient-descent assumptions, higher intra-step and inter-step alignment scores directly accelerate loss reduction.The paper notes that equal norms for weighted gradients are a practical assumption in PINNs, where loss-scale balancing is common.Source: §3 · p. 4
• Across five representative PDEs, first-order optimizers show early alignment oscillations near or below zero, while quasi-second-order methods maintain positive alignment and faster test-error convergence; SOAP has the highest scores.The optimizer comparison computes scores after each optimizer’s gradient transformations, capturing the actual update directions.Source: §3 · pp. 4–5, Figure 2 · p. 4

4 Intra-step Gradient Alignment of PINNs at Initialization

• The analysis uses a two-layer tanh network with small weight initialization to solve the one-dimensional Laplace equation.The authors state that this tractable setup captures key phenomena and that the results extend to more general PDEs.Source: §4 · p. 5
• At initialization, the alignment score converges to a binary random variable at infinite width, and no tested optimizer induces consistent intra-step alignment.The result assumes i.i.d. Gaussian weights with small ε; the optimizer set includes GD, Adam, Shampoo, and SOAP.Source: §4 · p. 5
• SOAP maintains a perfect inter-step alignment score of 1 throughout training.The authors identify its distinctive inter-step behavior as the subject of the following section’s theoretical explanation.Source: §4 · p. 5

5 Gradient Alignment in Quasi Second-Order Optimization

• Newton’s method permits a constant maximum learning rate independent of Hessian conditioning, while Hessian-approximating preconditioners relax learning-rate constraints and support larger steps with consistent optimization trajectories.The comparison is with vanilla gradient descent, whose learning-rate constraint depends on the Hessian condition number.Source: §5 · p. 6
• Under some assumptions, SOAP can be interpreted as approximating the Hessian as a preconditioner.The formal assumptions and proof are provided in Appendix D.Source: §5 · p. 6
• L-BFGS is unsuitable for large-scale or stochastic training, while first-order methods ConFiG [58] and DCGD [59] offer incremental improvements but remain limited by first-order updates.The section also notes that NGD is costly and does not scale well to challenging PDEs, where it can be hyperparameter-sensitive and often diverges.Source: §5 · p. 6
• SOAP reliably converged and outperformed NGD in accuracy and stability on the tested 2D Poisson and heat benchmarks, while NGD failed for some random seeds.The comparison reports relative L2 errors across MLP architectures in float32 and float64; the authors tested ten random seeds per method and reported the best results.Source: §5 · pp. 6–7, Table 2 · p. 7
• Quasi-second-order methods such as SOAP [1], Kron, and Muon promote directional gradient conflicts via preconditioning; SOAP consistently achieves the highest gradient alignment scores, likely due to its closer approximation to Newton’s method.Source: §5 · p. 7

Baselines

• The evaluation covers 10 representative PDEs, and comparisons focus on PINN methods scalable to large neural networks.The authors characterize their baseline pipeline as state-of-the-art and report that it outperforms most existing approaches on the Allen-Cahn benchmark.Source: §6 · pp. 7–8

State-of-the-art results

• SOAP achieves the best overall benchmark results, cutting relative error 6.4x on wave, 6.9x on Allen-Cahn, ~2x on KdV/Kuramoto-Sivashinsky, and an order of magnitude on Grey-Scott/Ginzburg-Landau.Table 3 compares Adam, Adam+L-BFGS, and SOAP using relative L2 error over the full spatial-temporal domain.Source: §6 · p. 8, Table 3 · p. 9

Complex fluid dynamics

• SOAP lowers relative L2 error from 32.4% to 3.99% for cavity flow (Re=5,000), from 20.4% to 3.20% for Kolmogorov flow (Re=10,000), and from 7.32% to 0.52% for Rayleigh-Taylor.The authors report this as the first successful PINN capture of turbulent dynamics at Re=10,000; its solution retains more high-wavenumber spectral energy than traditional solvers, including one using a 1024 × 1024 grid.Source: §6 · pp. 8–9, Figure 3 · p. 8

Ablation studies

• Across Wave, Burgers, Allen-Cahn, and KdV tests, SOAP improves accuracy over Adam with MLP, modified MLP, and PirateNet; PirateNet performs best overall.More frequent preconditioner updates improve accuracy with diminishing gains beyond frequency 2, while β1=0.99 performs best and β1=0.01 degrades results.Source: §6 · p. 9, Figure 4 · p. 10
• ConFiG shows hyperparameter sensitivity and inconsistent performance in these experiments, whereas SOAP’s optimization approach is reported as more robust.ConFiG is described as a method for addressing gradient conflicts in PINNs.Source: §6 · p. 10

Computational costs

• SOAP takes approximately 2x longer to train than baselines, but typically reaches approximately 10^-4 accuracy within 10,000 iterations.The authors suggest that its rapid initial convergence may permit up to 10x lower training time while maintaining competitive performance.Source: §6 · p. 10

7 Conclusion

• Across 10 PDE benchmarks, SOAP achieves state-of-the-art performance, confirming the effectiveness of quasi second-order optimizers.The conclusion reports this result without listing benchmark-specific values.Source: §7 · p. 10
• Gradient alignment scores quantify conflicts that theoretical analysis and empirical evidence show are prevalent during PINN optimization.Source: §7 · p. 10
• Second-order methods implicitly promote gradient alignment to mitigate conflicts, and the paper connects SOAP with Newton’s method.Source: §7 · p. 10
• Future work could develop more efficient preconditioned algorithms that preserve effectiveness at reduced computational cost.The conclusion presents this as an opportunity beyond SOAP’s handling of coupled physical constraints.Source: §7 · p. 10
• Gradient alignment and second-order preconditioning may benefit other competing-objective applications, though scaling to larger systems remains challenging.Source: §7 · p. 10

A Nomenclature

• The nomenclature uses uppercase letters for matrices and lowercase letters for their vectorized forms, with gradients and Hessians taken with respect to loss L.Source: Table 4 · p. 17

B Analysis of Intra-step Gradient Alignment

• The intra-step analysis uses a two-layer tanh network for the one-dimensional Laplace equation and assumes i.i.d. Gaussian weights initialized at small ε = o(1).The authors state that the analysis can be extended to other PDEs.Source: B Analysis of Intra-step Gradient… · p. 17
• The residual and boundary gradient norms scale as ∥∇θLr(θ)∥ = O(ε^7) and ∥∇θLbc(θ)∥ = O(ε), respectively.The analysis attributes this imbalance to the small-initialization regime.Source: B Analysis of Intra-step Gradient… · pp. 18–19
• At initialization, residual–boundary alignment converges to a binary random variable with zero expectation at infinite width, and GD, Adam, Shampoo, and SOAP fail to resolve the conflict.The prose links these failures to near-zero initial intra-step alignment scores in Figure 5.Source: Figure 5 · p. 18, B Analysis of Intra-step Gradient… · p. 20
• At early iterations, Adam behaves similarly to sign gradient descent, and the analysis applies it separately to the residual and boundary losses.This characterization motivates the examination of Adam’s initial updates in the small-initialization setup.Source: B Analysis of Intra-step Gradient… · p. 21
• At initialization, SOAP’s gradient conflict follows the Shampoo case because its transformed vector gradient is normalized before being mapped back.The passage states that this yields the same initial gradient conflict as Shampoo.Source: B Analysis of Intra-step Gradient… · p. 21

C Inter-step gradient alignment of preconditioned gradient descent

• For preconditioned gradient descent, successive gradients remain aligned when the second-order term is small, and a derived learning-rate upper bound ensures A(gt, gt+1) ≥ 1 − ϵ.The proposition considers a twice-differentiable loss and a positive definite preconditioner with uniformly bounded inverse.Source: C Inter-step gradient alignment of… · pp. 22–23
• For vanilla GD (P = I), the specialized learning-rate bound involves λmax(H), the largest eigenvalue of H.The passage identifies this as the vanilla-GD case of the derived bound.Source: C Inter-step gradient alignment of… · p. 23

• SOAP projects weights and gradients into the preconditioner’s eigenspace, applies Adam there, and transforms the updated weights back.Its covariance-based preconditioners are updated every f iterations in practice to reduce computational overhead.Source: D Connection between SOAP and Newton’s… · pp. 23–24
• Under the stated Gauss–Newton approximation, the rotated preconditioner is diagonal, supporting SOAP’s Newton-like interpretation.The analysis is limited to networks trained with cross-entropy loss and assumes the Gauss–Newton component approximates the true Hessian.Source: D Connection between SOAP and Newton’s… · pp. 24–25
• Under Proposition 6’s assumptions, SOAP’s update approximates Newton’s method in the original parameter space.Proposition 6 assumes a local minimum θ* and θ − θ* ∼ N(0, σ²I).Source: D Connection between SOAP and Newton’s… · pp. 25–26
• SOAP’s block-diagonal, layerwise Kronecker structure naturally promotes gradient alignment across optimization steps.The paper attributes this to approximating the block-diagonal Gaussian-Newton component of the Hessian in a rotated space.Source: D Connection between SOAP and Newton’s… · pp. 24, 26
• SOAP achieves better gradient alignment than Adam and Shampoo, whose Hessian approximations use different structures.Adam’s Newton approximation requires a diagonal Hessian; Shampoo instead uses the square root of the Gauss–Newton component as its preconditioner [86].Source: D Connection between SOAP and Newton’s… · p. 26

E Connection of Shampoo and Muon to Quasi-second-order Methods

• Adagrad uses accumulated gradient covariance as a preconditioner for parameter updates.The excerpt introduces parameters θ_t, gradients g_t, and a learning rate η, but omits the complete displayed equations.Source: E Connection of Shampoo and Muon to… · p. 26
• Shampoo tracks left and right gradient-statistic matrices and approximates the full-matrix Adagrad preconditioner with their Kronecker product.The approximation is given as (L_t ⊗ R_t)^(1/2), yielding a Shampoo update for parameter matrix W.Source: E Connection of Shampoo and Muon to… · p. 26
• Lemma 3 [87] bounds the full Adagrad preconditioner for rank-at-most-r gradient matrices, supporting an approximation of its preconditioned gradient.For conjugate p,q > 0 with 1/p + 1/q = 1, the corresponding approximation uses L_t^(1/p); p = q = 2 gives standard Shampoo.Source: E Connection of Shampoo and Muon to… · p. 27
• Shampoo’s preconditioner is also described as a Kronecker-product approximation to the Gauss–Newton component of the layerwise Hessian [86].Source: E Connection of Shampoo and Muon to… · p. 27
• Muon updates matrix-representable neural-network weights using momentum and Newton–Schulz iteration to approximately compute a matrix expression.The supplied excerpt truncates the expression after (M_t M_t^T.Source: E Connection of Shampoo and Muon to… · p. 27
• Without preconditioner accumulation, Muon is described as approximating the full-matrix Adagrad preconditioner, with additional orthogonalization benefits.Source: E Connection of Shampoo and Muon to… · p. 27

F Additional Lemma and Proof

• Lemma 4 [86], Corollary 2, is stated under the assumption that reshaping the Hessian tensor H_GN has rank 1.The supplied passage gives the assumption but not the lemma’s conclusion.Source: F Additional Lemma and Proof · p. 27

F.1 Proof of Proposition 1

• The cited passage contains the expression 2 + 2 v1 · v2, with no surrounding context provided in the supplied text.Source: §F.1 · p. 27

F.2 Proof of Proposition 2

• The single-step drop is derived using the L-smoothness descent lemma.The supplied passage introduces this step of the proof but does not include the resulting bound.Source: §F.2 · p. 27
• The proof obtains equal-norm formulas and a refined bound from the stated relation involving nλ2.The relation is truncated in the supplied passage, so its full conditions and formulas are unavailable.Source: §F.2 · p. 28
• The two-step cumulative drop applies L-smoothness with θ = θt−1 and ϕ = θt+1.This passage states the points at which smoothness is applied, but omits the full resulting bound.Source: §F.2 · p. 28
• The cited fragment states an inequality involving consecutive gradients, printed as t−1ht ≥ 1 M ⟨ht−1, ht⟩; it does not show how that inequality is used.Source: §F.2 · p. 28
• For η ≤ 1/(LM), the coefficient ηM(1 − LηM) is nonnegative, yielding monotonicity.The passage identifies this coefficient as the basis for the monotonicity result.Source: §F.2 · p. 28
• The proof rearranges an inequality involving Ainter to obtain the stated result.The supplied fragment does not show the full inequality or coefficient.Source: §F.2 · p. 28

Modified MLP

• Modified MLP uses dual input encoders with pointwise feature multiplication at each hidden layer, costing more computation but outperforming standard MLPs on PDE residual minimization.Source: §G.1 · pp. 28–29

PirateNet

• PirateNet [27] uses random Fourier features [91] to mitigate spectral bias and gated residual blocks initialized with α(l)=0, enabling shallow starts and least-squares fitting to prior data.The α(l) parameters control block nonlinearity: 0 gives an identity mapping, while 1 gives a fully nonlinear transformation.Source: §G.1 · p. 29

Exact imposition of periodic boundary conditions

• Periodic boundary conditions are imposed exactly using sine-cosine Fourier embeddings [70], extended to 2D and combined with time coordinates for time-dependent problems.Source: §G.1 · p. 30

Random weight factorization

• RWF [69] factorizes each weight matrix as diag(exp(s))·V and optimizes scale and direction parameters; µ=0.5 or 1, σ=0.1 are recommended for convergence and accuracy.Source: §G.1 · p. 30

G.2 Training pipeline

• Causal training partitions time into M segments and exponentially downweights later residuals using cumulative earlier losses; a moderate ε should let all temporal weights reach 1 by completion.Small ε may insufficiently enforce causality, while large ε can make optimization difficult by requiring extremely small early-time residuals.Source: §G.2 · pp. 30–31
• Self-adaptive learning-rate annealing balances weighted loss components by equalizing their gradient norms, with running-average weight updates that stabilize stochastic gradient descent.Updates occur at user-specified intervals, typically every 100–1000 iterations, with minimal computational overhead [21].Source: §G.2 · p. 31
• Curriculum training [42] decomposes difficult problems into subtasks; time-marching transfers each window’s final prediction as the next initial condition, at added retraining cost.Causal weighting is also applied within each window and complements time-marching because causality violations may persist within windows.Source: §G.2 · pp. 31–32

G.3 Data Generation

• Reference data use Chebfun in MATLAB and IncompressibleNavierStokes in Julia, with dt = 10^-4 and temporal downsampling.Source: §G.3 · p. 32

G.4 Hyper-parameters

• PirateNet [27] uses three residual blocks (9 layers), width 256, Tanh activations, and random weight factorization with μ = 1.0 and σ = 0.1; periodic boundaries are exact when applicable [69] [70].Source: §G.4 · p. 32
• Training uses 8,192 randomly sampled collocation points per iteration, a 5,000-step warm-up to 10^-3, then exponential learning-rate decay by 0.9, with causal tolerance 1.0.Loss weights update every 1,000 iterations; challenging benchmarks also use time-marching and curriculum learning. Decay is tailored to reach approximately 10^-7 by training end, and window counts are chosen empirically.Source: §G.4 · pp. 32–33
• Adam [71], Muon, and Kron use β1 = 0.9 and β2 = 0.999; SOAP [1] uses β1 = 0.99 and β2 = 0.999; Muon uses 10 Newton–Schulz iterations.Muon uses coefficients (2, −1.5, 0.5) for orthogonalization.Source: §G.4 · pp. 32–33

G.5 Computational Cost

• The JAX-PI [72] implementation ran all experiments on a single NVIDIA A6000 GPU, with runtime benchmarks reported in Table 7.Source: §G.5 · p. 33

G.6 Benchmarks

• Benchmark figures show reference-solution comparisons and Adam/SOAP training-loss and test-error trajectories for wave, Burgers, Allen–Cahn, and Korteweg–De Vries equations.The captions describe the plotted comparisons and trajectories but do not provide numerical values or establish which optimizer performs better.Source: Figure 6 · p. 34, Figure 7 · p. 35, Figure 8 · p. 35
• Allen–Cahn test errors achieve state-of-the-art performance among PINN variants, with accuracy improved by up to one order of magnitude.The comparison reports relative L2 test errors; the supplied passages do not provide individual table values.Source: §G.6 · p. 35, Table 8 · p. 36

NeurIPS Paper Checklist

• The checklist says the paper’s contributions include a gradient alignment score, second-order gradient alignment, SOAP’s Newton connection, and improved results across 10 PDE benchmarks.It states that the abstract and introduction accurately reflect these contributions and the paper’s scope.Source: NeurIPS Paper Checklist · p. 44
• SOAP requires approximately 2x longer training than baselines, although it achieves better convergence; the conclusion calls for more efficient preconditioned algorithms.The checklist identifies computational cost as a limitation discussed in the paper.Source: NeurIPS Paper Checklist · p. 45
• The paper gives assumptions and proofs for its theoretical results, including propositions on alignment score and cosine similarity, initialization alignment, and preconditioned gradient descent and Newton’s method.Complete proofs are provided in the appendix, as referenced in Sections 3–4.Source: NeurIPS Paper Checklist · p. 45
• The paper reports enough information to reproduce its main experiments, with baseline setup in Section 5 and experimental settings, hyperparameters, and data generation in the appendices.Source: NeurIPS Paper Checklist · p. 46
• The authors say they will release code upon acceptance and point to Section 5 for training and test details and Appendix G.5 for computational resources.The checklist separately says code implementing the methods will be released upon acceptance.Source: NeurIPS Paper Checklist · pp. 46, 47, 48, 50
• The checklist cites comparisons across PDEs and optimizers, Figure 4 ablations, and consistent improvements in Table 2 as support for statistical significance.It does not specify error bars, confidence intervals, or significance tests in this justification.Source: NeurIPS Paper Checklist · p. 47
• The checklist says the research conforms to the NeurIPS Code of Ethics and credits existing work and software, including JAX, Matplotlib, Chebfun, and NumPy.It also says the baseline implementation builds on appropriately cited prior work.Source: NeurIPS Paper Checklist · pp. 48–49
• The authors identify positive scientific and engineering impacts, acknowledge possible ill-intentioned uses, expect no specific negative impact, and mark safeguards as not applicable.They attribute the positive impacts to more accurate physical-system simulation and gradient-alignment principles applicable beyond scientific computing.Source: NeurIPS Paper Checklist · pp. 48–49
• The research is described as involving neither crowdsourcing or human subjects nor applicable IRB review, and the checklist marks core-method LLM use as not applicable.The LLM justification says the paper does not indicate LLMs as part of its core methods or research.Source: NeurIPS Paper Checklist · pp. 50–51