Gradient Surgery for Physics-Informed Neural Networks

• The general PDE setup uses operators N and B for PDE dynamics and boundary conditions, with source terms f, g, and h.Source: §1 · p. 2
• PINNs represent the solution u(x, t) with a parameterised neural approximation, computing the required spatial and temporal derivatives by automatic differentiation.Source: §1 · p. 2
• The PDE, boundary, and initial conditions define separate losses that a shared-parameter network jointly minimises.The PDE residual is evaluated at collocation points across the spatio-temporal domain; boundary and initial losses impose soft constraints at the spatial boundaries and initial time.Source: §1 · p. 2

• Angle-based conflict occurs when task gradients have a negative inner product, meaning their directions are misaligned by more than 90°.Misalignment can inhibit shared descent; opposing gradients can cancel updates and halt progress for both tasks.Source: §2 · pp. 2–3
• Magnitude-based conflict occurs when one task's substantially larger gradient norm dominates the update, potentially biasing it toward that task.Unlike angular conflict, magnitude imbalance does not necessarily mean the objectives oppose each other.Source: §2 · p. 3
• Gradient surgery modifies task-specific gradients before aggregation, using a strategy-specific combination rule to produce the layer update.Source: §2 · p. 3

• Loss-balancing methods estimate task weights using uncertainty, inverse losses, loss-change rates, gradient magnitudes, random sampling, or update effectiveness.These methods aim to prevent any task from dominating training.Source: §3 · p. 3
• Gradient-aggregation methods include Nash bargaining, multi-objective KKT directions, and approaches that decompose, decorrelate, align, or otherwise modify task gradients.Source: §3 · p. 3
• CONFIG is tailored to PINNs, using a normalised task-gradient space and pseudoinverse-based projection so physics constraints contribute positively to the shared update.SAM-GS instead targets magnitude conflicts with momentum-based stabilisation of gradient magnitudes across tasks.Source: §3 · p. 3

• The benchmarks cover 1D Burgers, 1D Schrödinger, 2D Kovasznay flow, and 3D Beltrami flow, spanning nonlinear convection, dispersive dynamics, and incompressible fluid behaviour.Source: §4 · p. 4
• The comparison includes loss-reweighting methods DWA, UW-SO, and FAMO, alongside gradient methods PCGrad, GradDrop, IMTL, CAGrad, Aligned-MTL, CONFIG, Nash-MTL, and SAM-GS.Source: §4 · p. 4
• Evaluation measures test MSE for initial conditions, boundary conditions, PDE residuals, and overall error, then aggregates performance using mean rank and relative improvement over vanilla joint training.Lower mean rank and relative improvement indicate better performance.Source: §4 · p. 4

5 Addressing Gradient Conflicts in PINNs

• The analysis measures gradient alignment with average cosine gradient angle and scale imbalance with average gradient magnitude ratio during vanilla PINN training using Adam.The cosine angle serves as a proxy for interference, while the magnitude ratio measures scale imbalance across tasks.Source: §5 · p. 4
• Across Fitting, Diffusion, and Total Diffusion, angle conflict is negligible, rises, then decreases, while magnitude similarity first rises and later deteriorates.The phase pattern is observed in interactions among different loss-term gradients within a single batch, not solely in batch variability.Source: §5 · p. 4, Figure 1 · p. 5
• PAM-GS detects magnitude conflict when Ψt < γ and equalises task gradients, while cos(Φt) < 0 triggers momentum-based reweighing for angle conflict.When neither conflict is detected, it directly aggregates task gradients; momentum tracks task-gradient dynamics and magnitude similarity over time.Source: §5 · p. 5
• The approach prioritises magnitude-conflict detection because it has been reported to substantially hinder training during fitting, while angle conflicts can decelerate convergence.Existing methods often couple the two conflict types; SAM-GS assumes they co-occur, which may apply similarity-based updates when magnitude imbalance dominates.Source: §5 · pp. 4, 6

Evaluation setup

• Methods, including baselines, are independently trained across five random weight initialisations, with parameters selected by hyperparameter optimisation.The evaluation averages results across the five initialisations.Source: §6 · p. 6

1D Burgers’ equation

• On Burgers, PAM-GS improves over SAM-GS and is competitive with CAGrad, whose lower MSE difference falls within observed variability (SD).The results highlight limitations of loss-balancing methods, particularly recent dynamic variants.Source: §6 · p. 6

1D Schrödinger equation

• On Schrödinger, PAM-GS ranks third in overall MSE but has the best aggregate relative improvement: ∆M% = −77.01 versus ∆M% = −57.00 for the lowest-MSE method.Results are reported separately for the real, imaginary, and magnitude components.Source: §6 · p. 7

3D Beltrami flow

• PAM-GS consistently performs strongest across most Beltrami metrics and constraint categories, with MR values approaching 1.Table 3 reports field-variable MSE and its decomposition across IC, BC, and PDE residual terms.Source: Table 3 · p. 6, §6 · p. 7

