PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Sifan Wang, Ananyae Kumar Bhartari, Bowen Li, Paris Perdikaris

arXiv:2502.00604v2cs.LGcs.AIphysics.comp-ph

TL;DR

PINNs must optimize competing physical and data-fitting objectives, but directional gradient conflicts remain insufficiently addressed. The paper introduces a multi-gradient alignment score and analyzes second-order preconditioning, showing that SOAP approximates Newton-style preconditioning and delivers strong results across PDE benchmarks. These findings connect gradient alignment with more reliable PINN optimization while leaving computational efficiency and larger-scale application as open challenges.

  • Problem

    PINNs face convergence and accuracy limitations, while directional conflicts among gradients from competing loss terms remain an important unresolved issue.

  • Method

    The paper introduces a multi-gradient alignment metric and theoretical and empirical analyses of second-order optimization, including SOAP as an approximation to the Newton preconditioner.

  • Results

    Across 10 PDE benchmarks, SOAP achieves state-of-the-art performance, including turbulent-flow results up to Reynolds number 10,000 and 2-10x accuracy improvements.

  • Takeaways & Limitations

    Gradient alignment and second-order preconditioning provide a framework for understanding and mitigating conflicts in PINN optimization.

  • Takeaways & Limitations

    SOAP requires approximately 2x longer training than baselines, and scaling alignment-based preconditioned methods to larger systems remains challenging.

Abstract

from arXiv · show

Multi-task learning through composite loss functions is fundamental to modern deep learning, yet optimizing competing objectives remains challenging. We present new theoretical and practical approaches for addressing directional conflicts between loss terms, demonstrating their effectiveness in physics-informed neural networks (PINNs) where such conflicts are particularly challenging to resolve. Through theoretical analysis, we demonstrate how these conflicts limit first-order methods and show that second-order optimization naturally resolves them through implicit gradient alignment. We prove that SOAP, a recently proposed quasi-Newton method, efficiently approximates the Hessian preconditioner, enabling breakthrough performance in PINNs: state-of-the-art results on 10 challenging PDE benchmarks, including the first successful application to turbulent flows with Reynolds numbers up to 10,000, with 2-10x accuracy improvements over existing methods. We also introduce a novel gradient alignment score that generalizes cosine similarity to multiple gradients, providing a practical tool for analyzing optimization dynamics. Our findings establish frameworks for understanding and resolving gradient conflicts, with broad implications for optimization beyond scientific computing.

1 Introduction

PINNs face convergence and accuracy limitations, particularly from directional conflicts among loss-term gradients. This work introduces gradient alignment analysis and shows that second-order methods, especially SOAP, address these conflicts while improving benchmark performance.

  • PINNs remain limited by convergence speed and accuracy despite advances in architectures, sampling, optimizers, and training strategies.
  • Directional conflicts between gradients from different loss terms remain a critical issue beyond imbalanced gradient magnitudes.
  • The proposed gradient alignment metric extends cosine similarity to quantify directional conflicts among multiple loss terms.
  • Higher gradient conflict scores are linked to slower convergence across various PDE systems.
  • Second-order optimizers enhance gradient alignment by implicitly preconditioning the loss landscape.
  • Across 10 PDE benchmarks, including turbulent flows with Reynolds numbers up to 10,000, the approach achieves 2-10x accuracy improvements.

2 Overview of PINNs

PINNs approximate PDE solutions with neural networks trained through composite losses that enforce physical constraints and fit data. Automatic differentiation supplies PDE residuals, while the loss combines residual, initial-condition, and boundary-condition objectives.

  • PINNs approximate PDE solutions with neural networks while enforcing physical constraints and fitting data through composite loss functions.
  • Automatic differentiation computes PDE residuals from the neural-network approximation.
  • The composite loss combines multiple competing objectives associated with the PDE and prescribed conditions.
  • Collocation points may be fixed mesh vertices or randomly sampled during each training iteration.

