PaperlayerView exact revision on arXiv

Source-linked AI summary

Variational objectives for amortized Bayesian inference in inverse problems: The role of posterior conditioning

Abhishek Srivastava, Arijit Hazra, Rajesh Dubbaku

arXiv:2609.25145v1stat.MLcs.CEcs.ITcs.LGcs.NAmath.ITmath.NA

TL;DR

Posterior accuracy in inverse problems may depend on the objective, especially when posterior uncertainty concentrates in weakly informed directions. This study compares KL, JS, and JSWA objectives and finds their relative performance depends on posterior conditioning across benchmark and physics-based problems.

  • Problem

    Bayesian inference for inverse problems is computationally challenging because posteriors are rarely closed form, forward-model evaluations can be expensive, and sampling struggles in high dimensions.

    Source: §1 · p. 3

  • Method

    The study compares reverse-KL, asymmetric-JS, and JSWA objectives, which replaces the reverse-KL regularizer with squared 2-Wasserstein distance while retaining forward-KL supervision.

    Source: §6 · p. 35

  • Results

    Across the tested problems, objective performance depends on posterior conditioning: KL performs best when well conditioned, while JS and JSWA become advantageous as ill-conditioning increases.

    Source: §1 · p. 4, §6 · p. 36

  • Takeaways & Limitations

    Posterior conditioning is a useful indicator for selecting a variational objective, and the results motivate geometry-adaptive variational inference.

    Source: §6 · p. 37

    The regularization weights were optimized against held-out MCMC references, and principled selection without such a reference remains open.

    Source: §6 · p. 36

Abstract

from arXiv · show

Variational autoencoders (VAEs) offer an efficient approach to amortized Bayesian inference for inverse problems, but posterior accuracy can depend strongly on the choice of variational regularization, particularly when the inverse problem contains weakly identified parameter directions. This study investigates three objectives: a reverse Kullback--Leibler formulation (VAE-KL), an asymmetric Jensen--Shannon formulation (VAE-JS), and a Jensen--Shannon--Wasserstein formulation (VAE-JSWA), which replaces the reverse Kullback--Leibler regularizer with the squared 2-Wasserstein distance while retaining forward-Kullback--Leibler posterior supervision. A full-covariance Gaussian encoder and a pre-trained physics-based surrogate are used for amortized posterior inference. A local linear--Gaussian analysis in the generalized Fisher basis is developed to characterize the variance-dependent gradients of the three objectives. The formulations are first evaluated using linear--Gaussian benchmarks with known posterior solutions and subsequently tested on nonlinear physics-based inverse problems, including an inverse problem governed by a linear ODE and two PDE-constrained problems. VAE-KL performs slightly better than the other formulations in the well-conditioned benchmark, where all three approaches yield comparable posterior approximations, whereas VAE-JSWA provides substantially lower posterior errors in the strongly ill-conditioned benchmark. The nonlinear physics-based problems exhibit a similar conditioning-dependent trend, with JS-based formulations providing greater benefit as posterior ill-conditioning increases. These results indicate that posterior conditioning is an important factor in selecting variational objectives and motivate geometry-adaptive variational inference for Bayesian inverse problems.

1 Introduction

The introduction asks whether variational objectives should differ across inverse problems with different posterior conditioning. It presents JSWA-VAE and reports conditioning-dependent performance across benchmarks.

1 Introduction

  • Unlike regularization methods that yield a point estimate, Bayesian inference produces a parameter distribution that supports uncertainty quantification for risk assessment, decision-making, and model validation.Source: §1 · p. 2
  • Bayesian inference can require repeated, costly forward-model evaluations, while sampling also struggles with high-dimensional spaces; variational inference instead frames posterior computation as optimization.Source: §1 · p. 3
  • The standard VAE uses reverse KL, while Goh et al. adapt it for inverse problems with an asymmetric Jensen–Shannon objective.Their modified framework treats latent space as parameter space and uses a pre-trained forward-model surrogate.Source: §1 · p. 3
  • Because inverse problems vary in identifiability and posterior conditioning, the introduction asks whether the variational objective affects posterior approximation across these regimes.Source: §1 · p. 3
  • The proposed JSWA-VAE targets PDE-governed inverse problems by replacing modified JS VAE regularization with a squared 2-Wasserstein distance.Source: §1 · p. 3
  • In the linear-Gaussian benchmark, KL is most accurate when well conditioned, while JS and JSWA become advantageous as weakly informed directions dominate uncertainty.The authors also compare the formulations with MCMC reference solutions on nonlinear physics-based problems using posterior mean, covariance, and Wasserstein metrics.Source: §1 · p. 4

