PaperlayerView exact revision on arXiv

Source-linked AI summary

Gradient-Update Mismatch: Rethinking Conflict-Free Training of Physics-Informed Neural Networks

Jing Xiao, Xinhai Chen, Qinglin Wang, Menghan Jia, Zhiquan Lai, Dongsheng Li, Jie Liu, Tiejun Li

arXiv:2609.01558v1cs.LG

TL;DR

PINN loss terms can induce conflicting gradients, and gradient surgery may lose conflict-freeness when optimizers transform its direction. The paper introduces GUM to describe this mismatch and GUA to project optimizer proposals back into the conflict-free cone, with optional state alignment. Experiments show GUM is prevalent and that GUA eliminates applied-update conflicts while improving PINN performance.

  • Problem

    PINN loss terms can produce conflicting gradients, and conflict-freeness of a gradient-surgery direction is not guaranteed after optimizer transformation.

  • Method

    GUA projects the optimizer proposal onto the current conflict-free cone and optionally adjusts optimizer state toward targets reconstructed from the applied update.

  • Results

    GUA achieves conflict-free applied updates across evaluated PINN settings and consistently improves various gradient surgery methods, while GUM is widespread across evaluated optimizers.

  • Takeaways & Limitations

    Conflict handling in PINN optimization should account for the optimizer proposal and applied update, not only the pre-transformation surgery direction.

  • Takeaways & Limitations

    The formulation uses local first-order gradient geometry and does not explicitly capture curvature, finite-step, or longer-horizon optimization effects.

Abstract

from arXiv · show

Training Physics-Informed Neural Networks (PINNs) requires jointly optimizing physics residual and initial/boundary condition loss terms, which often induce conflicting gradients. Gradient surgery methods mitigate this issue by constructing directions from loss-specific gradients to reduce conflict before optimizer transformation. However, even when the constructed direction is conflict-free, this property may not be preserved after optimizer transformation. Let $a_t$ denote the direction constructed by gradient surgery, $u_t$ the optimizer proposal, and $\mathcal{C}_t$ the conflict-free cone induced by the loss-specific gradients. We show that modern optimizers can transform $a_t$ through mechanisms such as historical state, adaptive scaling, preconditioning, or decoupled weight decay, so $a_t \in \mathcal{C}_t$ does not generally imply $u_t \in \mathcal{C}_t$. We refer to this optimizer-induced discrepancy in conflict-freeness between $a_t$ and $u_t$ as Gradient-Update Mismatch (GUM). Accordingly, we propose Gradient-Update Alignment (GUA), which projects $u_t$ onto $\mathcal{C}_t$ to obtain the aligned update $p_t$ and applies $p_t$ to the parameters. When the optimizer maintains internal state, GUA further adjusts this state toward targets reconstructed from the applied update. We conduct extensive experiments and find that GUM is widespread across momentum, adaptive, and curvature-based optimizers, with conflict rates reaching up to 86.3%. Across all PINN settings, GUA achieves conflict-free applied updates and consistently improves various gradient surgery methods, reducing the relative $L_2$ error by up to 98.2% in individual settings. Data and code are available at https://github.com/JingXiao10/GUA.

1 INTRODUCTION

PINN training is difficult because loss terms can produce substantially different gradient magnitudes and directions, while gradient surgery may not remain conflict-free after optimizer transformation. The paper formalizes this mismatch and proposes GUA to align the optimizer proposal before applying it.

  • PINNs combine PDE-residual and initial/boundary-condition losses whose gradients can differ substantially, creating conflicts that hinder training.
  • Loss balancing addresses gradient-magnitude imbalance but does not directly resolve directional conflicts, motivating gradient surgery methods.
  • Gradient surgery constructs a conflict-reduced direction before passing it to the optimizer, raising whether conflict-freeness survives optimizer transformation.
  • GUM occurs when a conflict-free constructed direction becomes a conflicting optimizer proposal through state, scaling, preconditioning, or weight-decay transformations.
  • GUA aligns the optimizer proposal with the conflict-free cone, applies the aligned update, and can adjust optimizer state toward targets reconstructed from that update.
  • Experiments find GUM widespread across momentum, adaptive, and curvature-based optimizers, while GUA achieves conflict-free applied updates and consistently improves gradient surgery methods.

2 RELATED WORK