3 Gradient Alignment in PINNs

Gradient conflicts in PINNs arise from magnitude imbalance and opposing directions among loss-term gradients, producing inefficient optimization trajectories. The paper defines alignment scores, links them to loss reduction, and shows that quasi-second-order methods improve alignment and convergence.

  • Gradient conflicts: PINNs experience magnitude-imbalanced and directionally opposing gradients, with dominant terms or conflicting objectives producing inefficient optimization.
  • Alignment metric: The alignment score extends cosine similarity from two vectors to multiple vectors and ranges from -1 to 1.
  • Alignment metric: Intra-step alignment measures conflicts among individual loss-term gradients, while inter-step alignment measures conflicts across consecutive optimization steps.
  • Alignment metric: The intra-step score evaluates conflicts across separate Navier-Stokes loss components rather than grouping them together.
  • Theory: Higher intra- and inter-step alignment scores directly accelerate loss reduction under the stated preconditioned-gradient assumptions.
  • Experiments: First-order optimizers show early alignment oscillations near or below zero, whereas quasi-second-order methods maintain higher positive alignment and faster test-error convergence.

4 Intra-step Gradient Alignment of PINNs at Initialization

The section analyzes gradient alignment near initialization for PINNs using a tractable two-layer tanh network solving the one-dimensional Laplace equation. It shows that all considered optimizers lack consistent intra-step alignment near initialization, while SOAP maintains perfect inter-step alignment during training.

  • Theoretical result: At initialization, the alignment score converges to a binary random variable in the infinite-width limit.This result is stated as Proposition 3 under the small-initialization assumptions.
  • Theoretical result: All optimizers fail to induce consistent intra-step alignment between ∇L_b and ∇L_r near initialization.The paper identifies this as a direct implication of the absence of a preferred direction and validates it empirically in Figure 5.
  • Optimization behavior: SOAP maintains a perfect inter-step alignment score of 1 throughout the entire training process.The paper presents this behavior as distinctive relative to the other optimizers and motivates the subsequent theoretical explanation.

5 Gradient Alignment in Quasi Second-Order Optimization

This section explains how preconditioned quasi-second-order optimization promotes gradient alignment and relaxes learning-rate constraints. It presents SOAP as a tractable Hessian-preconditioning approximation with strong accuracy and stability across PINN benchmarks, while contrasting its practical advantages with NGD and first-order conflict-resolution methods.

  • Theory: A preconditioned gradient-descent analysis establishes learning-rate conditions required to maintain alignment A(g_t, g_t+1) ≥ 1 − ϵ.The bound is then specialized to different preconditioners, including vanilla gradient descent and Newton’s method.
  • Theory: Newton’s method removes the Hessian-condition-number dependency from the maximum learning rate, allowing a constant bound regardless of problem conditioning.The paper interprets this as evidence that Hessian-approximating preconditioners relax learning-rate constraints while preserving trajectory consistency.
  • SOAP interpretation: SOAP can be interpreted, under stated assumptions, as approximating the Hessian as a preconditioner.The formal assumptions and proof are provided in Appendix D.
  • Practical limitations: NGD requires Fisher-matrix inversion, float64 precision, twice the memory, and 2-4x slower execution than float32.These practical limitations restrict NGD demonstrations to relatively simple benchmarks and contribute to failures on more challenging PDEs.
  • Benchmark comparison: SOAP reliably converges and outperforms NGD in accuracy and stability across the evaluated Poisson and Heat 2D benchmarks.The comparison uses MLPs with varying depth and width, ten random seeds per method, and both float32 and float64 precision.
  • Gradient alignment: Quasi-second-order methods promote directional gradient alignment through preconditioning while mitigating ill-conditioning and retaining computational tractability.SOAP achieves the highest gradient alignment scores, which the paper relates to its closer approximation to Newton’s method.
  • Benchmark results: Quasi-second-order optimizers show consistent performance improvements across a wide range of PDE benchmarks, with SOAP achieving the best overall results.The paper attributes SOAP’s suitability to PINN multi-objective optimization partly to its closer alignment with Newton’s method.