2 Mathematical background

This section establishes a Bayesian inverse-problem framework for amortized variational inference, compares KL-, JS-, and JSWA-based regularization, and uses Fisher information to characterize posterior geometry and conditioning. In whitened generalized Fisher coordinates, the eigenvalue spectrum describes how data information shapes posterior uncertainty and mean shrinkage, including distinctions between strongly informed and weakly informed directions.

2.1 Forward model and Bayesian formulation

  • A prior-regularized Bayesian inverse problem models observations with additive Gaussian noise, but PDE-based forward solves can make MCMC costly, motivating amortized inference.Under suitable regularity assumptions, the posterior is well defined and depends continuously on observations in the Hellinger metric [23, Theorem 4.2].Source: §2.1 · pp. 4–5

2.2 Variational inference and variational autoencoders

  • Variational inference optimizes a tractable posterior approximation, while a VAE encoder maps observations to posterior parameters and its standard ELBO combines reconstruction error with prior regularization.The choice and direction of divergence affect the resulting posterior approximation.Source: §2.2 · p. 5

2.3 Modified VAE for Bayesian inverse problems

  • The modified VAE treats latent variables as physical parameters and uses a physics-informed surrogate decoder, making reconstruction measure predicted-versus-observed physical quantities.The surrogate is pretrained and frozen, while a full-covariance encoder captures parameter correlations and produces amortized posterior estimates in one forward pass.Source: §2.3 · pp. 6–7

2.3.1 Distance measures and their properties

  • Reverse KL can be infinite under support mismatch and may favor a dominant mode; JS remains finite but saturates when distributions are widely separated.The support-mismatch pathology is less relevant for the nondegenerate Gaussian distributions considered here, which have common support; strong KL regularization can also drive posterior collapse.Source: §2.3.1 · pp. 7–8
  • JSWA retains reconstruction and forward-KL supervision but replaces reverse-KL regularization with squared 2-Wasserstein distance.Wasserstein distance retains sensitivity to geometric displacement and controls discrepancies in posterior means and variances; reverse KL can be mode-seeking, while JS is bounded but saturates under large support separation.Source: §2.3.1 · pp. 7–9

2.4 Fisher information and posterior geometry

  • Fisher information measures local data sensitivity, and its generalized eigenbasis separates strongly informed parameter directions from weakly informed ones.For nonlinear forward models, the local information matrix is obtained by linearizing at the posterior mode and forms the data-information component of the local Laplace approximation.Source: §2.4 · pp. 9–10
  • Strongly informed generalized Fisher modes have substantially reduced posterior variance, whereas weakly informed modes remain closer to the prior variance.The linear–Gaussian mode-wise posterior variances provide a reference for assessing learned variances.Source: §2.4 · p. 11
  • Generalized Fisher eigenvalues govern mode-dependent shrinkage of the posterior mean toward the data-informed estimate or the prior mean.When an eigenvalue is much less than one, the posterior mean is strongly shrunk toward the prior mean; this appropriately preserves prior information when observations are weak.Source: §2.4 · p. 11
  • Strongly informed directions combine weak posterior-mean shrinkage with strong posterior contraction, while weakly informed directions remain close to the prior in both mean and variance.The corresponding posterior standard deviation quantifies remaining uncertainty in each direction.Source: §2.4 · p. 11
  • A broad generalized Fisher eigenvalue spectrum produces an anisotropic posterior, and sufficiently large spread yields an ill-conditioned, elongated posterior ellipsoid.Narrow directions correspond to strong data information, while broad directions correspond to weak information.Source: §2.4 · p. 12

3 Computational formulation and implementation

This section turns the three VAE objectives into trainable losses for a Gaussian encoder and describes a shared two-stage training procedure. The formulations differ in their prior regularization and, for two of them, use paired-data supervision.

