Variance Balancing in Bayesian PINNs

Updated 10 April 2026

Variance balancing in Bayesian PINNs is an uncertainty-aware method that dynamically weights loss terms using learned or predicted variances.
It employs techniques like heteroscedastic likelihoods, empirical Bayes, repulsive ensembles, and error-bound models to automate loss weighting and quantify prediction uncertainty.
These methods accelerate convergence, optimize model selection, and improve predictive accuracy and generalization in solving PDEs.

Variance balancing in Bayesian Physics-Informed Neural Networks (PINNs) refers to a family of methods that dynamically weight the various loss contributions in PINN training, using either predicted or learned variances, such that data, physics, boundary, and other losses are optimally harmonized. This approach produces principled automatic trade-offs, avoids manual tuning of loss weights, quantifies uncertainty (aleatoric and/or epistemic), and typically accelerates convergence and improves generalization. Current methods derive from diverse Bayesian frameworks: heteroscedastic likelihoods, empirical Bayes, Bayesian ensembles with repulsion, Laplace or variational approximations, and the use of rigorous PINN error bounds to learn or prescribe loss variances.

1. Heteroscedastic Likelihood Formulation and Loss Design

Variance balancing in modern PINNs is commonly achieved by formulating the data/model error as a (potentially input-dependent) Gaussian likelihood. Instead of a scalar output $u_\theta(x)$ , the PINN produces both a predictive mean $\mu_\theta(x)$ and a log-variance head $s_\theta(x)$ , enforcing positivity with a softplus transformation: $\sigma^2_\theta(x) = \log(1 + e^{s_\theta(x)}) + \epsilon, \quad \epsilon=10^{-6}.$ For a target value $y$ at location $x$ , the negative log-likelihood (NLL) per point is

$\mathrm{NLL}(x,y) = \ln(\sigma^2_\theta(x) + 1) + \frac{(y - \mu_\theta(x))^2}{\sigma^2_\theta(x)},$

where the denominator adaptively down-weights large-residual regions by increasing local variance. This heteroscedastic NLL replaces or augments vanilla PINN quadratic terms in each loss component (initial, boundary, PDE residual). In the VI-PINN method, the MNLL is used both as a direct loss and as an auxiliary task, resulting in a total loss $\mathcal L_\mathrm{total} = \mathcal L_\mathrm{PINN} + \lambda \mathcal L'_\mathrm{MNLL}$ with $\lambda > 0$ (Shan et al., 2022).

This structure enables the network not only to fit the mean solution but also to generate internal uncertainty maps reflecting local model and data fit, which simultaneously serve as dynamic loss weights.

2. Bayesian Model Selection and Hyperparameter Control

Within a full Bayesian PINN, variances enter as hyperparameters of the Gaussian likelihoods associated with each loss term—typically one for data/boundary and one for physics/PDE residuals. By maximizing the model evidence via empirical Bayes or Laplace canonicization, optimal $\gamma$ (variance) parameters are determined: $\mu_\theta(x)$ 0 with $\mu_\theta(x)$ 1 and $\mu_\theta(x)$ 2 the data/boundary and PDE residuals, respectively. Maximizing marginal likelihood yields closed-form updates: $\mu_\theta(x)$ 3 Here $\mu_\theta(x)$ 4 is the loss Hessian at the MAP and $\mu_\theta(x)$ 5, $\mu_\theta(x)$ 6 are curvature contributions from data and physics terms. The system automatically balances the effective loss weights, upweighting terms with low residuals and/or higher information content, and vice versa; this obviates hand-tuning of loss weights (Graczyk et al., 2023). Complexity terms in the evidence penalize overfitting and deliver a model selection (Occam’s razor) principle.

3. Variance-Driven Ensemble and Replica Dynamics

Variance collapse is a major obstacle in ensemble and MCMC-based Bayesian PINNs: standard ensembles tend to converge to maximum a posteriori (MAP) solutions with negligible empirical variance, thereby underestimating uncertainty. The Repulsive Ensemble PINN (RE-PINN) injects a functional-space repulsive term via kernel density estimation: $\mu_\theta(x)$ 7 where each particle’s loss includes $\mu_\theta(x)$ 8. The repulsive gradient, $\mu_\theta(x)$ 9, disperses ensemble members in function and parameter space, sustaining nonzero variance and enabling empirical Bayes consistency as ensemble size grows. This enables credible uncertainty quantification and prevents ensemble collapse (Pilar et al., 22 May 2025).

Similarly, multi-variance replica exchange SGLD (MVRx-SGLD) exploits chains at different “temperatures” (i.e., different forward model fidelities/variances), balancing exploration and exploitation. An unbiased swap-rate corrects for variance differences between chains, accelerating convergence and improving posterior coverage (Lin et al., 2021).

4. Pseudo-Label and Selective Loss Inclusion via Posterior Variance

