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

1. Problem Statement

Inverse problems involve reconstructing hidden quantities of interest (QoI), denoted as $x$, from noisy, indirect observations $y$, typically modeled as $y = F(x) + \varepsilon$. These problems are often ill-posed, meaning small perturbations in $y$ can lead to large changes in the admissible $x$.

While classical variational methods (e.g., Tikhonov, Total Variation) and Bayesian approaches provide theoretical guarantees, they often require repeated evaluations of the forward operator $F$ and careful tuning of regularization parameters, making them computationally expensive for high-dimensional or nonlinear settings. Conversely, data-driven approaches like end-to-end neural networks offer fast inference but typically yield only point estimates, lacking the probabilistic structure necessary for Uncertainty Quantification (UQ).

Existing generative methods for UQ, such as score-based diffusion models, can approximate the posterior $p(x|y)$ but often sacrifice inference speed due to iterative sampling requirements. Furthermore, many latent-space approaches suffer from "dense" latent codes where information is entangled across all dimensions, obscuring the sources of uncertainty and making it difficult to interpret which data characteristics drive reconstruction errors.

2. Methodology: The vsPAIR Framework

The authors propose the Variational Sparse Paired Autoencoder (vsPAIR), a framework designed to provide fast, amortized inference alongside structured, interpretable uncertainty estimates. The architecture integrates three core components:

2.1 Paired Autoencoder Structure

Inspired by "Latent Twin" architectures, vsPAIR employs two separate autoencoders:

1. Observation Encoder: A standard Variational Autoencoder (VAE) encodes observations $y$ into a latent space $z_y$ using a Gaussian variational posterior $q_{\phi_y}(z_y | y)$.
2. QoI Encoder: A Sparse Variational Autoencoder (sVAE) encodes the quantities of interest $x$ into a latent space $z_x$ using a spike-and-slab variational posterior $q_{\phi_x}(z_x | x)$.
3. Latent Mapping: A learned mapping $M_{\leftarrow}$ connects the two latent spaces, transforming the distribution parameters of the observation encoder $(\mu_y, \sigma_y)$ to those of the QoI encoder $(\mu_x, \sigma_x, \omega_x)$.

This pairing allows the QoI encoder to learn features anchored to clean data ($x$) rather than directly from corrupted observations, separating the learning of signal features from the inversion process.

2.2 Sparse Latent Representations (sVAE)

To address the issue of entangled information, the QoI encoder utilizes a spike-and-slab distribution for its latent prior and posterior. For each latent dimension $i$, the distribution is a mixture of a point mass at zero (the "spike") and a continuous Gaussian (the "slab"):
$$q_{\phi}(z_i | x) = \omega_i \mathcal{N}(z_i | \mu_i, \sigma_i^2) + (1 - \omega_i)\delta_0(z_i)$$
Here, $\omega_i$ represents the probability that the $i$-th dimension is active. This structure forces the model to concentrate information into a small number of identifiable factors (active dimensions) rather than diffusing it across the entire latent space.

2.3 Training Objective

The framework is trained end-to-end by minimizing a composite loss function (Equation 14) that includes:

• ELBO for QoI ($L_{sVAE}$): Encourages reconstruction of $x$ from sparse latent codes.
• ELBO for Observations ($L_{VAE}$): Encourages reconstruction of $y$.
• Latent Mapping Loss ($L_M$): Minimizes the Mean Squared Error (MSE) between the true latent parameters of $x$ and the predicted parameters derived from $y$ via $M_{\leftarrow}$. An optional binary entropy penalty is added to encourage decisive gating (pushing $\omega$ toward 0 or 1).
• Beta Regularization ($L_\rho$): A term used to learn the prior sparsity level $\rho$.

Crucially, the framework does not require access to the forward operator $F$ during training; it relies solely on paired samples $(x, y)$.

3. Key Contributions

The paper outlines three primary contributions:

1. The vsPAIR Framework: A novel architecture combining probabilistic inference, sparsity, and paired latent mappings to enable fast, uncertainty-aware inversion without requiring the forward operator.
2. Experimental Validation: Demonstration of competitive reconstruction performance on ill-posed inverse problems, including blind inpainting, computed tomography (CT), and initial-condition inference for the heat equation.
3. Structured and Interpretable Uncertainty: Evidence that vsPAIR yields sparse latent representations where active dimensions correspond to core reconstruction features. Consequently, uncertainty is not diffused but localized to these specific features, allowing variance in estimates to be traced back to identifiable sources.

4. Experimental Results

The authors evaluated vsPAIR against deterministic baselines (PAIR), variational baselines (vPAIR), and sparse-only baselines (sVAE), as well as diffusion-based UQ methods (DPS).

• Gaussian Linear Model: In a simple test case with a known posterior, vsPAIR successfully learned a structured latent space where active dimensions collapsed to match the true dimensionality of the data. Perturbations in these active dimensions aligned with the eigenvectors of the true posterior covariance.
• Blind Inpainting (MNIST):
  • Reconstruction: vsPAIR achieved competitive MSE, though slightly higher than deterministic PAIR, consistent with the regularization trade-off inherent in VAEs.
  • Uncertainty Localization: Qualitative and quantitative analyses showed that vsPAIR could identify latent dimensions that selectively controlled corrupted regions. While vPAIR and vsPAIR showed similar "localization ratios" (variance concentrated in missing regions), vsPAIR demonstrated that its sparse, paired structure allowed for more distinct separation between global structure and local uncertainty.
  • Stability: vsPAIR maintained a stable subset of consistently active dimensions (approx. 11.5 per image), whereas the single-encoder sVAE showed high variability in active dimensions, suggesting the paired structure stabilizes feature learning.
• Computed Tomography (LoDoPaB-CT):
  • vsPAIR produced reconstructions competitive with TV-ADMM and faster than Diffusion Posterior Sampling (DPS).
  • Uncertainty estimates were concentrated on edges and structural features (e.g., holes, bone structures).
  • Unlike DPS, which sometimes introduced "hallucinated" fine details, vsPAIR exhibited conservative behavior, which the authors note may be favorable in medical applications where avoiding false positives is critical.
  • The model successfully concentrated variance into a sparse set of latent indices (approx. 1,383 active dimensions out of 4,096).
• Heat Equation (Nonlinear):
  • Applied to inferring initial conditions from later time states, vsPAIR achieved high-fidelity reconstructions (MSE 0.071).
  • Variance plots correctly concentrated around primary heat sources, and latent traversals showed that specific indices controlled macro-features of the reconstruction.

5. Significance and Claims

The paper claims that vsPAIR offers a principled path toward interpretable Uncertainty Quantification. Unlike standard UQ methods that provide aggregate measures of confidence (e.g., "the model is 80% sure"), vsPAIR aims to answer what is uncertain. By anchoring QoI representations to clean data and enforcing sparsity, the framework isolates uncertainty to specific, identifiable latent factors.

The authors position vsPAIR as particularly suitable for applications where:

• Interpretability is paramount: Users need to understand which features drive ambiguity (e.g., distinguishing pathology from artifacts in medical imaging).
• Forward models are unknown or expensive: The method requires only paired data, not the operator $F$.
• Fast inference is required: Once trained, vsPAIR provides amortized inference via a single forward pass, avoiding the iterative sampling costs of diffusion models.

The paper modestly acknowledges limitations: the model footprint is large (approx. 550M parameters for CT), and the activation probabilities ($\omega$) and variances ($\sigma$) of the latent distribution showed limited dependence on the specific observation $y$ in experiments, suggesting the model learns a reduced subspace rather than fully observation-specific UQ. The authors note that while the framework does not guarantee calibration with the true Bayesian posterior, its structure allows for the characterization of model-induced uncertainty in a reduced setting. Future work is identified as necessary to better calibrate the framework and refine the sVAE implementation for greater data dependence.