3.1 Implementable loss functions

  • The Gaussian encoder's reconstruction expectation is estimated by Monte Carlo sampling, while its KL and Wasserstein regularizers have closed-form expressions.The reconstruction term is evaluated through the frozen surrogate decoder.Source: §3.1 · p. 12
  • For paired training data, forward-KL supervision is equivalent, up to encoder-independent terms, to the negative log-density of the known parameter under the encoder posterior.The reconstruction expectation uses a single reparameterized draw per training sample.Source: §3.1 · pp. 12–13
  • Formulation III retains reconstruction and forward-KL supervision but replaces the reverse-KL prior regularizer with squared 2-Wasserstein distance.Formulation II adds supervised forward-KL to Formulation I; Table 2 compares the formulations across supervision, support mismatch, geometry, and posterior error bounds.Source: §3.1 · p. 13, Table 2 · p. 13

3.2 Training procedure

  • Training first pre-trains the decoder as a forward-model surrogate, then freezes it while training the encoder with the chosen objective.The decoder is trained by minimizing mean squared error over forward-model evaluations.Source: §3.2 · p. 13, §3.2.1 · p. 14

3.2.1 Stage 1: Surrogate pre-training

  • Surrogate pre-training monitors relative L2 validation error and proceeds only when the error is below 1%, after which decoder weights remain frozen.This threshold is checked on prior-sampled data.Source: §3.2.1 · p. 14, §3.2.2 · p. 14

3.2.2 Stage 2: Encoder training

  • The encoder predicts the posterior mean and full-covariance Cholesky factor, generates reparameterized samples, and receives gradients through the frozen decoder.Each formulation's objective is optimized with Adam; validation loss is used for early stopping.Source: §3.2.2 · p. 14
  • Regularization weights are selected separately by formulation using W2 distance to a reference MCMC posterior on held-out validation data.Formulation I searches β1 in {0.01, 0.1, 1.0}; Formulations II and III search β1 and β2 over the same values.Source: §3.2.2 · p. 14

4 Comparative theoretical analysis of posterior approximation

The theoretical analysis separates surrogate and variational approximation errors, then examines how objective gradients shape posterior variances across generalized Fisher modes. It finds that reverse-KL regularization is more sensitive than Wasserstein regularization as posterior variance becomes small.

4.1 Posterior approximation error bounds

  • The total posterior error separates encoder approximation error from surrogate-induced error, so improving the surrogate reduces only the latter component.The choice of variational objective determines how accurately the encoder approximates the surrogate-induced posterior.Source: §4.1 · p. 15
  • Under appropriate posterior regularity and stability assumptions, surrogate-induced posterior perturbation can be bounded using forward-model approximation error.The surrogate should be accurate enough that its contribution is small relative to variational approximation error.Source: §4.1 · p. 15
  • For a full-covariance Gaussian encoder, squared W2 error accounts for both posterior-mean error and covariance error.The covariance contribution is the squared Bures distance between covariance matrices.Source: §4.1 · p. 15

4.2 Comparative modal gradient analysis

  • Within the local linear-Gaussian approximation, the generalized Fisher basis supports comparing regularizer variance responses, and reconstruction sensitivity to modal variance is independent of current variance and proportional to the generalized Fisher eigenvalue.Source: §4.2 · p. 16
  • As a mode's posterior variance approaches zero, the Wasserstein regularizer's gradient grows more slowly than the reverse-KL gradient.This difference underlies the greater reverse-KL variance sensitivity in highly concentrated posterior directions.Source: §4.2 · p. 17
  • Each mode's variance equilibrium balances reconstruction, forward-KL correction toward the local posterior variance, and prior regularization toward prior variance.The forward-KL contribution is negative below the local posterior variance and positive above it.Source: §4.2 · p. 17

5 Numerical results

The numerical results assess how variational-objective performance varies with posterior conditioning across linear–Gaussian benchmarks and nonlinear inverse problems, including one-dimensional heat conduction and two-dimensional problems.

