Dual-Balancing for Physics-Informed Neural Networks

TL;DR

Physics-informed neural networks (PINNs) solve PDEs by minimizing a composite loss over the PDE residual and boundary/initial constraints, but gradient imbalance often hinders convergence. The authors introduce Dual-Balanced PINN (DB-PINN), which combines inter-balancing (gradient-distribution alignment between the PDE residual and condition-fitting losses) with intra-balancing (allocation of weight according to the difficulty of fitting each condition) and a Welford-based online weight update to stabilize training. DB-PINN demonstrates faster convergence and lower prediction errors on Klein-Gordon, Wave, and Helmholtz benchmarks, outperforming gradient-statistics-based weighting methods across multiple statistics (mean, std, kurtosis). The approach reduces weight spikes, eliminates hyperparameter tuning for updates, and enhances robustness for multi-physics and stiff PDE settings.

Abstract

Physics-informed neural networks (PINNs) have emerged as a new learning paradigm for solving partial differential equations (PDEs) by enforcing the constraints of physical equations, boundary conditions (BCs), and initial conditions (ICs) into the loss function. Despite their successes, vanilla PINNs still suffer from poor accuracy and slow convergence due to the intractable multi-objective optimization issue. In this paper, we propose a novel Dual-Balanced PINN (DB-PINN), which dynamically adjusts loss weights by integrating inter-balancing and intra-balancing to alleviate two imbalance issues in PINNs. Inter-balancing aims to mitigate the gradient imbalance between PDE residual loss and condition-fitting losses by determining an aggregated weight that offsets their gradient distribution discrepancies. Intra-balancing acts on condition-fitting losses to tackle the imbalance in fitting difficulty across diverse conditions. By evaluating the fitting difficulty based on the loss records, intra-balancing can allocate the aggregated weight proportionally to each condition loss according to its fitting difficulty level. We further introduce a robust weight update strategy to prevent abrupt spikes and arithmetic overflow in instantaneous weight values caused by large loss variances, enabling smooth weight updating and stable training. Extensive experiments demonstrate that DB-PINN achieves significantly superior performance than those popular gradient-based weighting methods in terms of convergence speed and prediction accuracy. Our code and supplementary material are available at https://github.com/chenhong-zhou/DualBalanced-PINNs.

Figures (5)

• Figure 1: Illustration of the proposed DB-PINN. The flow of different losses is represented via different colors (PDE residual loss: blue; IC loss: purple; BC loss: green). The yellow dashed box signifies inter-balancing, where the gradients ($\nabla_\theta$) of each loss are first calculated, and then gradient distribution discrepancies are evaluated by the ratio ($/$) of $\nabla_\theta \mathcal{L}^r$ to $\nabla_\theta \mathcal{L}^i$ under certain gradient statistical metrics. Total gradient ratio $\mathcal{G}$, also considered as the aggregated weight, is computed by a summation ($\sum$). The pink dashed box signifies intra-balancing, where BC and IC losses are used to calculate their difficulty indexes (i.e., $\mathcal{I}^{bc}$ and $\mathcal{I}^{ic}$). The gray arrow implies that we sort condition losses by fitting difficulty, and finally, the aggregated weight can be proportionally allocated to BC and IC losses to derive their respective weights.
• Figure 2: Loss and weight curves during the training on solving the Klein-Gordon equation.
• Figure 3: The Klein-Gordon equation. (a-b) Point-wise absolute errors. (c) Mean prediction error curves. (d-f) Comparison between reference solutions (blue) and predictions (red).
• Figure 4: The Wave equation. (a-b) Point-wise absolute errors. (c) Mean prediction error curves. (d-f) Comparison between reference solutions (blue) and predictions (red).
• Figure 5: The Helmholtz equation. (a-b) Point-wise absolute errors. (c) Mean prediction error curves. (d-f) Comparison between reference solutions (blue) and predictions (red).