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

arXiv: https://arxiv.org/abs/2510.06684 · HTML: https://mlsanity.com/paper/2510.06684/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 22 of its 40 references that are in the map.
- First posted on arXiv 8 Oct 2025.

## Builds on (most cited first, 22 of 22)

- [Gradient Surgery for Multi-Task Learning](https://mlsanity.com/paper/2001.06782/) (2020)
- [GradNorm: Gradient Normalization for Adaptive Loss Balancing in Deep Multitask Networks](https://mlsanity.com/paper/1711.02257/) (2017)
- [Multi-Task Learning as Multi-Objective Optimization](https://mlsanity.com/paper/1810.04650/) (2018)
- [When and why PINNs fail to train: A neural tangent kernel perspective](https://mlsanity.com/paper/2007.14527/) (2020)
- [End-to-End Multi-Task Learning with Attention](https://mlsanity.com/paper/1803.10704/) (2018)
- [Conflict-Averse Gradient Descent for Multi-task Learning](https://mlsanity.com/paper/2110.14048/) (2021)
- [B-PINNs: Bayesian Physics-Informed Neural Networks for Forward and Inverse PDE Problems with Noisy Data](https://mlsanity.com/paper/2003.06097/) (2020)
- [Identifying and attacking the saddle point problem in high-dimensional non-convex optimization](https://mlsanity.com/paper/1406.2572/) (2014)
- [Deep Learning without Poor Local Minima](https://mlsanity.com/paper/1605.07110/) (2016)
- [Quantifying total uncertainty in physics-informed neural networks for solving forward and inverse stochastic problems](https://mlsanity.com/paper/1809.08327/) (2018)
- [Adversarial Uncertainty Quantification in Physics-Informed Neural Networks](https://mlsanity.com/paper/1811.04026/) (2018)
- [Residual-based attention in physics-informed neural networks](https://mlsanity.com/paper/2307.00379/) (2023)
- [Learning in Modal Space: Solving Time-Dependent Stochastic PDEs Using Physics-Informed Neural Networks](https://mlsanity.com/paper/1905.01205/) (2019)
- [Solving Inverse Stochastic Problems from Discrete Particle Observations Using the Fokker-Planck Equation and Physics-informed Neural Networks](https://mlsanity.com/paper/2008.10653/) (2020)
- [ConFIG: Towards Conflict-free Training of Physics Informed Neural Networks](https://mlsanity.com/paper/2408.11104/) (2024)
- [Multi-Objective Meta Learning](https://mlsanity.com/paper/2102.07121/) (2021)
- [RoPINN: Region Optimized Physics-Informed Neural Networks](https://mlsanity.com/paper/2405.14369/) (2024)
- [The loss landscape of deep linear neural networks: a second-order analysis](https://mlsanity.com/paper/2107.13289/) (2021)
- [Mitigating Gradient Bias in Multi-objective Learning: A Provably Convergent Stochastic Approach](https://mlsanity.com/paper/2210.12624/) (2022)
- [Initialization-enhanced Physics-Informed Neural Network with Domain Decomposition (IDPINN)](https://mlsanity.com/paper/2406.03172/) (2024)
- [Complex Physics-Informed Neural Network](https://mlsanity.com/paper/2502.04917/) (2025)
- [Convolution-weighting method for the physics-informed neural network: A Primal-Dual Optimization Perspective](https://mlsanity.com/paper/2506.19805/) (2025)

## Abstract (arXiv)

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∞ norms. Moreover, AutoBalance is orthogonal to and complementary with other popular PINN methodologies, amplifying their effectiveness on demanding benchmarks.