Prior PINN work addresses training difficulty through loss balancing, while related multi-task methods use gradient surgery to reduce conflicts among loss-specific gradients. The paper focuses on optimizer transformations that occur after this surgery step.

  • PINNs are difficult to train because loss terms can have imbalanced magnitudes, poor conditioning, and different convergence rates.
  • Loss-balancing methods manually or adaptively reweight physical constraints but do not directly address conflicts in gradient directions.
  • Gradient surgery methods modify loss-specific gradients to reduce interference, with PCGrad, CAGrad, and IMTL-G using different conflict-handling principles.
  • Modern optimizers transform surgery-produced directions using mechanisms including historical state, adaptive scaling, preconditioning, and decoupled weight decay.

3 GRADIENT–UPDATE MISMATCH AND GRADIENT–UPDATE ALIGNMENT

The paper defines conflict-free geometry and GUM, then analyzes when optimizer transformations preserve that geometry. GUA projects optimizer proposals into the conflict-free cone and optionally aligns optimizer state with the applied update.

  • 3.1 PRELIMINARIES: MULTI-LOSS OPTIMIZATION: The conflict-free cone contains directions whose inner products with every loss-specific gradient are nonnegative, so no loss increases to first order.
  • 3.1 PRELIMINARIES: MULTI-LOSS OPTIMIZATION: Gradient surgery produces a direction that an optimizer transforms, together with internal state, into the proposal determining the actual parameter-update direction.
  • 3.2 GRADIENT–UPDATE MISMATCH: GUM is the event that the constructed direction lies in the conflict-free cone while the optimizer proposal lies outside it.
  • 3.2 GRADIENT–UPDATE MISMATCH: Generalized GUM uses Ru > Ra when optimizer transformation produces a net increase in the frequency of conflicting proposals.
  • 3.3 CONFLICT PRESERVATION UNDER OPTIMIZER TRANSFORMATIONS: Fixed-state affine transformations preserve conflict-freeness only under exact conditions; positive scalar scaling preserves the cone, whereas coordinate-wise scaling need not.
  • 3.3 CONFLICT PRESERVATION UNDER OPTIMIZER TRANSFORMATIONS: Historical state, adaptive scaling, preconditioning, and weight decay can move a conflict-free direction outside the conflict-free cone.
  • 3.4 GRADIENT–UPDATE ALIGNMENT: GUA projects the optimizer proposal onto the cone, applies the resulting update, and softly aligns optimizer state when projection changes the proposal.
  • 3.4 GRADIENT–UPDATE ALIGNMENT: The aligned update is the closest conflict-free update under the chosen metric and has smaller first-order loss-change violations than a conflicting proposal for sufficiently small steps.

4 EXPERIMENTS

Experiments test whether optimizer transformations reintroduce conflicts after gradient surgery and whether GUA improves accuracy, alignment, state handling, efficiency, and task scaling across PINN settings.

  • Q1: Optimizer-Induced GUM: GUM is widespread: M-SGD produces conflict rates Ru of 51.3%–86.3%, whereas plain SGD preserves Ru = 0 after ConFIG.Nonzero Ru after ConFIG's Ra = 0 isolates conflicts introduced by optimizer transformation.
  • Q2: PINN Performance: GUA eliminates update-level conflicts, yielding Rp = 0 across all PINN settings.It improves accuracy across gradient surgery methods and PDE benchmarks, with relative L2 error reductions from 11.6% to 98.2%.
  • Q3: Direction Alignment Mechanism: Direction alignment drives the benefit: p_t has lower worst relative one-step loss change than u_t in 95.1% and than norm-matched q_t in 93.4% of comparisons.The norm-matched control preserves u_t's direction while matching p_t's norm, isolating direction alignment.
  • Q4: State Alignment Mechanism: Soft state alignment improves performance beyond direction alignment in most settings, while exact alignment is less effective and can be unstable.Aligning both Adam moments exactly degrades performance in some settings and is numerically unstable in one Heat-MS setting.
  • Q5: GUA Overhead: GUA improves accuracy with additional training-time overhead but little peak GPU memory overhead across PINN settings.The projection avoids forming a p × p matrix and has parameter-dimensional Gram-matrix cost O(m2p).
  • Q6: Task-Cardinality Scaling: GUA remains effective from 2 to 40 jointly optimized tasks, improving F1 in six of seven settings with only 4.2% additional memory cost at 40 tasks.Training-time overhead becomes more pronounced as task counts grow, while performance remains comparable at 20 tasks.

5 CONCLUSION