Variance balancing also appears as a mechanism for controlling which model points contribute to the physics loss at any given training step. Bayesian Pseudo Label PINN (B-PL-PINN) uses posterior variance estimated from MCMC samples to define a trust-region over the input domain: new collocation points are added as pseudo-labeled only when their predicted variance is below a threshold. The PDE residual loss is enforced only on points within proximity to labeled regions and where uncertainty is low, adaptively expanding the physics-loss support as confidence grows. This selective enforcement sharply improves convergence and propagation of information into regions of the domain traditionally susceptible to “physics loss dominance” or instability (Innerebner et al., 2 Jul 2025).

5. Heteroscedastic and Aleatoric Uncertainty Based on PINN Error Bounds

A complementary approach employs rigorous PINN error bounds to construct aleatoric variance estimates for uncertainty quantification. Instead of learning variance, heteroscedastic error bounds computed from the residuals and known stability estimates are promoted to observation noise models in a two-step Bayesian workflow: $s_\theta(x)$ 0 where $s_\theta(x)$ 1 is a rigorously derived upper bound on the solution error at $s_\theta(x)$ 2. This input-dependent variance is injected into the Bayesian neural net’s likelihood, yielding pointwise calibrated uncertainty, and is used both in forward and in (cosmological) inverse problems (Flores et al., 9 May 2025). Empirically, this substantially reduces miscalibration area and improves sharpness of the predictive posterior versus homoscedastic or purely residual-based baselines.

Method	How Variance is Used	Loss Structure
VI-PINN (Shan et al., 2022)	Learned, heteroscedastic head	MNLL in loss, aux task
Bayesian-PINN (Graczyk et al., 2023)	Evidence-maximizing variance hyperparameters	Data + PDE NLL, Laplace approx
RE-PINN (Pilar et al., 22 May 2025)	Ensemble variance, repulsive kernel	Repulsive ensemble loss
B-PL-PINN (Innerebner et al., 2 Jul 2025)	Posterior variance for pseudo-label	Selective PDE loss
Error-bound BNN (Flores et al., 9 May 2025)	Analytic error bounds	Heteroscedastic BNN

6. Empirical Effects: Convergence, Calibration, and Generalization

Variance balancing consistently accelerates convergence, sharpens generalization, and produces calibrated uncertainties. VI-PINNs achieve up to 10× lower $s_\theta(x)$ 3 error than baseline PINNs in classic PDE (advection, Burgers, wave, Poisson, C-Diff) benchmarks (Shan et al., 2022). Empirical Bayes balancing via evidence maximization obviates the need for hand-tuned weights and always finds the Pareto-optimal assignment of loss importances (Graczyk et al., 2023). Repulsive ensembles close the gap with fully Bayesian posteriors (Wasserstein distance to true posterior shrinks with ensemble size), retaining nonzero uncertainty bands even in data-rich regions (Pilar et al., 22 May 2025). Error-bound driven BNNs decrease miscalibration area by up to 80% compared to vanilla PINN predictors (Flores et al., 9 May 2025). Sensitivity studies demonstrate optimal CRPS at moderate collocation strength, and confirm that overly strict PDE imposition suppresses useful variance (Ramirez et al., 7 Jan 2026).

7. Practical Algorithmic Implementation

Variance balancing methods are compatible with standard PINN workflows but require architectural or optimization-layer modifications. In VI-PINNs, the mean and variance heads are trained jointly, with softplus transformations and cap constants to regularize $s_\theta(x)$ 4 terms. Bayesian PINNs require iterative optimization over both model parameters and variance hyperparameters (either via closed-form evidence maximization or gradient-based routines); the Hessian trace may be estimated via low-rank or stochastic schemes. Repulsive ensembles necessitate kernel density estimation in high-dimensional function space; KDE bandwidths are set by median heuristics and repulsion is typically activated after initial physics convergence. Pseudo-labeling strategies entail a nested, interleaved MCMC and model expansion schedule, with local variance used as a gating function. For error-bound BNNs, error bounds must be computed analytically or numerically, and may be clipped/floored for stability.

Empirical recommendations emphasize: initial training with deterministic PINN, gradual activation of variance or ensemble-based mechanisms, testing sharpness/calibration metrics, and sensitivity analyses on the relative size of data vs. physics loss populations.

In summary, variance balancing in Bayesian PINNs subsumes a broad class of uncertainty-aware, adaptively weighted architectures and training strategies. Rigorous formulation and empirical validation confirm its central role in reliable, efficient, and interpretable PDE solution by neural networks. For implementation, approaches such as VI-PINN, empirical Bayes PINN, RE-PINN, B-PL-PINN, and error-bound-driven BNNs exemplify the state of the art (Shan et al., 2022, Graczyk et al., 2023, Pilar et al., 22 May 2025, Flores et al., 9 May 2025, Innerebner et al., 2 Jul 2025, Ramirez et al., 7 Jan 2026).