6 Experiments

Experiments across challenging PDE benchmarks show that SOAP consistently improves PINN accuracy, including turbulent and coupled fluid dynamics, while exposing computational trade-offs and hyperparameter sensitivities.

  • 10 challenging PDE benchmarks span wave propagation, shock formation, chaotic systems, reaction-diffusion, fluid dynamics, and heat transfer.
  • SOAP achieves the best overall results among the evaluated optimizers across a wide range of PDE benchmarks.The comparison uses relative L2 error over the entire spatial-temporal domain.
  • 6.4x lower relative error occurs on the wave equation, while Allen-Cahn improves by 6.9x and Korteweg-de Vries and Kuramoto-Sivashinsky improve by approximately 2x.
  • Complex fluid dynamics: 3.99% relative L2 error is achieved for lid-driven cavity flow at Reynolds number 5,000, reduced from 32.4%.The model captures secondary and tertiary corner vortices and agrees with benchmark results.
  • Complex fluid dynamics: 0.52% relative L2 error is achieved for Rayleigh-Taylor instability, compared with 7.32% for the best baseline.The model captures mushroom-shaped structures and both linear-growth and nonlinear-development phases across multiple time steps.
  • Complex fluid dynamics: 3.20% relative L2 error is achieved for turbulent Kolmogorov flow at Reynolds number 10,000, versus 20.4% with the baseline.The predictions reproduce large-scale structures and smaller eddies, with higher spectral energy at high wavenumbers than traditional solvers using a 1024 × 1024 grid.
  • Ablation studies: SOAP consistently improves accuracy across standard MLP, modified MLP, and PirateNet backbones on four benchmark PDEs.
  • Ablation studies: More frequent preconditioner updates improve results with diminishing gains beyond update frequency 2, while β1 = 0.99 outperforms β1 = 0.01.SOAP requires approximately 2x longer training time than baselines, although it reaches approximately 10^-4 accuracy within the first 10,000 iterations.

7 Conclusion

The paper frames gradient conflicts as a central PINN optimization issue and combines alignment metrics, theoretical analysis, and quasi-second-order optimization to address them. Across 10 PDE benchmarks, SOAP achieves state-of-the-art performance, while the authors identify computational efficiency and scaling as remaining challenges.

  • The work quantifies directional gradient conflicts, shows their prevalence during PINN optimization, and links second-order optimization to implicit gradient alignment.
  • More efficient preconditioned algorithms and scaling to larger systems remain open research directions.

B Analysis of Intra-step Gradient Alignment

The analysis studies intra-step gradient conflicts in PINNs near initialization, showing that the alignment score has no preferred direction and that standard optimizers fail to resolve these conflicts initially.

  • Analytical setup: The analysis uses a two-layer tanh network for the one-dimensional Laplace equation under small initialization to obtain tractable effective gradient dynamics.The authors state that the analysis can be extended to other PDEs.
  • Initialization and effective gradients: Under small Gaussian initialization, residual and boundary gradients have different leading orders, with ∇θLr(θ) scaling as ε^7 and ∇θLbc(θ) as ε.The normalized parameters are introduced by scaling weights with ε^-1.
  • Optimizer comparison: GD, Adam, Shampoo, and SOAP therefore fail to induce consistent intra-step alignment between residual and boundary gradients near initialization.This theoretical prediction agrees with the near-zero initial intra-step scores shown in Figure 5.
  • Initialization and effective gradients: At initialization and infinite width, the intra-step alignment score converges to a binary random variable with expectation zero.The result is established for GD, Adam, Shampoo, and SOAP, with an optimizer-dependent constant in the characterization.
  • Inter-step alignment analysis: For preconditioned gradient descent, successive-gradient alignment remains close to 1 when the second-order term is small, imposing an upper bound on the learning rate.The bound is derived from a second-order Taylor expansion involving the Hessian and preconditioner.