The paper identifies Gradient–Update Mismatch as a failure of conflict-freeness after optimizer transformation and proposes GUA to enforce conflict-free geometry at the applied-update level.

  • 5 CONCLUSION: GUA projects the optimizer proposal onto the current conflict-free cone before applying it to the parameters.The resulting aligned update is conflict-free by construction.
  • 5 CONCLUSION: Affine conflict preservation depends on the transformation’s linear and additive components, while the nonlinear characterization uses ray-wise behavior and Jacobians.The locally Lipschitz formulation also covers piecewise-smooth transformations almost everywhere.
  • 5 CONCLUSION: GUM occurs when a gradient-surgery direction in the conflict-free cone is transformed by the optimizer into a proposal outside that cone.The paper characterizes both affine and nonlinear optimizer transformations that preserve conflict-freeness.
  • 5 CONCLUSION: The affine result is a specialization of the general nonlinear characterization for optimizer transformations.The linear map and offset recover the affine conditions from the broader criterion.

A.3 PROOF OF PROPOSITION 2

The proposition establishes that GUA’s projection provides a least-change feasible correction and removes positive first-order loss-change terms for conflicting optimizer proposals.

  • A.3 PROOF OF PROPOSITION 2: For conflicting proposals, the aligned update has no positive first-order loss-change term.This follows from membership in the conflict-free cone and the smoothness bounds.
  • A.3 PROOF OF PROPOSITION 2: The strict one-step improvement inequality holds for all sufficiently small admissible step sizes when both trial points remain in the assumed neighborhood.The proposition assumes smooth losses and neighborhood containment for the two trial points.
  • A.3 PROOF OF PROPOSITION 2: GUA’s metric projection produces the closest conflict-free update to the optimizer proposal under the chosen metric.The aligned update is closer than the proposal to every feasible direction, including the constructed direction when feasible.
  • A.3 PROOF OF PROPOSITION 2: If the aligned update has strictly positive projections onto every loss gradient, all losses decrease simultaneously for sufficiently small steps.On the cone boundary, the guarantee is limited to eliminating positive first-order loss-change terms for corresponding losses.
  • A.3 PROOF OF PROPOSITION 2: Norm matching cannot restore conflict-freeness because positive rescaling preserves the direction’s conflict status.Changing the update direction is necessary even when the control and aligned update have the same Euclidean norm.

B.1 PDE BENCHMARKS

The evaluation spans six PDE benchmarks with varied equations, dimensions, time dependence, and flow structure, using shared fully connected PINN architectures across methods.

  • B.1 PDE BENCHMARKS: The benchmark suite covers steady and time-dependent PDEs, scalar and multi-component systems, low- and high-dimensional domains, and incompressible flows.Protocols primarily follow PINNacle and ConFIG.
  • B.1 PDE BENCHMARKS: The suite includes Heat-MS, Schrödinger, Poisson-5D, Burgers, and Kovasznay flow among the specified PDE problems.The passages describe their domains, equations, or boundary and initial conditions.
  • B.1 PDE BENCHMARKS: Beltrami flow adds an unsteady three-dimensional incompressible-flow benchmark with velocity, pressure, and divergence equations.Its default setting uses Re = 1 and ν = 1.
  • B.1 PDE BENCHMARKS: All PINN models use fully connected MLPs with a shared trunk, coupling residual, boundary, and initial-condition losses through one parameter vector.Output channels depend on the PDE state and may represent scalar, complex, multi-species, or flow variables.
  • B.1 PDE BENCHMARKS: Architectures are fixed across baseline methods and their +GUA variants, with matched Xavier initialization for each random seed.This isolates GUA’s effect from architectural differences.

B.3 TRAINING SETTINGS

The experiments standardize training, optimizer, evaluation, and conflict diagnostics while comparing multiple gradient-surgery methods and GUA variants under controlled settings.

  • B.3 TRAINING SETTINGS: Adam is the default optimizer with β1 = 0.9, β2 = 0.999, ϵ = 10^-8, unit loss weights, and full-batch training.Training points are generally resampled each epoch using Latin-hypercube sampling.
  • B.3 TRAINING SETTINGS: The default learning-rate schedule uses η0 = 10^-3 and ηmin = 10^-4, with a Heat-MS 3-loss exception to avoid ConFIG instability.That exception sets both learning-rate endpoints to 10^-4.
  • B.3 TRAINING SETTINGS: The optimizer-induced GUM experiments fix ConFIG and vary the optimizer, measuring how the constructed direction becomes the optimizer proposal.No GUA is applied in those experiments.
  • B.3 TRAINING SETTINGS: The PINN comparison evaluates PCGrad, CAGrad, IMTL-G, A-MTL, UPGrad, and ConFIG before and after GUA.All methods first construct a direction from loss-specific gradients and pass it to the same base optimizer.
  • B.3 TRAINING SETTINGS: Conflict diagnostics report the sequence Rg → Ra → Ru → Rp alongside relative L2 error.The stages measure raw gradients, constructed direction, optimizer proposal, and aligned update; Rp is unavailable for baselines.
  • B.3 TRAINING SETTINGS: Algorithm 1 computes loss-specific gradients, constructs a direction, obtains the optimizer proposal, projects it, applies the aligned update, and optionally aligns optimizer state.The update is applied as θt+1 ← θt − ηtpt.