2D Kovasznay flow

• PAM-GS achieves the best overall performance across Kovasznay constraints, while phase-adapted interventions yield stable improvement over SAM-GS.The reported MSE covers boundary conditions and the PDE residual.Source: §6 · p. 7, Table 4 · p. 7

Loss landscape

• PAM-GS yields the shortest trajectory in the Kovasznay loss landscape, where optimisers traverse flat regions, sharp valleys, and high-curvature areas differently.The visual uses a 2D PCA subspace of parameter space.Source: §6 · p. 7, Figure 3 · p. 8

Cross-benchmark performance

• PAM-GS performs particularly well in overall MSE on Kovasznay and Beltrami and remains among the top three methods on Burgers and Schrödinger.CAGrad leads on Burgers and Schrödinger but ranks lower on Kovasznay and Beltrami; Aligned-MTL shows the opposite pattern.Source: §7 · p. 8

Task-balanced performance

• By MR and ∆M%, PAM-GS achieves the best results on three of four benchmarks and remains top-two on Burgers.These metrics account for relative performance across tasks, unlike MSE alone.Source: §7 · p. 8

Conflict-aware design

• PAM-GS consistently outperforms SAM-GS, supporting the benefit of explicit conflict detection.The paper attributes PAM-GS's cross-benchmark competitiveness to its PINN-specific, conflict-aware design being less sensitive to problem-dependent conflict patterns.Source: §7 · p. 8

Findings

• Across multiple PINN benchmarks, angle- and magnitude-based conflicts alternate, with reduced occurrence of one type coinciding with the other.The conclusion presents this pattern as motivation for treating the conflict types differently.Source: §8 · p. 8

Proposed method

• PAM-GS explicitly accounts for the different conflict types encountered during PINN optimisation.Its behaviour is selected according to detected conflict types.Source: §8 · p. 8

Implication

• The conclusion emphasizes characterizing conflicts throughout training rather than treating them as static, linking this approach to more effective and balanced optimisation over tasks.Source: §8 · p. 8

1D Burgers Equation

• The 1D Burgers benchmark uses viscosity 0.01/π on x ∈ [−1, 1], t ∈ [0, 1], with a sine initial condition and zero endpoint values.No closed-form solution is available, so PhiFlow numerical solutions serve as reference data.Source: A Partial Differential Equation… · p. 10

1D Schrödinger Equation

• The nonlinear Schrödinger benchmark evolves h(x, t) = u(x, t) + iv(x, t) on x ∈ [−5, 5], t ∈ [0, π/2], with periodic boundaries and h(x, 0) = 2 sech(x).DeepXDE numerical results provide the benchmark because no analytical solution is available under these conditions.Source: A Partial Differential Equation… · p. 10

2D Kovasznay Flow

• Kovasznay flow is a steady two-dimensional incompressible Navier–Stokes solution, with Re = 40 on (x, y) ∈ [−0.5, 1] × [−0.5, 1.5].The flow uses exact-solution Dirichlet boundary values.Source: A Partial Differential Equation… · p. 11

3D Beltrami Flow

• Beltrami flow is an unsteady three-dimensional Navier–Stokes solution with velocity components and pressure, using a = d = 1 and Re = 1.Its domain is x, y, z ∈ [−1, 1], t ∈ [0, 1], with boundary and initial conditions prescribed from the analytical solution.Source: A Partial Differential Equation… · pp. 10–11

Training protocol

• Each model uses four 50-neuron tanh hidden layers, Xavier initialisation, Latin hypercube sampling, and Adam; learning rates follow cosine annealing except in extended training.Extended training uses a fixed 10^-4 learning rate; otherwise, cosine annealing runs from 10^-3 to 10^-4 with a 100-epoch warm-up.Source: B Training Details · p. 11
• Training budgets are 3 × 10^4 epochs for Burgers and 10^5 for the other benchmarks, while PAM-GS fixes β1 = 0.9 and β2 = 0.99 and varies γ by benchmark.Interior, boundary, and initial-condition collocation counts vary by benchmark; PAM-GS γ values are 0.6, 0.4, 0.4, and 0.7 for Burgers, Schrödinger, Kovasznay, and Beltrami.Source: B Training Details · p. 11

Runtime and Memory Analysis

• PAM-GS has no substantial computational or memory differences from the evaluated methods despite maintaining additional momentum statistics.The analysis averages memory use and training time across three independent 2D Kovasznay runs over 1,000 steps on an NVIDIA H200 GPU.Source: B Training Details · p. 11

Table label

• The supplied table label identifies Table 5 as Table B.Source: Table 5 · p. 12

Runtime and memory

• The Table B caption identifies peak GPU memory and training time as the quantities reported for evaluated methods.Source: C Comprehensive Results for the 3D… · p. 12

3D Beltrami results

• Table 6 compares gradient-surgery methods on 3D Beltrami flow using MSE values in units of 10^-5, with lower values indicating better performance.Source: Table 6 · p. 12