Disentangled Deep Priors for Bayesian Inverse Problems

Abstract: We propose a structured prior for high-dimensional Bayesian inverse problems based on a disentangled deep generative model whose latent space is partitioned into auxiliary variables aligned with known and interpretable physical parameters and residual variables capturing remaining unknown variability. This yields a hierarchical prior in which interpretable coordinates carry domain-relevant uncertainty while the residual coordinates retain the flexibility of deep generative models. By linearizing the generator, we characterize the induced prior covariance and derive conditions under which the posterior exhibits approximate block-diagonal structure in the latent variables, clarifying when representation-level disentanglement translates into a separation of uncertainty in the inverse problem. We formulate the resulting latent-space inverse problem and solve it using MAP estimation and Markov chain Monte Carlo (MCMC) sampling. On elliptic PDE inverse problems, such as conductivity identification and source identification, the approach matches an oracle Gaussian process prior under correct specification and provides substantial improvement under prior misspecification, while recovering interpretable physical parameters and producing spatially calibrated uncertainty estimates.

• The paper introduces a hierarchical prior using disentangled deep generative models to separate physical parameters from residual variability.
• It establishes a theoretical framework where block-diagonal covariance in the latent space facilitates precise uncertainty quantification.
• Empirical results on elliptic PDE inverse problems show that the AuxVAE method improves RMSE and credible interval coverage.

Disentangled Deep Priors for Bayesian Inverse Problems

Introduction

"Disentangled Deep Priors for Bayesian Inverse Problems" (2604.02304) proposes a hierarchical prior construction leveraging disentangled deep generative models, particularly the Auxiliary-Variable VAE (AuxVAE) and its extensions, for high-dimensional Bayesian inverse problems. The approach is distinguished by partitioning the latent space into auxiliary variables that are explicitly aligned with known, interpretable physical parameters, and residual variables that encode remaining, typically non-parametric, variability. This partition enables posterior inference and uncertainty quantification directly in terms of physically relevant quantities, while preserving the flexibility of deep generative modeling for capturing complex field structure.

Structured Priors via Disentangled Generative Models

Classical priors for high-dimensional inverse problems, such as GP priors with Matérn or SE kernels, often fail to capture non-Gaussian statistics and intricate structures like sharp interfaces or localized features. Recent advances in DGMs—VAEs, normalizing flows, and score-based methods—have enabled richer prior families via neural generators $G_\theta(z)$ that map tractable latent distributions into the high-dimensional parameter space. However, the main shortcoming is entanglement: the latent coordinates $z$ are rarely aligned with meaningful physical factors, impeding statistical interpretation and physically-grounded UQ in the posterior.

The paper systematically addresses this by constructing a hierarchical prior over $(u, z_{\mathrm{rec}})$, with $u$ denoting user-specified physical parameters and $z_{\mathrm{rec}}$ capturing residual factors:

• The generator $$G_\theta(z) = G_\theta(z_{\mathrm{aux}}, z_{\mathrm{rec}})$$ is trained with the twin goals of explicit alignment (for $z_{\mathrm{aux}}$) and independence (between $z_{\mathrm{aux}}$ and $z_{\mathrm{rec}}$), enforced via polynomial-correlation regularizers.
• The prior is hierarchical: $u \sim p(u)$, $z$0 (soft-anchored), $z$1, $z$2.

This construction induces an implicit, non-Gaussian prior over $z$3 with explicit separation between the physically interpretable portion and a data-driven residual.

Theoretical Analysis: Local Posterior Decoupling

The authors derive conditions under which the approximately disentangled latent structure yields block-diagonal (weakly coupled) uncertainty in the posterior, illuminating the regimes where encoder-side disentanglement translates into posterior interpretability.

• By linearizing $z$4 around representative latent points, they show that the induced prior covariance in the tangent space splits additively into the $z$5 and $z$6 blocks, with cross-covariance controlled by the overlap of the generator’s tangent subspaces (quantified by an orthogonality metric $z$7).
• At the likelihood level, blockwise independence in the posterior requires the forward model’s Jacobian to be block-diagonal in the guided and residual directions under the observation noise metric.
• Encoder-side independence from training objectives (i.e., low cross-correlation in polynomial features) is necessary but not sufficient—generator-side diagnostics of tangent space overlap or direct penalization ($z$8) are essential for posterior-level interpretability.

Latent-Space Inference Mechanics

Posterior inference is performed in the low-dimensional latent space $z$9, yielding substantial computational advantages and interpretability gains:

• MAP and MCMC (HMC, NUTS) are performed over this hierarchical posterior, with the field parameters recovered by mapping posterior samples through the trained generator.
• Gradients are computed via automatic differentiation through $(u, z_{\mathrm{rec}})$0 and the forward operator $(u, z_{\mathrm{rec}})$1, enabling scalable but exact posterior exploration.
• The resulting posteriors permit direct uncertainty quantification on physical parameters $(u, z_{\mathrm{rec}})$2 and, separately, on the residual nonparametric structure.

Empirical Results: Elliptic PDE Inverse Problems

The approach is benchmarked on two canonical elliptic PDE inverse problems—a well-specified conductivity recovery problem and a misspecified source identification task—using the 2D Poisson/heat equation discretized on fine grids.

Conductivity Identification with GP-Structured Priors

For Gaussian process-sampled log-conductivity fields, the AuxVAE prior is competitive with an oracle GP baseline (with known hyperparameters). All four methods (AuxVAE, plain VAE, GP-fixed, GP-hierarchical) achieve low RMSE and near-nominal coverage, but only the AuxVAE and GP baselines recover the true hyperparameters within credible intervals.

Figure 1: Marginal posterior distributions of the GP hyperparameters $(u, z_{\mathrm{rec}})$3 from the AuxVAE, with ground-truth indicated; all true values lie within the $(u, z_{\mathrm{rec}})$4 CI.

This illustrates that the hierarchical structure does not degrade performance in well-specified cases.

Source Identification under Prior Misspecification

In the local Gaussian-bump source identification problem (where standard GP priors are misspecified), disentangled priors yield major performance improvements:

• AuxVAE produces lower RMSE and substantially improved credible interval coverage versus GP baselines.
• The physically interpretable parameters $(u, z_{\mathrm{rec}})$5 for the source are accurately inferred with correct calibration.

  Figure 2: Marginal posterior distributions for the source parameters $(u, z_{\mathrm{rec}})$6 inferred by the AuxVAE.

AuxVAE expresses spatially-aware posterior uncertainty: credible interval widths adapt to regions where residual structure is not captured by the parametric prior.

Latent Posterior Summaries

Figure 3: Hamiltonian Monte Carlo trace plots for the AuxVAE auxiliary coordinates, demonstrating effective mixing and convergence.

High effective sample sizes for the disentangled latent space are achieved, facilitating practical inference with high-dimensional fields, and outperforming hierarchical GP baselines in terms of computational efficiency and mixing.

Practical and Theoretical Implications

This work establishes that carefully structured, disentangled generative priors offer a reliable route to interpretable and physically calibrated uncertainty quantification for high-dimensional inverse problems, even under prior misspecification. By enabling explicit separation of uncertainty between named physical coordinates and flexible residuals, practitioners can perform inference and interpret posteriors in terms of domain-relevant parameters, improving both the utility of Bayesian UQ for scientific analysis and the efficiency of gradient-based samplers.

The hierarchical approach is not reliant on the correctness of classical parametric priors—in settings where the field distribution deviates sharply from GP models, the learned prior maintains calibration and accuracy, while plain VAEs and GP-hierarchical approaches degrade sharply.

Figure 4: Summary comparison of field RMSE and 95% CI coverage across all methods; AuxVAE attains the best reconstruction and correct identification among non-oracle approaches.

Limitations and Future Directions

Key limitations are the dependence on the completeness of the offline dataset—if the generator does not encounter required structure during training, posterior inference may be systematically biased. The approach enforces alignment of latent variables with a user-specified $(u, z_{\mathrm{rec}})$7; it does not automatically discover the correct set of physical drivers. Moreover, since each likelihood evaluation involves a forward solve, gains from latent-space reduction are contingent on forward model cost.

Potential extensions include:

• Incorporating hybrid priors blending explicit parametric structure with the learned residual generator, or explicit model error terms.
• Embedding these priors into large-scale PDE-based inverse problems (e.g., subsurface, geothermal), and using the latent structure for optimal experimental design.
• Enforcing generator-side block-diagonalization penalties during training to guarantee strict tangent-space disentanglement for improved posterior separation.

Conclusion

This work formalizes the use of auxiliary-variable-guided deep generative models as structured, hierarchical priors for Bayesian inverse problems, achieving an overview of interpretability and expressivity in deep-learning-enabled UQ. The methodology yields posterior distributions over physically meaningful parameters and provides statistically grounded UQ, as demonstrated in elliptic PDE benchmarks, surpassing classical and unconstrained deep-learning baselines in robustness, calibration, and computational efficacy. This framework offers a general paradigm for integrating domain knowledge, statistical disentanglement, and deep generative modeling in scientific Bayesian inference (2604.02304).