D Connection between SOAP and Newton’s method

This section connects SOAP to Newton’s method through rotated-space preconditioning and a Gauss–Newton approximation, explaining its gradient-alignment behavior relative to Adam and Shampoo.

  • Rotated-space optimization: SOAP performs optimization in a transformed space aligned with the preconditioner’s principal directions, then maps the update back to the original parameter space.The procedure projects weights and gradients, applies Adam in rotated coordinates, and transforms the result back.
  • Hessian approximation: Assuming the Gauss–Newton component approximates the true Hessian, the equivalent rotated-space preconditioner is diagonal.Its diagonal structure follows from the Kronecker-factorized representation in principal-curvature coordinates.
  • SOAP and Newton equivalence: Under the stated assumptions, SOAP’s update approximates Newton’s method.The connection is established by combining the rotated-space preconditioning analysis with the Adam update characterization.
  • Hessian approximation: SOAP effectively approximates the block-diagonal Gaussian Newton component layer by layer through Kronecker factorization.The paper links this structure to naturally promoted gradient alignment across optimization steps.
  • Comparison with Adam and Shampoo: Adam requires a highly restrictive diagonal-Hessian assumption, whereas Shampoo uses the square root of the Gauss–Newton component as its preconditioner.These structural differences are presented as explaining SOAP’s better gradient alignment than both methods.

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

The section places Shampoo and Muon within a family of quasi-second-order methods by relating their matrix preconditioners to Adagrad covariance structure and fractional Hessian powers.

  • Adagrad foundation: Adagrad uses accumulated gradient covariance as a preconditioner for online parameter updates.Its preconditioner is constructed from the accumulated outer products of gradients.
  • Shampoo: Shampoo approximates the full-matrix Adagrad preconditioner with a Kronecker product of two statistical matrices, transforming the update for vectorized parameters into a matrix update.The approximation is expressed through layerwise left and right factors.
  • Shampoo: When p = q = 2, the generalized fractional-preconditioner construction yields the standard Shampoo update.The parameters satisfy 1/p + 1/q = 1.
  • Shampoo: Shampoo’s preconditioner is described as a Kronecker product approximation of the layerwise Gauss–Newton component of the Hessian.This interpretation supplies its quasi-second-order connection.
  • Muon: Muon uses Newton–Schulz iterations to approximate a matrix preconditioner and, without accumulation, can be viewed as approximating H^-1/2.Its updates operate on neural-network weight matrices with momentum.

F.2 Proof of Proposition 2

The proof analyzes cumulative loss decrease under smoothness and relates the resulting bound to inter-step gradient alignment.

  • Proof structure: The proof uses L-smoothness to establish a single-step loss-drop bound.The argument begins from the descent lemma.
  • Proof structure: It then applies L-smoothness across two steps to derive a cumulative decrease involving consecutive gradient alignment.The proof explicitly uses the inner product between successive update directions.
  • Monotonicity condition: For η ≤ 1/(LM), the relevant loss-drop expression is nonnegative, yielding monotonicity.The displayed expression is η/M (1 − LηM) ≥ 0.

G.1 Architectures