5.1 Linear Gaussian benchmark with exact analytical posterior

  • In the ill-conditioned benchmark, JSWA achieves the lowest posterior-mean, covariance, and W2 errors: 0.0732, 0.0128, and 0.079, respectively.Across 500 test observations, its W2 error is approximately 21.7% lower than JS and 35.2% lower than KL.Source: Table 4 · p. 19, §5.1 · p. 19
  • For the ill-conditioned benchmark’s dominant weak-information mode, JSWA recovers variance with 3.0% error, versus 12.5% for JS and 15.5% for KL.This mode accounts for approximately 87% of total posterior variance, explaining why its improved recovery strongly affects global covariance accuracy.Source: Table 5 · p. 20, §5.1 · p. 20
  • In the well-conditioned benchmark, KL has the lowest mean, covariance, and W2 errors: 0.0221, 0.092, and 0.0263, respectively.Its W2 distribution is shifted toward smaller values, although all three objectives remain relatively close.Source: Table 7 · p. 23, §5.1 · p. 23, Figure 5 · p. 23
  • Across the benchmark regimes, JSWA’s advantage is strongest when uncertainty is concentrated in weakly informed directions, whereas KL performs best when posterior uncertainty is distributed across several modes.In the ill-conditioned case, forward-KL shifts the JS equilibrium toward the target variance but shifts the JSWA equilibrium upward from an underestimated variance.Source: §5.1 · pp. 20, 24

5.2 One-dimensional steady-state heat conduction

  • In the nonlinear one-dimensional heat-conduction inverse problem, with posterior condition numbers from 48.3 to 68.6, KL gives the most accurate distributional approximation.The JS and JSWA W2 CDFs are nearly coincident and shifted toward larger errors than KL; the heat equation is linear in temperature, but the parameter-to-observation map is nonlinear.Source: §5.2 · p. 28
  • For the representative one-dimensional case, VAE-KL has the lowest relative mean error (0.309%), covariance error (4.486%), and W2 distance (0.00212).VAE-JS has relative mean error 0.347%, covariance error 23.827%, and W2 distance 0.00360; VAE-JSWA has relative mean error 0.710%, covariance error 9.595%, and W2 distance 0.00436. Across the ensemble, VAE-KL also has lower median W2 than VAE-JSWA by approximately 0.007.Source: Table 9 · p. 29, §5.2 · p. 29

5.3 Two-dimensional steady-state heat conduction with spatially varying conductivity

  • In the ensemble-level posterior comparison, VAE-JSWA is shifted toward smaller W2 values over most of the distribution and reduces median W2 by 0.0250 versus VAE-JS and 0.0628 versus VAE-KL.The median posterior condition number is approximately 1929.3, with observed values ranging from approximately 1779.7 to 2253.9.Source: §5.3 · p. 30
  • For the representative two-dimensional conductivity case, VAE-JSWA has the lowest posterior mean error (3.68%), covariance error (24.114%), and W2 distance (0.2111).VAE-JS has corresponding errors of 11.11%, 45.762%, and 0.4173; VAE-KL has 15.09%, 108.088%, and 0.6577. JSWA also has the lowest L2 posterior-mean field error: 0.59%, versus 0.89% for JS and 1.07% for KL.Source: Table 10 · p. 31, §5.3 · pp. 31–32, Figure 11 · p. 32

5.4 Two-dimensional nonlinear steady-state heat conduction with joint parameter inference

  • VAE-JS and VAE-JSWA substantially outperform VAE-KL in ensemble W2, while JSWA improves on JS by only 0.0044 at the median.The median W2 reduction from VAE-KL to VAE-JSWA is approximately 0.1636; the posterior condition number has a median of approximately 11526.3 and ranges from 1939.5 to 328168.8.Source: §5.4 · p. 33, Figure 12 · p. 34
  • For the representative nonlinear PDE case, VAE-JSWA has the lowest L2 posterior-mean field error (0.13%) and posterior mean, covariance, and W2 errors (3.52%, 3.91%, and 0.0359).The L2 field errors are 0.47% for VAE-KL and 0.17% for VAE-JS; JSWA's improvement over JS in W2 is more pronounced for this realization than at the ensemble median.Source: §5.4 · pp. 33–34

6 Conclusions and future work

The conclusions show that the best variational objective depends on posterior conditioning, with JSWA strongest when uncertainty concentrates in weakly informed directions. The appendices describe the finite-element heat model, its parameterization, and the MCMC reference procedure.

