Variational Sparse Paired Autoencoders (vsPAIR) for Inverse Problems and Uncertainty Quantification

TL;DR

The paper introduces vsPAIR, a variational sparse paired autoencoder framework for inverse problems that jointly encodes observations with a standard VAE and quantities of interest with a sparse VAE, connected by a learned latent map to enable fast amortized inference and calibrated uncertainty. It contributes a beta hyperprior on sparsity and a hard-concrete spike-and-slab mechanism to learn sparse latent representations, along with theoretical results in a reduced setting and empirical validation on MNIST blind inpainting and low-dose CT (LoDoPaB-CT). The approach yields interpretable, structured uncertainty by anchoring QoI representations to clean data and concentrating information into a subset of latent factors, while maintaining competitive reconstructions. The work highlights potential applications in areas where understanding the source and location of uncertainty is crucial (e.g., medical imaging, geophysics), and discusses limitations such as computational overhead and the need for further theoretical and empirical exploration in nonlinear settings.

Abstract

Inverse problems are fundamental to many scientific and engineering disciplines; they arise when one seeks to reconstruct hidden, underlying quantities from noisy measurements. Many applications demand not just point estimates but interpretable uncertainty. Providing fast inference alongside uncertainty estimates remains challenging yet desirable in numerous applications. We propose the Variational Sparse Paired Autoencoder (vsPAIR) to address this challenge. The architecture pairs a standard VAE encoding observations with a sparse VAE encoding quantities of interest, connected through a learned latent mapping. The variational structure enables uncertainty estimation, the paired architecture encourages interpretability by anchoring QoI representations to clean data, and sparse encodings provide structure by concentrating information into identifiable factors rather than diffusing across all dimensions. We also propose modifications to existing sparse VAE methods: a hard-concrete spike-and-slab relaxation for differentiable training and a beta hyperprior for adaptive sparsity levels. To validate the effectiveness of our proposed architecture, we conduct experiments on blind inpainting and computed tomography, demonstrating that vsPAIR is a capable inverse problem solver that can provide interpretable and structured uncertainty estimates.

Figures (12)

• Figure 1: Deterministic paired autoencoder (PAIR) architecture where QoI and observation are related by a forward process, $F$: $y = F(x) + \varepsilon$. Separate Autoencoders are trained to encode the QoI ($x$) and observation ($y$), as well as surrogate forward $M^\to$ and inversion $M^\gets$ inversion operators that operate on the latent encoding.
• Figure 2: Schematic of the vsPAIR framework. The QoI is encoded via an sVAE producing sparse latent representations $z_x$, while observations are encoded via a standard VAE. The latent mapping $M^\gets$ operates on distribution parameters, mapping $(\mu_y, \sigma_y)$ to $(\hat{\mu}_x, \hat{\sigma}_x, \hat{\omega}_x)$.
• Figure 3: Reconstruction and variance comparison across variational methods. Test cases are selected at the 25th, 50th, and 96th percentile of reconstruction error averaged across all three methods, corresponding to easy, medium, and hard cases respectively. For each case, we show the corrupted observation $y$, ground truth $x$, mean reconstruction $\hat{x}$ averaged over 30 samples from the QoI latent distribution, and the pixel-wise variance across those samples. vsPAIR is shown on the top row, vPAIR on the middle row, and sVAE on the bottom row. Higher variance in yellow indicates regions where the model produces diverse reconstructions, while lower variance in purple to black indicates confident predictions. All variance maps share a common normalized color scale.
• Figure 4: Sparse latent representations across methods. Top: corrupted observation $y$, ground truth $x$, and reconstructions from each method. Bottom: absolute value of each latent dimension $|z_j|$ obtained through the inversion path from $y$. Each bar chart is annotated with the number of nonzero dimensions.
• Figure 5: Latent dimension perturbation analysis for the first test example in \ref{['fig:latent_codes']}. For each method, we select the three latent dimensions with highest sensitivity, measured by reconstruction MSE change under perturbation. For sparse methods, we restrict to dimensions active in all 30 stochastic samples. Each image shows the mean reconstruction across 30 samples with the specified perturbation applied to that dimension. Perturbation ranges from $-400\%$ to $+400\%$ of the original value; the center column in boxes shows the unperturbed reconstruction.

Theorems & Definitions (2)

• Theorem 1
• Theorem 2