AutoBalance: An Automatic Balancing Framework for Training Physics-Informed Neural Networks

TL;DR

This work addresses training PINNs under composite losses with heterogeneous curvatures by introducing AutoBalance, a post-combine framework that assigns independent adaptive optimizers to each loss term and aggregates their updates. Theoretical analysis on a curvature-imbalanced quadratic model and empirical Hessian studies show improved conditioning and stability, while experiments across 1D/2D PDE benchmarks demonstrate robust, superior performance over existing balancing methods. AutoBalance is shown to be orthogonal to various PINN variants, enhancing their accuracy without requiring new balancing hyperparameters. The approach offers a practical, generalizable tool for multi-task optimization in physics-informed learning, with potential extensions to broader multi-task contexts and efficiency-focused refinements.

Abstract

Physics-Informed Neural Networks (PINNs) provide a powerful and general framework for solving Partial Differential Equations (PDEs) by embedding physical laws into loss functions. However, training PINNs is notoriously difficult due to the need to balance multiple loss terms, such as PDE residuals and boundary conditions, which often have conflicting objectives and vastly different curvatures. Existing methods address this issue by manipulating gradients before optimization (a "pre-combine" strategy). We argue that this approach is fundamentally limited, as forcing a single optimizer to process gradients from spectrally heterogeneous loss landscapes disrupts its internal preconditioning. In this work, we introduce AutoBalance, a novel "post-combine" training paradigm. AutoBalance assigns an independent adaptive optimizer to each loss component and aggregates the resulting preconditioned updates afterwards. Extensive experiments on challenging PDE benchmarks show that AutoBalance consistently outperforms existing frameworks, achieving significant reductions in solution error, as measured by both the MSE and $L^{\infty}$ norms. Moreover, AutoBalance is orthogonal to and complementary with other popular PINN methodologies, amplifying their effectiveness on demanding benchmarks.

Figures (7)

• Figure 1: Empirical analysis of Hessian properties for the 2D Helmholtz equation. (Left) At initialization, the Hessian spectra for the residual and boundary losses exhibit strong heterogeneity. (Middle) After 30,000 epochs, the AutoBalance-preconditioned Hessian for the residual loss shows a larger minimum eigenvalue compared to both the original and Adam-preconditioned Hessians. (Right) During training, AutoBalance consistently maintains a substantially lower effective condition number for the boundary loss than standard Adam or the original Hessian.
• Figure 2: AutoBalance exhibits an emergent auto-balancing property on the 2D Helmholtz problem.(Left) The norm ratio of the raw gradients (interior/boundary) is highly imbalanced, whereas the ratio of the update vectors from AutoBalance consistently remains near the ideal balance of 1.0. (Right) The raw gradients become increasingly anti-aligned (negative cosine similarity), indicating task conflict. In contrast, the update vectors from AutoBalance maintain a positive alignment, ensuring constructive updates.
• Figure 3: MSE history of AutoBalance and baseline methods for three PDE benchmarks: 1D reaction-diffusion system (Left), 2D Helmholtz equation (Middle), and 2D Poisson inverse problem (Right).
• Figure 4: Heatmap of the 1D reaction-diffusion system. Top Row: Exact solution. Other Rows: Predicted solutions by Auto-AdamW ( Second Row), DWA ( Third Row), and IMTL-G ( Bottom Row), along with their corresponding point-wise errors.
• Figure 5: Heatmap of the 2D Helmholtz equation. Top Row: Exact solution. Other Rows: Predicted solutions by Auto-AdamW ( Second Row), DWA ( Third Row), and IMTL-G ( Bottom Row), along with their corresponding point-wise errors.

Theorems & Definitions (5)

• Definition 1
• Theorem 1: Convergence of Iterates in Euclidean Norm
• Corollary 1.1
• proof
• proof