6 Conclusions and future work

  • At κ(Σ*_post) ≈9, VAE-KL has the lowest errors; at κ(Σ*_post) ≈8000, VAE-JSWA has lower posterior mean, covariance, and W2 errors than VAE-KL and VAE-JS.In the strongly ill-conditioned benchmark, the weakest Fisher mode accounts for approximately 87% of total posterior variance.Source: §6 · p. 36
  • Across nonlinear problems, VAE-KL leads for the moderate-conditioning ODE case, while JSWA leads in the more anisotropic PDE cases and both JS-based formulations outperform KL in the final problem.Median posterior covariance condition numbers are approximately 55.7 for the ODE problem and 1929.3 for the two-dimensional conductivity inversion.Source: §6 · p. 36
  • The nonlinear results do not establish a quantitative conditioning-scaling relationship for general nonlinear forward models because the cases lack a common global Fisher basis.They provide evidence that conditioning-dependent behavior from controlled analysis can persist in nonlinear physics-based inverse problems.Source: §6 · p. 36
  • Future work proposes geometry-adaptive objective weights, while noting that weight selection without an MCMC reference remains open; richer posterior families and more efficient surrogates are also proposed.The current Gaussian encoder cannot represent genuinely multimodal posteriors, and the experiments cover dimensions d ≤9.Source: §6 · p. 36

A Finite Element Formulation

  • The heat model uses temperature as the state, conductivity and source as model inputs, and prescribed temperatures on left and right boundaries with insulated top and bottom boundaries.Temperature-dependent conductivity makes those cases nonlinear.Source: A Finite Element Formulation · p. 41

