Uncertainty Visualization via Low-Dimensional Posterior Projections

TL;DR

This work tackles uncertainty visualization in high-dimensional inverse problems by learning the projected posterior distribution within a low-dimensional, input-adaptive affine subspace. It introduces PPDE, a conditional energy-based model that, given a measurement $\mathbf{y}$ and a subspace $\mathcal{A}(\mathbf{y})=\{\mathbf{x}_0(\mathbf{y}),\mathbf{W}(\mathbf{y})\}$, outputs the projected posterior density $p_{\mathbf{v}|\mathbf{y}}(\mathbf{v}|\mathbf{y})$ for $\mathbf{v}=\mathbf{W}(\mathbf{y})^T(\mathbf{x}-\mathbf{x}_0(\mathbf{y}))$. The subspace is typically anchored at the MMSE predictor with directions given by NPPC's top posterior PCs, enabling informative 1D/2D visualizations. The method uses a two-part architecture with a heavy feature extractor and a small conditional MLP, trained via a multilevel contrastive divergence to model $p_{\mathbf{v}|\mathbf{y}}(\mathbf{v}|\mathbf{y})$ and evaluated on MNIST, CelebA-HQ, ImageNet, and biology datasets. PPDE outperforms Gaussian approximations and diffusion-based KDE in terms of log-likelihood and speed, providing a practical framework for uncertainty visualization in high-dimensional imaging while acknowledging current limits on higher-dimensional projections.

Abstract

In ill-posed inverse problems, it is commonly desirable to obtain insight into the full spectrum of plausible solutions, rather than extracting only a single reconstruction. Information about the plausible solutions and their likelihoods is encoded in the posterior distribution. However, for high-dimensional data, this distribution is challenging to visualize. In this work, we introduce a new approach for estimating and visualizing posteriors by employing energy-based models (EBMs) over low-dimensional subspaces. Specifically, we train a conditional EBM that receives an input measurement and a set of directions that span some low-dimensional subspace of solutions, and outputs the probability density function of the posterior within that space. We demonstrate the effectiveness of our method across a diverse range of datasets and image restoration problems, showcasing its strength in uncertainty quantification and visualization. As we show, our method outperforms a baseline that projects samples from a diffusion-based posterior sampler, while being orders of magnitude faster. Furthermore, it is more accurate than a baseline that assumes a Gaussian posterior.

Figures (24)

• Figure S1: Informed uncertainty visualization. Point estimation methods receive a distorted image and output only a single solution, e.g., the MMSE estimator (top). NPPC nehme2023uncertainty complements MMSE estimators with input-adaptive uncertainty directions (principal components of the posterior) without modeling output likelihood (middle). Our method (bottom) learns the input-adaptive projected posterior distribution, facilitating a likelihood-informed uncertainty visualization.
• Figure S2: Notations illustration. An image $\bm{x}$ is projected onto a measurement-adaptive affine subspace $\mathcal{A}(\bm{y})$, and represented by its projection coefficients $\bm{v}=[v_1,v_2]^{\top}$. The restored image within the subspace is denoted by $\bm{x}_{\text{restored}}$.
• Figure S3: Posterior slicing vs. projection. (a) Schematic illustrating the difference between the slice and the projection of a high-dimensional manifold onto a low-dimensional subspace. (b) Sliced and projected posterior comparison for a 1D affine subspace defined by the posterior mean and the first PC in the task of digit inpainting. The sliced posterior has been computed using an EBM.
• Figure S4: Architecture overview. The degraded image $\bm{y}$ is first fed to a pre-selected subspace extractor that outputs an input-adaptive subspace $\mathcal{A}(\bm{y})=\{\bm{W}(\bm{y}),\bm{x}_0(\bm{y})\}$. Afterward, the degraded image $\bm{y}$ and the extracted subspace are fed to a feature extractor that outputs a feature vector $h(\bm{y})$. The resulting $h(\bm{y})$ modulates a lightweight MLP that outputs $p(\bm{v}|\bm{y})$ for any query $\bm{v}$. The resulting projected distribution can then be navigated to visualize posterior uncertainty.
• Figure S5: Gaussian mixture denoising. (a) Prior distribution $p_{\mathbf{x}}{(\bm{x})}$, and a sample from the joint distribution $p_{\mathbf{x},\mathbf{y}}{(\bm{x},\bm{y})}$. (b) Full posterior distribution $p_{\mathbf{x}|\mathbf{y}}{(\bm{x}|\bm{y})}$, and the selected 1D subspace $\mathcal{A}(\bm{y})=\{\bm{x}_0(\bm{y}),\bm{W}(\bm{y})\}$. (c) Sliced $f^{\text{sliced}}(\bm{v}|\bm{y})$ (purple) vs. projected $p_{\mathbf{v}|\mathbf{y}}(\bm{v}|\bm{y})$ (green) posterior distribution along the 1D subspace (dashed gray line in (b)). (d) Comparison of the GT and the learned projected posterior.