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

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.

W1. The novelty of this paper seems a bit limited...

While SOAP is used as the primary optimizer, our contributions extend significantly beyond the application of an existing algorithm.

• 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.

• 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.

• Our empirical results—including the first successful PINN applications to high-Reynolds-number Kolmogorov flows $\text{Re} = 10^4$ .

SOAP 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.

W2. The provided explanation of SOAP's superior performance is not very convincing...

Our rationale for SOAP’s superior performance is as follows:

• 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.

• In Theorem 1, we show that SOAP approximates Newton’s method under certain assumptions.

• 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.

In 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.

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.

We 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.

It 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.

We 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.

Q2. Why a very high inter-step alignment score close to 1 is possible and desired?

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.

On desirability: Intuitively, 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.

To 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.

Although 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.

Q3. Alignment Scores in Korteweg-De Vries

While our theoretical insights and metrics such as alignment scores explain general trends across problems, they are not sufficient to capture all problem-specific dynamics.

There are several potential reasons for this discrepancy:

• 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.

• 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.

• 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.

Overall, 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.

Q4. provide figures for relative L2-error vs. wall-clock time.

Thank 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.

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.

W1. The problem of loss balancing in PINNs has been already studied and solved in several literature.

While 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:

• Type I: Conflicts due to imbalanced gradient magnitudes across loss terms.
• Type II: Conflicts due to gradients having opposing directions.

Type 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.

Moreover, 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.

Our work makes three key contributions in this regard:

• We show that directional gradient conflicts (Type II) are pervasive during PINNs training, as evidenced by the proposed gradient alignment score (Figure 5).

• We prove that Type II conflicts are intrinsic to PINNs optimization—even for simple problems like the Poisson equation (Proposition 2).

• 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).

Therefore, 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.

W2, W3, Q1. Unimpressive Results from SOAP and Comparison with Second-Order Methods

We 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.

Indeed, 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. However, 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.

To 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.

We 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.

We 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.

Poisson 2D

Heat 2D

These results highlight several key limitations of NGD:

• Extreme sensitivity to precision: NGD's performance deteriorates significantly under float32, often leading to divergence.
• Inconsistent scalability: Increasing network depth/width does not always improve accuracy for NGD and often leads to worse performance or instability.

In 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.

W4. Combined Techniques in PINN Training May Confound SOAP’s Effect

To 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.

While 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.

Q2. Concerns About SOAP’s Accuracy and Sensitivity to Ill-Conditioned Hessians

We 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.

Regarding 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.

In 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.

[1] Understanding and mitigating gradient flow pathologies in physics-informed neural networks, SIAM SISC

[2] ConFIG: Towards Conflict-free Training of Physics Informed Neural Networks, ICLR 2025

[3] Dual cone gradient descent for training physics-informed neural networks. Neurips 2024

[4] Achieving High Accuracy with PINNs via Energy Natural Gradient Descent, ICML 2023

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.

W1. Some datasets are not very challenging ...

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.

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.

Importantly, 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.

W2. Didn't show why second-order methods can help with inter-step gradient alignment.

In Proposition 3, we prove that preconditioned gradient descent with preconditioner $P$ maintains high inter-step alignment, specifically
$A(g_t, g_{t+1}) \geq 1 - \epsilon$

provided the learning rate satisfies
$\eta \leq \sqrt{\frac{2\epsilon \|g_t\|^2}{\|H P^{-s} g_t\|^2}}.$

For Newton’s method, where $P = H$ and s = 1, this upper bound simplifies to

$\eta_{\max} = \sqrt{2\epsilon}$,

removing 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.

W3. The method-wise contribution is limited by borrowing SOAP, which affects the novelty of this work.

While SOAP is used as the primary optimizer, our contributions extend significantly beyond the application of an existing algorithm.

• 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.

• 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.

• Our empirical results—including the first successful PINN applications to high-Reynolds-number Kolmogorov flows $\text{Re} = 10^4$ .

SOAP 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.

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.

We 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:

• 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).
• 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.
• 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.

Importantly, 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.

Q1. The authors only focus on type II of the gradient conflicts...

Our decision to study Type II gradient conflicts was deliberate and grounded in the following rationale:

• Type I conflicts (i.e., imbalanced gradient magnitudes) have been extensively studied in the PINN literature:

  • Prior work, including Wang et al. [8,9] and many follow-ups, has shown that such conflicts are prevalent.
  • A wide range of adaptive loss weighting schemes (e.g., GradNorm, NTK-based weighting) has been developed to address them.

• In contrast, Type II conflicts (i.e., directional gradient misalignment) remain largely underexplored:

  • To the best of our knowledge, only a few works ([5,6]) have explored these issues.
  • 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.
  • 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).

Importantly, 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.

In 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.

Q2. Why does specifically SOAP (and sometimes Muon) perform better than others?

Formally, 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. While 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.

In summary, SOAP performs better because it more closely approximates Newton’s method, achieving superior alignment and optimization efficiency.

---

[1] Characterizing possible failure modes in physics-informed neural networks, Neurips 2021

[2] Separable Physics-Informed Neural Networks, Neurips 2023

[3] Mitigating Propagation Failures in Physics-informed Neural Networksusing Retain-Resample-Release (R3) Sampling, ICML 2023

[4] PINNsFormer: A Transformer-Based Framework For Physics-Informed Neural Networks, ICLR 2024

[5] ConFIG: Towards Conflict-free Training of Physics Informed Neural Networks, ICLR 2025

[6] Dual cone gradient descent for training physics-informed neural networks. Neurips 2024

[7] PIG: Physics-Informed Gaussians as Adaptive Parametric Mesh Representations, ICLR 2025

[8] Understanding and Mitigating Gradient Flow Pathologies in Physics-Informed Neural Networks, SISC

[9] When and why PINNs fail to train: A neural tangent kernel perspective, JCP

[10] Vyas, Nikhil, et al. "Soap: Improving and stabilizing shampoo using adam." arXiv preprint arXiv:2409.11321 (2024).

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.

W1. The theoretical analysis focuses mainly on 1D Laplace equations at initialization.

We 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.

W2. What's the results of Adam+ L-BFGS as it is a very useful technique in training PINNs.

The 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.

Since 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.

For 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.

W3. The sensitivity of performance to learning rates in SOAP vs. Adam is not deeply discussed.

To 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.

Both 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.

[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.