A.1 Karhunen–Lo`ve Parameterization

  • The spatial conductivity is represented by a truncated Karhunen–Loève expansion of its logarithm, with truncation chosen to retain a prescribed fraction of total variance.The parameter-to-observation map sends the finite-element temperature solution to m temperature observations; the same forward map is used for the VAE decoder and MCMC reference.Source: §A.1 · p. 42

B.1 Adaptive Metropolis Algorithm

The Adaptive Metropolis algorithm samples the posterior with a Gaussian random-walk proposal and updates its covariance online from the chain history.

B.1 Adaptive Metropolis Algorithm

  • Adaptive Metropolis updates the proposal covariance every K iterations after t>d, using the empirical covariance of all samples generated so far.The default adaptation interval is K=100.Source: §B.1 · pp. 43–44
  • The proposal begins at user-supplied covariance C0 and draws Gaussian random-walk proposals; failed covariance decomposition triggers a fallback to 0.1I.Source: §B.1 · p. 44
  • Empirical covariance is rescaled by sd=2.382/d and regularized with ϵI, where ϵ=10^-6, with extra eigenvalue regularization near singularity.The regularization maintains positive definiteness and improves numerical stability.Source: §B.1 · p. 44

B.2 Convergence Criteria

MCMC runs are accepted as converged only when both effective sample size and rank-normalized split-R-hat criteria are satisfied.

B.2 Convergence Criteria

  • Convergence diagnostics include trace plots, autocorrelation, ESS, and rank-normalized split-R-hat across four independent chains.Source: §B.2 · pp. 44–45
  • A run is accepted only when both ESS and R-hat criteria are satisfied.Source: Table 12 · p. 45
  • The rank-normalized split-R-hat threshold is R-hat<1.01 for each parameter after burn-in.The diagnostic compares between-chain and within-chain variability after rank normalization and chain splitting.Source: §B.2 · p. 45
  • Each parameter must have a minimum effective sample size of 400, estimated from post-burn-in chain autocorrelation.Source: §B.2 · p. 45

B.3 Comparison Metrics

The comparison metrics assess posterior agreement in mean, spread, covariance structure, and Gaussian Wasserstein discrepancy. The Fisher-basis analysis describes posterior variance and mean shrinkage by information mode, with nonlinear results interpreted locally.

B.3 Comparison Metrics

  • The W2 metric compares the Gaussian VAE posterior with a moment-matched Gaussian fitted to postburn-in MCMC samples.For nonlinear problems, it is not computed directly against the empirical MCMC distribution.Source: §B.3 · p. 45
  • Full-covariance error captures diagonal variance and off-diagonal correlation errors that marginal spread metrics alone miss.The combination of marginal spread and covariance metrics gives a more complete picture of posterior covariance recovery.Source: §B.3 · pp. 45–46

C.1 Posterior Covariance in the Fisher Basis

  • In prior-whitened coordinates, posterior covariance is diagonal in the generalized Fisher basis, with modal variances determined by Fisher information.The analysis also gives the corresponding posterior mean and modal shrinkage relations.Source: §C.1 · pp. 46–47
  • Weakly identified modes remain close to the prior mean, whereas strongly identified modes approach the true parameter component.For nonlinear forward models, replacing the forward operator with its Jacobian yields a local Fisher-information and Laplace-covariance approximation, not an exact posterior variance.Source: §C.1 · p. 47

C.2 Derivation of Loss Functions and Modal Gradients

The appendix derives closed-form loss expressions and variance-gradient formulas to support the analysis of the three VAE objectives.

C.2 Derivation of Loss Functions and Modal Gradients

  • The derivations provide closed-form expressions for Gaussian divergence measures and the reconstruction loss used in the computational formulation.Source: §C.2 · p. 47
  • The appendix derives gradient expressions for the objectives in the whitened generalized Fisher basis.Source: §C.2 · p. 47
  • These modal gradients are used to analyze how the different VAE formulations respond to variance changes.Source: §C.2 · p. 47

C.2.1 Derivation of Implementable Loss Functions

This subsection presents implementable Gaussian loss terms, including KL divergences, squared 2-Wasserstein distance, and reconstruction loss.

C.2.1 Derivation of Implementable Loss Functions

  • The Gaussian KL expression supplies both reverse-KL prior regularization and a forward-KL term approximated using paired parameters from the target posterior.The reverse-KL implementation drops an additive constant.Source: §C.2.1 · p. 48
  • For Gaussian distributions, the squared 2-Wasserstein distance includes a covariance-dependent squared Bures term that vanishes exactly when the covariances match.The implementation evaluates the full-covariance expression using eigendecomposition.Source: §C.2.1 · p. 48
  • For the dimensions considered, n ≤15, the O(n^3) computation is negligible relative to the neural-network forward pass.Source: §C.2.1 · p. 48
  • The Gaussian reconstruction loss is the negative log-likelihood, evaluated with the frozen surrogate and one reparameterized posterior sample per iteration.Source: §C.2.1 · p. 48

C.2.2 Derivation of the Reconstruction Gradient in the Whitened Fisher Basis

The reconstruction-gradient derivation isolates variance effects in a locally mean-matched encoder and shows that sensitivity varies with Fisher information.

C.2.2 Derivation of the Reconstruction Gradient in the Whitened Fisher Basis

  • The local analysis uses a linear-Gaussian model with isotropic observation noise, diagonal encoder covariance in the whitened Fisher basis, and a calibrated encoder mean.Source: §C.2.2 · p. 49
  • Mean matching isolates the reconstruction loss's variance-dependent contribution and does not assert that the posterior mean generally equals the realized true parameter.Source: §C.2.2 · p. 49
  • The reconstruction contribution is constant with respect to encoder variance within each Fisher mode, while its magnitude is proportional to that mode's Fisher eigenvalue.Thus reconstruction is more variance-sensitive in strongly informative directions and less sensitive in poorly informative directions; additional variance-dependent behavior comes from forward-KL and prior-regularization terms.Source: §C.2.2 · p. 50

C.2.3 Derivation of the Forward-KL Gradient

The subsection isolates encoder-variance dynamics in a local linear-Gaussian setting, using the generalized Fisher basis to analyze individual modes. The forward-KL gradient drives each mode’s variance toward its exact posterior value.

C.2.3 Derivation of the Forward-KL Gradient

  • With encoder and exact posterior means matched, the derivation isolates how the forward-KL objective depends on encoder posterior variance.The analysis uses the benchmark’s linear-Gaussian setting.Source: §C.2.3 · p. 50
  • In the generalized Fisher basis, diagonal encoder covariance makes the variance-dependent objective contributions separable across modes.The analysis considers one mode’s scalar variance while holding its Fisher eigenvalue fixed.Source: §C.2.3 · p. 50
  • The forward-KL term is nonnegative, vanishes at the exact posterior variance, and corrects modal variance in either direction toward that value.Its gradient increases variance below the posterior value and decreases it above.Source: §C.2.3 · p. 51
Loading 2609.25145v1…