Research map / Probabilistic and scientific modeling / Uncertainty and Gaussian processes
Hierarchical Uncertainty Exploration via Feedforward Posterior Trees
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes posterior trees, a feedforward U-Net that outputs a hierarchical tree-valued summary of an inverse problem's posterior in a single forward pass, matching hierarchical K-means over diffusion/GAN posterior samples at a fraction of the cost.
Problem
Visualizing and summarizing the high-dimensional posterior distribution of ill-posed imaging inverse problems so users can explore plausible solutions efficiently, since posterior sampling plus clustering is too slow.
Approach
Train a network with a hierarchical (amortized) oracle/winner-takes-all loss so that its K^d outputs form a tree of hierarchical centroidal Voronoi cells, with parent nodes computed bottom-up as probability-weighted means of children; leaf probabilities come from an MLP head. Includes gradient-scaling annealing to prevent hypothesis collapse and a weighted sampler for imbalanced posteriors.
Key results
- Posterior trees (single NFE, ~7 ms on 1080Ti) achieve PSNR and NLL comparable to hierarchical K-means over 100 DDNM/DDRM/RePaint samples, e.g., colorization d=2 PSNR 26.4 vs 26.1 for DDNM.
- The method is reported to be 10^2-10^6 times faster than the sampling-based baselines in NFE (10,000 for DDNM, 457,000 for RePaint vs 1).
- On a 2D toy mixture-of-Gaussians denoising task, the learned trees match the analytical-posterior-based approximate ground truth in cluster centers and probabilities.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace DDNM posterior sampling and clustering with posterior trees | FFHQ colorization · FFHQ colorization | optimal PSNR (d=2): 26.1$\pm$3.6 -> 26.4$\pm$4.0 | Helped |
| Replace DDNM posterior sampling and clustering with posterior trees | FFHQ colorization · FFHQ colorization | NFE: 10000 -> 1 | Helped |
| Replace RePaint posterior sampling and clustering with posterior trees | FFHQ mouth inpainting · FFHQ mouth inpainting | optimal PSNR (d=2): 20.7$\pm$2.3 -> 20.4$\pm$2.2 | Hurt |
| Replace RePaint posterior sampling and clustering with posterior trees | FFHQ mouth inpainting · FFHQ mouth inpainting | NFE: 457000 -> 1 | Helped |
| Replace RePaint posterior sampling and clustering with posterior trees | FFHQ eyes inpainting · FFHQ eyes inpainting | optimal PSNR (d=2): 20.6$\pm$2.8 -> 19.8$\pm$2.4 | Hurt |
| Replace RePaint posterior sampling and clustering with posterior trees | FFHQ eyes inpainting · FFHQ eyes inpainting | NFE: 457000 -> 1 | Helped |
| Use hierarchical rather than flat clustering (ablation) | Posterior trees on the toy denoising example · Toy denoising | representation of rare posterior modes | Helped |
| Use a binary hierarchical tree rather than a flat tree with the same number of leaves (ablation) | Posterior trees on digit inpainting · MNIST digit inpainting | faithfulness of the '4' cluster | Helped |
| Use smaller tree degree and greater depth (ablation) | Posterior trees on digit inpainting · MNIST digit inpainting | emphasis on weaker posterior modes | Helped |
| Use weighted sampling with Adam instead of standard sampling with Adam (ablation) | Posterior trees for image inpainting · Image inpainting | quality of low-likelihood leaves | Helped |
| Use SGD instead of Adam with standard sampling (ablation) | Posterior trees for image inpainting · Image inpainting | plausibility of low-likelihood leaves | Helped |
| Fully share parameters across prediction leaves instead of using separate decoders (ablation) | Posterior trees for image inpainting · Image inpainting | prediction diversity | Hurt |
Surprising findings
- GAN-based sampler MAT collapses to a single mode, producing little sample variance, so its hierarchical K-means trees show little variance.
- Adam with standard sampling leads to implausible predictions in low-likelihood leaves, requiring the proposed weighted sampler.
Limitations
- Hyperparameters K and d are task-dependent with no rule of thumb, and balanced fixed-degree trees may be suboptimal.
- Number of output leaves is limited since the whole tree is amortized into one forward pass; deeper trees would need iterative inference.
- Prototypes are cluster centers averaging multiple plausible solutions, so shallow-tree prototypes are not expected to lie on the image manifold.
Experimental setup
Datasets: MNIST Edges2Shoes CelebA-HQ FFHQ AFHQ(v2) bioimage migrating cells dataset
Benchmarks: FFHQ 100-image subset (colorization, mouth inpainting, eyes inpainting)
Metrics: optimal PSNR negative log-likelihood neural function evaluations
Baselines: DDNM DDRM RePaint MAT Flat K-means Hierarchical K-means
Builds on: multiple choice learning hierarchical K-means U-Net
Closest papers by meaning
- Uncertainty Visualization via Low-Dimensional Posterior Projections (CVPR 2024 · similarity 0.90)
- Provable Posterior Sampling with Denoising Oracles via Tilted Transport (NeurIPS 2024 · similarity 0.86)
- A Symmetry-Aware Exploration of Bayesian Neural Network Posteriors (ICLR 2024 · similarity 0.85)
- Divide-and-Conquer Posterior Sampling for Denoising Diffusion priors (NeurIPS 2024 · similarity 0.85)
- Consistency Posterior Sampling for Diverse Image Synthesis (CVPR 2025 · similarity 0.85)
- NPN: Non-Linear Projections of the Null-Space for Imaging Inverse Problems (NeurIPS 2025 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.