The architectures combine modified MLP and PirateNet components with Fourier features, dual encoders, gating, residual blocks, and random weight factorization. These choices target complex PDE solutions, spectral bias, stable deep training, and improved residual minimization.

  • Modified MLP: The modified MLP uses two parallel encoders whose features are merged through point-wise multiplication at each hidden layer.This design demonstrates superior performance in minimizing PDE residuals compared with standard MLPs.
  • PirateNet: PirateNet transforms coordinates with random Fourier features to improve learning of high-frequency components and multiscale features.The embedding mitigates spectral bias by improving the eigenfunction frequency of the Neural Tangent Kernel.
  • PirateNet: Each PirateNet residual block contains three dense layers, dual gating operations, and an adaptive residual connection controlled by trainable α(l) parameters.α(l) = 0 gives an identity mapping, whereas α(l) = 1 gives a fully nonlinear transformation.
  • PirateNet: Initializing α(l) = 0 starts the network effectively shallow and allows learned depth growth during training.The resulting linear structure also permits direct incorporation of prior solution data through least-squares fitting.
  • Periodic boundary conditions: Periodic boundary conditions are encoded through Fourier embeddings so network outputs inherently satisfy the specified periodicity.The construction extends from one-dimensional to higher-dimensional periodic domains and can concatenate time coordinates for time-dependent problems.
  • Random weight factorization: Random weight factorization parameterizes each weight matrix as diag(exp(s)) · V and optimizes the scale and factor matrices directly.The implementation uses exponential parameterization to maintain non-zero scale factors across varied magnitudes.

G.2 Training pipeline

The training pipeline combines causality-aware weighting, curriculum and time-marching strategies, adaptive loss balancing, a PirateNet backbone, and comparisons among several optimizers. These procedures address temporal difficulty, gradient imbalance, and computational trade-offs across PDE benchmarks.

  • Causal training: Causal training exponentially decreases later temporal weights according to cumulative residuals from preceding time segments.The causality parameter ϵ must be tuned moderately because very small or large values can impair optimization.
  • Learning rate annealing: Learning-rate annealing dynamically equalizes gradient norms across weighted loss components to prevent bias toward any particular term.Weights are updated as running averages, typically every 100-1000 iterations, with minimal computational overhead.
  • Curriculum training and time-marching: Curriculum training decomposes complex PINN optimization into sequential sub-tasks for time-dependent PDEs and singular perturbation problems.Temporal domain decomposition passes each window’s final prediction as the next window’s initial condition.
  • Architecture and configuration: The experiments use PirateNet with three residual blocks, nine total layers, width 256, Tanh activations, and random weight factorization.Exact periodic boundary conditions are enforced when applicable.
  • Training protocol: Training uses 8,192 randomly sampled collocation points per iteration, warm-up from 0 to 10^-3 over 5,000 steps, and exponential decay.Loss weights are updated every 1,000 iterations using a moving average.
  • Optimizers: The optimizer comparison includes Adam, SOAP, Muon, and Kron with specified momentum and matrix-iteration configurations.SOAP uses β1 = 0.99 and β2 = 0.999, while Adam, Muon, and Kron use β1 = 0.9 and β2 = 0.999.

G.6 Benchmarks

The benchmark suite covers ten PDEs spanning wave propagation, shocks, chaotic systems, reaction-diffusion, fluid dynamics, and heat transfer. The reported experiments compare reference solutions, predictions, training behavior, test errors, accuracy, and runtime across PINN methods.

  • Benchmark visualizations: The figures compare reference solutions with model predictions and show training-loss or test-error trajectories for Adam and SOAP across the PDE cases.These comparisons cover Wave, Burgers, Allen-Cahn, Korteweg–De Vries, Kuramoto-Sivashinsky, Grey-Scott, Ginzburg-Landau, and Lid-driven Cavity equations.
  • Accuracy results: An improvement of up to one order of magnitude in accuracy is reported on the Allen-Cahn benchmark against different PINN variants.The comparison uses relative L2 test errors.
  • Benchmark scope: 10 representative PDE benchmarks span wave propagation, shock formation, chaotic systems, reaction-diffusion, fluid dynamics, and heat transfer.The benchmark setup includes PDE parameters, initial and boundary conditions, numerical implementations, and supplementary visualizations.
  • Computational costs: SOAP requires approximately 2x longer training time than baselines while achieving better convergence.The conclusion identifies more efficient preconditioned algorithms as a remaining need.
Loading 2502.00604v2…