UQ-VAE: Uncertainty Quantification in VAEs

• UQ-VAE is a neural variational inference framework that estimates posterior means and covariances for inverse problems while delivering explicit uncertainty quantification.
• It employs a Jensen–Shannon divergence-based objective with a tunable hyperparameter to adaptively balance data fit and uncertainty estimation.
• Extensive benchmarks demonstrate UQ-VAE’s scalable performance and reliable uncertainty assessment in high-dimensional Bayesian inverse problems and reduced-order modeling.

Uncertainty Quantification Variational Autoencoders (UQ-VAE) define a family of neural variational inference methods that provide data-driven estimation of posterior means and covariances for unknown parameters in inverse problems, while simultaneously delivering explicit, instance-wise uncertainty quantification. UQ-VAEs generalize the conventional variational autoencoder formulation beyond the classical ELBO, adapting objective functions for faithful posterior approximation, improved control of epistemic uncertainty, and robust performance in high-dimensional or computationally constrained settings. The methodology has been applied for Bayesian inverse problems and uncertainty-aware reduced-order modeling, with documented theoretical guarantees and extensive empirical validation across physical sciences and engineering (Goh et al., 2019, Tonini et al., 14 Sep 2025, Zighed et al., 29 Mar 2025).

1. Mathematical Framework and Bayesian Inverse Problem Formulation

UQ-VAE operates within the Bayesian inverse problem paradigm. Let $\theta \in \mathbb{R}^D$ denote unknown model parameters with prior $p(\theta)$. Given forward map $\mathcal{F}:\mathbb{R}^D \rightarrow \mathbb{R}^O$ and observational noise $E$, the observed data satisfy

$Y = \mathcal{F}(\theta) + E\,,$

with $p(E)$ typically Gaussian. The likelihood is $p(y | \theta) = p_E(y - \mathcal{F}(\theta))$. The posterior density is then

$p(\theta|y) \propto p(y|\theta)\,p(\theta)\,.$

The goal is to approximate $p(\theta|y)$ efficiently, providing access to both the posterior mean and covariance, with scalable computation for high dimensions and expensive forward models (Goh et al., 2019, Tonini et al., 14 Sep 2025).

2. UQ-VAE Objective and Divergence-Based Variational Inference

Unlike the standard VAE, which maximizes an ELBO based on the zero-avoiding KL divergence, UQ-VAE introduces a Jensen–Shannon divergence (JSD) family objective:

$\JS_\lambda(q\|p) = \lambda\, \KL(q \| \lambda q + (1-\lambda)p) + (1-\lambda)\, \KL(p \| \lambda q + (1-\lambda)p)$

with $p(\theta)$0, interpolating between $p(\theta)$1 (ELBO, as $p(\theta)$2) and $p(\theta)$3 (as $p(\theta)$4). The UQ-VAE training objective for a variational posterior $p(\theta)$5 is (up to constants):

$p(\theta)$6

where the three terms correspond to posterior-data fit, likelihood misfit, and a prior regularizer (Goh et al., 2019).

A main innovation is the role of the hyperparameter $p(\theta)$7, which regulates the balance between mean-fitting and variance estimation. Small $p(\theta)$8 promotes greater posterior variance (zero-forcing), while larger $p(\theta)$9 tightens the posterior (zero-avoiding). This adaptive weighting allows the encoder to automatically learn uncertainty magnitudes that reflect both data fit and prior information.

3. UQ-VAE Network Architectures and Training Procedure

The encoder $\mathcal{F}:\mathbb{R}^D \rightarrow \mathbb{R}^O$0 is a neural network mapping $\mathcal{F}:\mathbb{R}^D \rightarrow \mathbb{R}^O$1 to $\mathcal{F}:\mathbb{R}^D \rightarrow \mathbb{R}^O$2; $\mathcal{F}:\mathbb{R}^D \rightarrow \mathbb{R}^O$3 is typically diagonal for scalability. Encoder architectures may use 5 hidden layers with 500 units per layer (ReLU activation), outputting both mean and log-variance parameters (Goh et al., 2019).

The decoder corresponds either to the actual forward map $\mathcal{F}:\mathbb{R}^D \rightarrow \mathbb{R}^O$4 (if available) or to a learned surrogate neural network $\mathcal{F}:\mathbb{R}^D \rightarrow \mathbb{R}^O$5, typically with 2 hidden layers (500 units) (Goh et al., 2019).

A typical training loop processes batches of $\mathcal{F}:\mathbb{R}^D \rightarrow \mathbb{R}^O$6, performing a forward pass for mean and variance inference, reparameterized sampling (the standard VAE “reparameterization trick”), evaluation of loss terms (including the divergence-weighted negative log-posterior and KL), and backpropagation with Adam optimizer.

Subsequent extensions, such as those in (Tonini et al., 14 Sep 2025), introduce architectural and algorithmic enhancements for large-scale or nonlinear forward models, including a two-stage (decoder, then encoder) procedure when $\mathcal{F}:\mathbb{R}^D \rightarrow \mathbb{R}^O$7 is not directly accessible.

4. Loss Function Innovations and Training Efficiency

Conventional UQ-VAE ELBO-based loss functions may become prohibitively expensive in high-dimensional parameter regimes, especially when requiring large Monte Carlo samples for expectation terms due to the curse of dimensionality. To address this, (Tonini et al., 14 Sep 2025) introduces a new loss function $\mathcal{F}:\mathbb{R}^D \rightarrow \mathbb{R}^O$8 that eliminates sample mean terms and replaces expectation-based moments with column-wise perturbed affine expansions at $\mathcal{F}:\mathbb{R}^D \rightarrow \mathbb{R}^O$9:

$E$0

where $E$1 is the Cholesky factor, and $E$2 evaluates the forward map for each perturbed column. For affine $E$3, this reduces to a deterministic computation. The new loss leads to a factor $E$4–$E$5 reduction in training time and improved scalability, with theoretical convergence to the MAP mean and Laplace covariance as $E$6 (Tonini et al., 14 Sep 2025).

5. Adaptive and Calibration Properties

UQ-VAE possesses an adaptive optimization property: $E$7 is learned to optimally modulate the tradeoff between the data-fitting term and the entropy penalty in the loss. As training proceeds, the uncertainty estimates adapt to the information content in the observations and the data regime (e.g., increasing with noise, decreasing with more data) (Goh et al., 2019). Tuning of $E$8 or $E$9 hyperparameters provides further control, with low $Y = \mathcal{F}(\theta) + E\,,$0 or $Y = \mathcal{F}(\theta) + E\,,$1 yielding more conservative uncertainty envelopes. This self-calibrating mechanism has been empirically validated to yield posterior means frequently closer to the truth than classical MAP, and posterior variances that track expected trends with respect to noise, dataset size, and parameter identifiability (Goh et al., 2019, Tonini et al., 14 Sep 2025).

6. Numerical Benchmarks and Applications

UQ-VAE methodologies have been validated in diverse settings:

• 2D elliptic PDE inverse problems: UQ-VAE achieves relative mean-squared errors of 18.5% on test parameters (for $Y = \mathcal{F}(\theta) + E\,,$2) with well-calibrated 3-$Y = \mathcal{F}(\theta) + E\,,$3 bounds. Posterior mean and variance estimates degrade gracefully as noise increases, and respond as expected to training set size (Goh et al., 2019).
• Poisson, nonlinear, and 0D cardiocirculatory benchmarks: Novel loss $Y = \mathcal{F}(\theta) + E\,,$4 yields mean errors below $Y = \mathcal{F}(\theta) + E\,,$5 in $Y = \mathcal{F}(\theta) + E\,,$6 dimensional Poisson inverse problems, and below $Y = \mathcal{F}(\theta) + E\,,$7 for low-dimensional 0D cardiovascular models. Training times decrease by up to $Y = \mathcal{F}(\theta) + E\,,$8 relative to prior MC-based schemes (Tonini et al., 14 Sep 2025).
• Uncertainty-aware reduced-order modeling: In the UP-dROM framework (Zighed et al., 29 Mar 2025), a convolutional VAE with latent space transformer models transient fluid flows. The VAE quantifies trajectory-wise uncertainty; ensemble statistics on decoded forecasts give uncertainty fields and confidence intervals that align with observed test errors (Pearson $Y = \mathcal{F}(\theta) + E\,,$9 between predicted uncertainty and MSE). This enables adaptive sampling strategies in parameter space, empirically reducing errors across bifurcation diagrams for the Navier–Stokes and Kuramoto–Sivashinsky systems.

7. Extensions, Theoretical Guarantees, and Practical Considerations

Theoretical results for the $p(E)$0 loss provide consistency: the mean and covariance output by the trained encoder recover the MAP estimator and Laplace covariance in the small perturbation limit under regularity assumptions (differentiability of $p(E)$1, independent Gaussian priors and noise, uniqueness of the MAP). This positions UQ-VAE as a variational generalization of the Laplace method (Tonini et al., 14 Sep 2025).

Recommended practices include normalization of parameter and data input spaces, careful initialization near the prior, and scaling weights for stable training at small $p(E)$2. The approach retains performance even with low-regularity posteriors or limited training data; if the forward map is unknown or unavailable, a two-stage decoder–encoder strategy is employed (Tonini et al., 14 Sep 2025).

Potential extensions include adoption of more expressive variational families (e.g., normalizing flows), non-Gaussian likelihoods, and integration with hierarchical/multimodal VAEs for richer nonparametric posterior approximations. In the absence of sufficient training set diversity or quality, uncertainty estimates may falter, highlighting the dependence of UQ fidelity on representative $p(E)$3 ensembles (Goh et al., 2019, Tonini et al., 14 Sep 2025).

8. Summary of UQ-VAE Properties

Aspect	UQ-VAE Approach	References
Inverse problem prior	$p(E)$4	(Tonini et al., 14 Sep 2025, Goh et al., 2019)
Variational posterior	Diagonal/full-covariance Gaussian	(Goh et al., 2019, Tonini et al., 14 Sep 2025)
Loss function	JSD ($p(E)$5 family), MC or surrogate loss	(Goh et al., 2019, Tonini et al., 14 Sep 2025)
Forward model $p(E)$6	True or neural network surrogate	(Goh et al., 2019, Tonini et al., 14 Sep 2025)
Scalability	High-dimensional, efficient at large $p(E)$7	(Tonini et al., 14 Sep 2025)
Uncertainty quantification	Posterior mean and covariance from NN encoder	(Goh et al., 2019, Tonini et al., 14 Sep 2025)
Applications	Inverse problems, UQ-ROM, physical simulators	(Zighed et al., 29 Mar 2025, Tonini et al., 14 Sep 2025)

UQ-VAE thus constitutes a rigorous, computationally scalable, and uncertainty-aware variational inference framework for scientific inverse problems and reduced-order modeling, with demonstrated versatility across linear and nonlinear systems, scalable computational properties, and clear statistical calibration of both mean and uncertainty (Goh et al., 2019, Tonini et al., 14 Sep 2025, Zighed et al., 29 Mar 2025).