Hierarchical Uncertainty Exploration via Feedforward Posterior Trees

TL;DR

Ill-posed imaging inverse problems yield multi-modal posteriors that are difficult to inspect. The authors introduce posterior trees, a single-pass neural model that outputs a depth-$d$, $K$-ary tree of prototypes and associated probabilities to visualize $p_{\mathbf{x}|\mathbf{y}}(\mathbf{x}|\mathbf{y})$ at multiple granularities. The method builds a bottom-up tree whose root matches the MMSE estimator $\hat{\mathbf{x}}_{\mathrm{MMSE}}=\mathbb{E}[\mathbf{x}|\mathbf{y}]$, using an amortized oracle loss and an annealed regularization to prevent collapse, plus a weighted sampling scheme to handle imbalanced posteriors. Across denoising, colorization, inpainting, and bioimage translation, posterior trees achieve competitive performance with diffusion-based baselines while offering dramatically faster inference ($\approx$7 ms per image on a GPU) and enabling efficient, interactive uncertainty exploration in practice.

Abstract

When solving ill-posed inverse problems, one often desires to explore the space of potential solutions rather than be presented with a single plausible reconstruction. Valuable insights into these feasible solutions and their associated probabilities are embedded in the posterior distribution. However, when confronted with data of high dimensionality (such as images), visualizing this distribution becomes a formidable challenge, necessitating the application of effective summarization techniques before user examination. In this work, we introduce a new approach for visualizing posteriors across multiple levels of granularity using tree-valued predictions. Our method predicts a tree-valued hierarchical summarization of the posterior distribution for any input measurement, in a single forward pass of a neural network. We showcase the efficacy of our approach across diverse datasets and image restoration challenges, highlighting its prowess in uncertainty quantification and visualization. Our findings reveal that our method performs comparably to a baseline that hierarchically clusters samples from a diffusion-based posterior sampler, yet achieves this with orders of magnitude greater speed.

Figures (21)

• Figure 1: Hierarchical decomposition of the minimum-MSE predictor into prototypes in the task of mouth inpainting. The predicted tree explores the different options of bigger/smaller lips, mouth opening/closing, round/square jawline, etc.
• Figure 2: Method overview. Our model $\mathcal{T}({\bm{y}};{\bm{\theta}})$ receives a degraded image ${\bm{y}}$ and predicts $\{\hat{{\bm{x}}}_{k_1,\dots,k_d}\}_{k_1,\dots,k_d=1}^K$, the bottom $K^d$ leaves, and their probabilities $\{\alpha_{k_1,\dots,k_d}\}_{k_1,\dots,k_d=1}^K$ (faint blue box; illustrated here for $K=3$ and $d=2$). Next, the tree is iteratively constructed from the bottom up using weighted averaging, until we reach the root node which is the minimum MSE predictor $\hat{{\bm{x}}}_{\text{MMSE}}$. During training, starting from the root, the ground truth ${\bm{x}}$ is propagated through the tree until it reaches the leaves (dashed red lines). At tree level $d$, $x$ is compared to its immediate $K$ children nodes, and the MSE loss to the nearest child is added to the loss trajectory.
• Figure 3: 2D Gaussian mixture denoising. \ref{['fig:gmm-task']} Underlying signal prior $p_{{\mathbf{x}}}({\bm{x}})$ (blue heatmap), and training samples $({\bm{x}}_i,{\bm{y}}_i)\sim p_{{\mathbf{x}},{\mathbf{y}}}({\bm{x}},{\bm{y}})$. \ref{['fig:gmm-flat-kmeans']}$K$-means with $K=4$ applied to 10K samples ${\bm{x}}_i\sim p_{{\mathbf{x}}|{\mathbf{y}}}({\bm{x}}|{\bm{y}}_t)$, for a given test point ${\bm{y}}_t$ (red circle). The resulting cluster centers (blue markers) partition the underlying posterior $p_{{\mathbf{x}}|{\mathbf{y}}}({\bm{x}}|{\bm{y}}_t)$ (red heatmap), resulting in cluster probabilities $p({\bm{y}}_t)$. \ref{['fig:gmm-hier-kmeans']} Hierarchical $K$-means applied twice with $K=2$ on 10K samples ${\bm{x}}_i\sim p_{{\mathbf{x}}|{\mathbf{y}}}({\bm{x}}|{\bm{y}}_t)$. At depth $d=1$, the posterior is partitioned by the dashed blue line (blue triangles mark cluster centers). The resulting half spaces are subsequently halved by the dashed orange and green lines respectively. \ref{['fig:gmm-ours']} Posterior trees (ours) with degree $K=2$ and depth $d=2$. Note that in all cases the estimated posterior mean $\hat{{\bm{\mu}}}({\bm{x}}|{\bm{y}}_t)$ (black star) coincides with the analytical mean ${\bm{\mu}}({\bm{x}}|{\bm{y}}_t)$ (red star), while in (c)-(d) the lowest density mode is better represented. $\mathcal{T}({\bm{y}}_t)$/$p({\bm{y}}_t)$ are drawn at the bottom of (b)-(d).
• Figure 4: Diverse applications of posterior trees. The predicted trees represent inherent task uncertainty: e.g.,\ref{['fig:shoes']} Refining the mean estimate by color, grouping similar colors, while still depicting unlikely ones (e.g., the blue boot); \ref{['fig:colorization']} Presenting various plausible colorizations varying by hat color, skin tone, and background; and \ref{['fig:eye_inpainting']} Exploring the diverse options of eyebrows/eyeglasses.
• Figure 5: Bioimage translation. Here we explored posterior trees for the task of translating the image of a tissue from one fluorescent dye to another. The resulting trees expose important information regarding uncertain cells (yellow/red arrows), e.g., ones that do not consistently appear in all branches, and additionally explore different plausible cellular morphology consistent with the input.