C.1 PROJECTION METRIC ABLATION

The projection metric determines which conflict-free update is selected and affects numerical robustness, not merely feasibility. Across evaluated settings, the Adam metric is at least competitive in accuracy and more stable on difficult cases.

  • The Adam metric provides accuracy comparable to or better than the Euclidean metric across evaluated PINN settings.
  • On Heat-MS, Euclidean projection is partially unstable with two losses and entirely non-finite with three, while Adam projection remains stable in both settings.
  • Both metrics enforce the same conflict-free constraints but select different feasible updates because they measure deviation from the optimizer proposal in different geometries.
  • The Adam metric accounts for optimizer-induced coordinate anisotropy through second-moment statistics, whereas the Euclidean metric treats coordinates uniformly.
  • The Adam metric is used by default because it is accuracy-competitive and offers better numerical robustness in difficult settings.

C.3 OPTIMIZER-INDUCED GUM RESULTS ON KOVASZNAY, BELTRAMI, AND POISSON-5D

Additional PINN benchmarks show that optimizer-induced GUM persists beyond the main-text subset. Across the evaluated configurations, gradient surgery can produce conflict-free directions that become conflicting optimizer proposals, while GUA is evaluated through complete error and conflict-rate summaries.

  • ConFIG yields Ra = 0, while momentum, adaptive, and curvature-based optimizers can produce nonzero post-optimizer conflict rates Ru.
  • Table 7 tracks conflict rates through raw gradients, gradient surgery, and optimizer transformation using the sequence Rg →Ra →Ru.
  • The complete Q2 results report relative L2 errors and conflict-rate transitions for each evaluated method–benchmark configuration.

C.4.1 QUANTITATIVE PINN RESULTS FOR Q2

GUA improves quantitative PINN performance across evaluated base methods and settings while addressing optimizer-induced conflict after gradient surgery. The experiments compare complete relative-error results, conflict-rate transitions, optimizer families, and update-level behavior across PDE benchmarks.

  • Adding GUA reduces the mean relative L2 error across all evaluated base methods and PINN settings.
  • Adam+GUA improves over Adam in all ten settings without pre-optimizer gradient surgery.
  • PCGrad+GUA improves relative L2 error by 34.1%, 85.5%, 22.1%, 20.5%, and 42.5% across the reported settings.
  • IMTL-G+GUA achieves improvements of 98.2%, 67.4%, 23.8%, 51.5%, and 70.9% across the reported settings.
  • UPGrad+GUA improves relative L2 error by 22.4%, 85.6%, 23.6%, 17.3%, and 40.6% across the reported settings.
  • GUA projects the optimizer proposal onto the conflict-free cone under an optimizer-dependent metric and can align optimizer state toward targets reconstructed from the applied update.
  • GUA improves displayed Burgers and Schrödinger settings across RMSProp, AdamW, and SOAP, including SOAP despite its alignment-oriented preconditioning.
  • Figure 11 visualizes projected two-loss training-step directions after rotating each directional slice into a basis defined by the loss-specific gradients.

D MULTI-TASK LEARNING DETAILS

The multi-task benchmark uses CelebA with nested attribute subsets and applies GUA only to shared parameters after task-gradient aggregation and optimizer proposal. Task-specific heads remain updated by the native optimizer.

  • CelebA experiments vary task cardinality using the first m attributes in official annotation order, producing nested task subsets.
  • Images are resized to 64 × 64, converted to tensors, and used without additional data augmentation.
  • Each attribute is trained independently with binary cross-entropy.
  • GUA modifies shared parameters after task-gradient aggregation and optimizer proposal, while task-specific classification heads use the native optimizer.
Loading 2609.01558v1…