Hierarchical Uncertainty Exploration via Feedforward Posterior Trees
Authors: Elias Nehme
Rotem Mulayoff
Tomer Michaeli
Venue: NeurIPS 2024 poster (2024)
Keywords: Uncertainty Quantification
Explainable Computer Vision
Inverse Problems
Computational Imaging
Hierarchical Clustering
Primary area: probabilistic_methods
TL;DR: We predict a tree-valued hierarchical summarization of the posterior distribution for any input measurement, in a single forward pass of a neural network

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. Code and examples are available at our [webpage](https://eliasnehme.github.io/PosteriorTrees/).

Hierarchical Uncertainty Exploration via Feedforward Posterior Trees
Elias Nehme
Rotem Mulayoff CISPA Helmholtz Center for Information Security rotem. mww.aof com com
Electrical and Computer Engineering Technion- Institute of Technology sel illasseecampus technion. ac il
Tomer Michaeli
Electrical and Computer Engineering Technion-Israel Institute of Technology tomer m@ee. technion.c..... ac. il
Abstract
When solving ill-posed inverse problems, one often desires to explore the space potential solutions rather than be presented with single plausible reconstruction. Valuable insights into these feasible solutions and their associated probabilities are embedded in the posterior distribution. However, when confronted with data high dimensionality (such as images), visualizing this distribution becomes 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 the posterior distribution for any input measurement, of 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 baseline that hicrarchically clusters samples from diffusion- based posterior sampler, yet achieves this with orders of magnitude greater speed, Code and examples are available our webpage.
a single forward pass
1 Introduction
Communicating prediction uncertainty is a crucial element in advancing the reliability of machine learning models. This is particularly important in the realm of imaging inverse problems, in which given input can typically be associated with multitude of plausible solutions. In such cases, it is advantageous to equip users with efficient tools for exploring and visualizing the of admissible solutions. Such tools may be of high value especially in safety-criticll domains, such as scientific and medical image analysis [1 [1-4], where inaccuracies in predictions could impact human lives.
Information about the plausible solutions and their respective likelihoods is encapsulated within the posterior distribution. However, dithessinnsion posteriors challenging to visualize. One popular approach for communicating uncertainty to generate samples from the posterior [5-14]. Yet, in complex domains with high uncertainty levels, users may need sift through hundreds of posterior samples per input to be able to confidently validate or refute suspicions about the unobserved ground-truth image [15]. Several methods have been proposed for generating a concise of samples that highlight posterior diversity [15, 16]. However, choosing this set to be small is often insuflicient for summarizing the posterior, while choosing be large hinders the user's ability to properly
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Input
Neural network
p=0.50
p=o
p=0.13
p=0. 00
p=o 04
p=0.06 p=0.14
p=0.30
p=0.13
p=
p=o o7
M
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 1:
inspect all solutions in the set. Other methods proposed to visualize uncertainty by allowing the user to navigate along the principal components (PCs) of the posterior [17-20]. While more interactive in nature, the projection onto a low dimensional principal subspace commonly accounts for only a small fraction of the reconstruction error.
In this work, we propose to model posterior uncertainty through hierarchical clustering. by using tree-structured outputs. In stark contrast to all existing methods, our approach supports efficient user interaction, allowing examination of only a small number of hypotheses en route to accepting or refuting hypothesis about the unknown image. Specifically, our model receives degraded measurement and outputs a hierarchy of predictions that cluster the solution set at multiple levels of granularity, each accompanied by their probability. This allows us to visualize the likelihood of different posterior modes, facilitating an informed visual exploration posterior uncertainty (Fig. 1). Our method, which we coin posterior frees, is a generic technique for uncertainty visualization that is seamlessly transferable across tasks and datasets. Compared the popular practice of training point estimators with mean square eπror (MSE) minimization, This is because it adds a minor additional computational burden, while simultancously accompanying the point estimate with its decomposition to the constituting clusters, revealing the underlying modes of variation and disclosing posterior uncertainty.
method is an ideal drop-in replacement.
We showcase our method on multiple inverse problems in imaging, demonstrating the practical benefit of tree- veluud predictions. We also quantitatively compare our learned posterior trees to two-step baseline that generates samples using a diffusion- posterior sampler and then applies hierarchical K-means clustering. Our approach achieves a comparable performance across tasks and datasets while enjoying a significantly faster runtime.
2
Related Work
The common avenue for probing data uncertainty (also known as aleatoric uncertainty [6-8, 21- 23]) involves posterior sampling using conditional generative models like [5-14, 24-27], with the dominance of ecore-based/diffusion models [28-31] in recent years. While posterior sampling theoretically offers a sizable solution set per input, allowing useful uncertainty estimates through summarization methods like PCA [18] or K-means, its implementation proves impractically slow for contemporary state-of-the-an models. Despite attempts to enhance speed [32-35], the strategy remains plagued by extended run times.
In an effort to expedite inference, some studies proposed approximating posteriors with simpler distributions considering pixel correlations, such as correlated Gaussians [36-39]. More recently, two approaches emerged, directly outputting the top principal components (PCs) of the posterior distribution without imposing distributional assumptions [17, 19]. Additionally, subsequent work by

---

Yair et al. [20] visualized the projected posterior distribution onto the subspace spanned by the top PCs, unveiling its multi-modal nature at heightened uncertainty levels.
Multiple choice learning (MCL) [40-50] was proposed as a framework that can provide a multi-modal set of predictions when confronted with ambiguous inputs. The idea was originally proposed by Guzman-Rivers et al. [40] using structured support vector machines, and was later introduced to deep models in [42, 43]. The main working principle of MCL is to train an architecture with multiple output heads using the oracle loss, such that for a given input, only the head that is closest to the desired output gets updated. This winner-takes--a strategy encourages each head to specialize in different mode of the output space, effectively predicting the centers of a Voronoi tessellation of all possible outputs [45]. Later works such as [48, 50] proposed to complement each prediction in the set with its likelihood. Nonetheless, existing multi- hypothesiss methods are mainly tailored to classification settings, and are still limited in their ability to organize the predicted set of hypotheses effectively. Categorizing the solution set hierarchically, as we propose here, can significantly accelerate uncertainty exploration. For example, it allows iterative user interaction, focusing user efforts on the examination of a small number of hypotheses (not necessarily the most likely ones) to confirm/refute a suspicion about the signal. Moreover, the hierarchical structure can regularize the number of prototypes devoted to high-density modes, highlighting rare cases.
3
Method
Our goal is to predict a clean signal x E X C Rd=, given some degraded measurements y E ) C R4,, We assume that a and y are realizations of random vectors x and y with an unknown joint distribution Px,y (x. y). from which our predicted trees is therefore to visualize the uncertainty in the posterior distribution Px ( y via a hierarchical set of few prototypes.
have training set of i.i.d. samples D = {( w.) The objective of
3.1 Multiple Output Prediction
Many image restoration methods output a single prediction d for any given input A common choice is to aim for the posterior mean ã Exxly = y). which is the predictor that minimizes the MSE. However, a single prediction does not convey to the user the uncertainty in the restoration, especially if the solution set is multi-modal. When confronted with such ambiguous tasks, it is more natural to predict a small set of prototypical restorations .t......k" } that portray the different options. In MCL, this is achieved by minimizing the so-called maacle/winne-ttakss.ee loss given by
LMCL (z, { = min f(x. t.),
i=l,..,KK
where .(: )) is a loss function that measures prediction distance. Note that LMCL. is minimized as long one restoration in the predicted set is close æc This encourages prediction diversity, such that each restoration specializes in a different mode of the posterior. Rupprecht et al. [45] proposed probabilistic interpretation of MCL, showing that the optimal minimizers of Eq. (1) form a centroidal Voronoi tessellation (CVT) of the posterior, where { are given by the conditional means/centroids of the resulting K Voronoi cells (see Appendix E). discrete setting where we are given samples from Pxly (xly), CVT with the Euclidean distance is equivalent to K -means clustering.
Our goal here is generalize the concept of predicting a set of restorations. Given input y. instead (y)). we propose to output an input-adaptive tree T(y) representing hierarchical clustering of the posterior Pxly (mly). Specifically, nodes to correspond to a hierarchical CVT of the posterior (analogous to a hierarchical K K-means process in the discrete setting), where the children of each node constitute CVT for the cluster represented by the node. Such a tree can organize the different restorations in a manner that facilitates their navigation by a user and can also accompany each node with its relative cluster probability.
of predicting an unordered set { (v),... : .
want the tree
3.2 Posterior Trees
As discussed above, training a model to predict tree of degree K and depth d = 1 with the oracle loss, results in a CVT of the posterior Pxly (xly) for every y. Let [A ( denote the resulting Voronoi cells at d for some input y y, such that uk Ax X, with X e 1d being the
3

---

MMMRS
Bottom buakding Bug
Tree root Posterior mean
(Available
7
only in training)
022
6
9
GT
Leaves + probs.
1111
0122
01.3
0g21
02.3
63.1
032
$3,3
(
7
9 11,2
9 31.3
Input
Neural net 11.1
F2.1
122.2
22,3
353.1
JF.22
2r,33
Figure 2: Method overview. Our model Tyy; 0) receives degraded image y and predicts (āx,.....k ) .==1...u=1 the bottom Kt leaves, and their probabilities {ūkk,.... }K box; illustrated here for K = and d = 2). Next, the tree is iteratively constructed from the bottom up using weighted averaging, until we reach the @MSEE During training, starting from the root, the ground truth is propagated through the tree until reaches the leaves (dashed red lines). At tree level is compared its immediate K children nodes, and the MSE loss to the nearest child is added to the loss trajectory.
Jk=....k.=1 (faint blue
node which is the minimum MSE predictor
output space. Training with l loss, i.e. f(tt, x)) ||0 x the optimal K predicted leaves and probabilities are given by [45, 50]
( = [ = v. a (y) F(x Au(v)yy y) == K.
We now want to extend the tessellation to the children of each node at the next level. Let {Au., (y)) denote the Voronoi cells at d = 2, partitioning/ressellainng the Voronoi cell As(y) from level into K sub cells such that UK Ak,q (y) = (x))). We want the K22 leaves and their associated probabilities outputted by our model be given by
(1) = E xxly = y. (k)
a k (3) = P( (u)ly)ly = y), k. 1,.., K.
For simplicity of notation, in the following, we omit the dependence of the Voronoi cells on and write A instead. Equations (2) and (3) expose an interesting connection between tree levels. First, recall that since {Au.oJ o a tessellation, satisfies UK - = Ag. and Aky nAkz
for (ki, qi) * (k2, 42). Invoking the law of total expectation we can write
E x yx Av)
Σ E y,x AC, Ak P Acgly UXXE Ak).
(4)
Equation (4) can be simplified by noting that A& n A.Akig = Ar..g hence the expectation in the summand given by E|xly = y.x Ar.gle Moreover, we further simplify the probability in Eq. (4)
by writing
P P v, P(XE Auly y. PR,) Akgly =
P(x A&
P(X Akglyy =
(5)
Σx P(XE = y)
where the first equality follows from Bayes' rule and the second equality is by the law of total probability and the fact that Ak.g C As. which implics that P(x Asly y. Akg)) = 1.

---

Substituting Eq. (5) and the simplified expectation in Eq. (4) with the notations of Eqs. (2) and (3), we obtain that
K
C (y)
(1)
L kg (y)
(6)
2-=1 ,
4=1
K
(7 = a (3))
k= -=....,,K K,
(7)
where Eq. (7) is by the law of total probability. Note that Eqs. (6) and (7) expose the redundancy in predicting both {0t (z) (y)) and {#%,, (V), 0 (2)). as given the latter, we can compute the former in closed form. Furthermore, although our derivation assumed a tree of depth d = 2, the result trivially generalizes to tree of depth d. Given the leaves and their probabilities at level d. {#t (v), 6% .k....k [y) 1.......,.11 we can compose their parents at the upper d levels, by employing Eqs. (6) and (7) iteratively from the bottom up.
Here, this is precisely the property we exploit to output an imput apputaadapiive tree T(y) of degree K and depth d, using a single model. Specifically, we train model 7 Tyy: 0) e (Xa (3: 9), ag 0)) that outputs Kd prediction leaves (yy;; 0) (f...::: }>]]>] and their accompanying probabilities dea (y; 0) {3: {)} K (see Fig. 2). Next, starting from the bottom leaves at depth d, the rest of the tree is composed level by level, recursively employing Eqs. (6) and (7) from the bottom up.
During training, the tree hierarchy is enforced through our loss function. Given some input y: the model Tyy: 0) outputs are used as explained earlier to construct the full tree, where the root of the tree approximates the posterior mean E y/J. Next, the label x; is first compared the tree root with an & loss to ensure our decomposition recovers the MMSE estimator. At each successive level, the label x; compared to the children of the parent with minimal loss, and the closest child is added to the loss trajectory (Fig. 2). This process is repeated until we arrive at the bottom leaves.
Note that the described training scheme serves two purposes. First, induces the desired tree structure by employing an amortized hierarchical version of the oracle loss in Eq. (1). Second, the loss employed on the root node enables predicting branch probability, mitigating the need for supervising the probabilities with explicit targets, See Appendix D for more details on our training algorithm and hierarchical oracle loss function.
3.3 Architecture
The approach presented in Section 3.2 is general, and can be used to augment any architecture outputting Kd images and probabilities. Here, we choose to adapt the U-Net architecture [51] as generic established choice for image-to-image regression tasks. The number of output images is K' to accommodate tree of degree and depth In case all features are shared the architecture, this amounts to changing the number of output filters at the last layer. However, in challenging tasks such as image inpainting, we found that sharing all parameters leads to reduced prediction diversity. In such cases, to trade off feature sharing and output diversity, the decoder, such that each prediction has a disjoint set of channels learned separately from other outputs. As for the skip connections from the encoder, the concatenated features are interleaved equally per level such that each prediction has an equal share. In addition to the output images, our architecture also predicts Kd probabilities (see Figs. 2 and A2). This is achieved by (global) average pooling all feature maps from the decoder, and feeding their concatenation to an additional lightweight multi-layer perceptron (MLP), with four linear layers and softmax at the output.
group convolutions [52] in
3.4 Preventing Tree Collapse
Vanilla training with the oracle loss is known to suffer from "hypotheses collapse", where some of the predictions are implausible [40, 42, 43, 45, 46, 50, 53, 54]. This phenomenon occurs when some predictions are initialized worse than others in stochastic training, and therefore receive little to no gradient updates due to not being chosen in the oracle loss. This results degenerate outputs that are not encouraged to produce meaningful results. Previous works proposed various regularizations to remedy this undesired effect. For example, Rupprecht et al. [45] relaxed the arg min operator with a small constant E, such that predictions with non mon-minimal loss get updated with e-scaled gradients.
5

---

polr 34
poor gA.
T
Tim
PlaU 10.3 0.18, D.S, DURT
P13 BLOS 0.12 060 0019)
11.8 09,0 ATI
(a) Denoising task
(b) "Flat" K-means
(e)Hiierrchical K-means
(d) Ours
Figure 2D Gaussian mixture denoising. (a) Underlying signal prior px(x) (blue heatmap),
and training samples (I,, ys) ~ Px,y(2; y). (b) K-means with K 4 applied to 10K samples
2; ~ (xly/ for a given test point yv: (red circle). The resulting cluster centers (blue markers) partition the underlying posterior Pxly (e)))) (red heatmap), resulting in cluster probabilities p(y:). (c) Hierarchical K K-means applied twice with K = 2 10K samples E; Pxly (l..). At depth = 1, the posterior is partitioned by the dashed blue line (blue triangles mark cluster centers). The resulting half spaces
subsequently halved by the dashed orange and green lines respectively. (d) Posterior trees (ours) with degree = 2 and depth = 2. Note that in all cases the estimated posterior mean jo( (z) (3)3)) (black star) coincides with the analytical mean fu 3/) (red star), while in (c)-(d) the lowest density mode is better represented. (p(())) are drawn at the bottom of (b)-(d).
Moreover, Makansi et al. [54] proposed an evolving version the oracle loss, where the top-k predictions are updated in each step, with & being annealed down to k = 1 during training, such that it starts with equal weights to all predictions and gradually shifts to updating only the best one towards convergence. Here, employ a strategy that combines the simplicity of [45] with the adaptivity of [54]. Specifically, we scale the gradients of non-performing predictions with a constant e that is annealed during training according
EO exp 4/]+
(8)
where t is the epoch number, and we fixed = 1 and to = 5 for all experiments. At the beginning of training, this has the benefit of bringing all predictions to reasonable starting point, alleviating sensitivity to initialization and top prediction domination. As training progresses, this regularization decayed, and only the relevant prediction is chosen for each sample, hence converging to the desired MCL behavior (see Appendix B).
3.5 Weighted Sampling
Albeit the regularization mentioned in Section 3.4, for tasks with highly imbalanced posteriors, training with Adam [55] leads to trees with large disparity in leaf quality. This is because less likely leaves get chosen with lower frequency during training, resulting in transient gradients that highly affect the adaptive normalization in Adam. This hypothesis was also validated in matched settings by using SGD with an appropriate learning rate which resulted in significantly slower, yet improved convergence (see Appendix A.3). Therefore, to keep the speed advantage provided by Adam, in the task of image inpainting, we opted for weighted sampler during training. The purpose this sampler is to roughly balance out the overall number of occurrences the bottom leaves during training. This is achieved by keeping track of an association matrix A RRK xN where Ai,, counts the number of times sample æ, was associated with leaf c, over some time window. This allows us to estimate the conditional probability Pe|x (c; (z;). ), and subsequently derive an optimized sample probability achieve roughly uniform marginal leaf probability across the entire training set. The loss of each sample is then adjusted to account for this intervention, avoiding tampering with the original posterior probabilities. See Appendix C for more details.
4 Experiments
Here we demonstrate our method in several settings. In all experiments except for the toy example, we used variants of the U-Net architecture [51], with a custom number of output images, adjusted to accommodate Kd predictions according to the desired tree layout. Morcover, explained in

---

Section 3.3, each model was also complemented with a small MLP for predicting leaf likelihood. Full details regarding the architectures, optimizer, and per-task settings are Appendix A.
Toy Example. As a warm-up, we demonstrate posterior trees on a 2D denoising task. Here, sampled from a mixture of four Gaussians (arranged in a rhombus-like layout), and y is a noisy version of x. The prior distribution Px (x(x)) and exemplar samples from (x, y) are presented in Fig. 3(a). For this simple case, the posterior distribution Pxly (arly) can be calculated analytically (see Appendix G for the derivation) and is also a mixture of Gaussians. Therefore, this task can serve as a sanity check enabling us to benchmark our results. To demonstrate our method on this toy example, train model (3y: 0) outputting tree of depth d = 2 and degrec K = 2. The results are compared against an approximate "ground truth" in Fig. 3(c). Note that even though we have closed-fort expression for the posterior density, we still have to approximate the "ground truth" posterior tree by applying hierarchical K-means clustering to samples for every y. As seen in Fig. 3(d), our method resulted in highly accurate posterior trees both in terms of cluster centers and cluster likelihoods. Moreover, this result was achieved while being presented only with a single posterior sample per y in training. In contrast, the approximate ground truth computed with 10K samples. Moreover, it is important to note that besides trivial cases, our result could not have been achieved with "flat" K-means clustering with K = 4 (Fig. 3()).. This is because K-means tends to focus on high-density modes, whereas our hierarchical trees better convey rare cases of the posterior. See Appendix for more details and examples.
Handwritten Digits, Edges-+Shoes, and Human Faces. Figure 2 demonstrates posterior trees on inpainting the top 70% of handwritten digits from the MNIST dataset. As can be seen, at depth d 1, the learned tree exposes the two likely modes averaged in the mean estimate å being either a "7" or "9", In addition, deeper tree levels, the different modes are further refined to reveal intricate intra-digit variations. More examples are available in Appendix M. We also applied posterior trees to the edges-to-shoce dataset taken from pix2pix [56. 57]. Here, the task is to convert an image of black and white edges to an output RGB image of a shoe. As shown in Fig. 4(a), our tree is capable of representing diverse shoe colors, in an adaptive manner to the input contours.
Next, we tested posterior trees on face images from the CelebA-HQ dataset, using the split from CelebA [58]. Figure 4(b) demonstrates our method applied to image colorization, where the input is a grayscale image and the desired output is its RGB version. The resulting trees hierarchically refine the output predictions by background, skin tone, and hat color. We also tested our method the task of image inpainting. Figures I and 4(c) show the resulting trees for mouth/eye inpainting respectively. For example, the predicted tree in Fig. 4(c) explores the different options of eye-opening/closiee seebrow-raising/lowee
eyeglasses, etc. This demonstrates that even shallow trees can still depict diverse reconstruction characteristics, showcasing the benefit of outputting multiple predictions.
Bioimage Translation. Transforming images from one domain to match the statistics of images from another is commonly referred as the task of image-to-image translation [57]. In the realm of bioimaging, such transformations were utilized to predict fluorescence from bright-fiedd images [1]. "virtually stain" unstained tissue [3], and transfer the images of one fluorescent dye to appear as they were imaged with another [59]. The common scenario in image-to- to-gaee translation is the task being highly posed with wide range of plausible transformations satisfying the desired statistics. For general image editing/style transfer, this is desired property adding to the artistic excitement; However, in bioimaging, this is highly problematic as the result often informs a downstream task with high stakes, and therefore output uncertainty should be communicated.
Here, we applied our method to a dataset of migrating cells, imaged in a spinning-din disk microscope, simultancously with two different fluorescent dyes (one staining the nuclei and one staining actin filaments) [59]. The task was to predict the image of one fluorescent dye (nuclear stain) from another (actin stain). Figure 5 demonstrates the predicted tree for a 128 x 128 test patch. The tree conveys important information to the user exposing uncertain cells, and exploring optional cell shapes. These can for example affect downstream tasks such as cell counting and morphological cell analysis.
7

---

GT
Inpal
MMSE
p=0.31
pa0001
p=t.09
p=0.16
p=0.10
p=0.05
0.001
p=0.40
p=b.1
p=0.00
p=0.03
p=0.05
(a) Edges
GT
Input
MMSE
p=0.47
p=0.25
p=O 08
p=0.19
p=0.13
p=0..4
Joz d.08
p=0.08
p=0.00
p=0.05
p=0.044
(b) Face colorization
GT
Inpot
MMSE
p=0.24
p=0.03
pm0.13
p=0.07
p=0.07
p=d.19
p=0.24
p=0.21
p=o 07
p=d 03
(c) Eyes inpainting
Figure 4: Diverse applications of posterior trees. The predicted trees represent inherent task uncertainty: egg., (a) Refining the mean estimate by color, grouping similar colors, while still depicting unlikely ones (e.g., the blue boot); (b) Presenting various plausible colorizations varying by hat color, skin tone, and background; and (c) Exploring the diverse options of eyebrows/eyeglassee.

---

GT
Inpot
pand.37
p=m
p=0.01 p=0.27
p=0.05
p=0.02
p=0.08
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.
Table Comparison to proposed baseline on 100 test images from the FFHQ dataset. Hierarchical
K K-means was applied to 100 posterior samples per test image. Runtime is reported as both the speed of forward pass (sec) and the memory usage (GB) required to infer a single test image with a batch of on an A6000 GPU. The NLL at the root node (d 0) is trivial and therefore omitted. Blue and
Red indicate best and second best respectively.
Optimal PSNR (t)
NLL (L)
Speed (sec 4)
Memory
Task
Method
d= 0
d= 1
d d=2
d=1
d=
(GB L)
DDNM
24.6+3.8
25.5++3.6
26.1+:3.6
0.9+0.44
2.0
340
18500
Color.
DDRM
22.71:2.7
24.011:.0
24.813.0
1.1:0.4
2.010.5
3700
Ours
24.6:t..1
25.7=13.9
26.41:000
0.910.4
2.0:0.7
0.014
1.3
DDNM
19.342.6
19.91222
20.242.0
1.1:0.4
2.120.6
340
18500
Mouth
RePaint
19.8:227
20.5+12.3
20.7+22.3
1.0.1..3
1.900.4
15538
845450
inpaint
MAT
19.2.12.3
19.4122.3
19.6122.3
1.2.1..3
2.510.66
15
300
Ours
20.1:12.4
20.5.12.3
20.4.12.2
0.911..4
2.0.0.88
0.014
1.3
DDNM
19.3::2.7
19.6:12.5
19.5
1.1:0.7
2.1:0.6
340
18500
Eyes inpaint
RePaint
20.1+-2.8
20.4++2.8
20.6+:2.8
1.040.4
1.9+0.5
15538
845450
MAT
19.0++3.0
19.24:3.0
19.3+3.0
1.240.3
2.4+0.6
15
300
Ours
19.6:12.7
19.94-2.5
19.8
1.340.6
2.5
0.014
1.3
4.1 Quantitative Comparisons
Baseline. Building on the notion of clustering from Section 3.1, we propose simple baseline to benchmark our results. As shown in Fig. 3(d), our predicted trees are constructed out of prototypes that yield a hierarchical clustering of the posterior. In discrete settings, this is equivalent to applying Hierarchical K-means
samples from (xll). Therefore, given method to sample from the posterior (e.g., using [7, 8, 10, [1], natural baseline our case a two-step procedure: (i) Generate N. samples from the posterior, and (ii) Apply hierarchical K-means d times to generate tree of degree K and depth d. Note that due to computational reasons, the comparison performed over random subset of 100 test images from FFHQ, where for each image we generate N, = 100 samples, and apply hierarchical K-means with K = and d = 2. To ensure a fair comparison, at each clustering step, K K-means run times, keeping only the clusters with the best objective.
Metrics. To compare our predicted trees to the proposed baseline, we opted for two different metrics. The first metric the PSNR (equivalent to MSE) between the ground truth test image and the tree nodes along the optimal path starting from the root and ending at the leaves. Intuitively, an accurate posterior clustering implies our tree nodes should maximize the PSNR, representing the ground truth test image with increasingly higher accuracy (lower MSE) function of (optimal)

---

node depth. The second metric we adopt is the sample negative log- log-likeihhod (NLL) (using the natural logarithm) of the ground truth test image under the predicted posterior partitioning. Here, an accurate tree is expected to maximize the sample log-likelihood. 'hii metric serves the purpose of verifying the predicted probabilities, as in practice we do not have access to the ground truth posterior distribution, and estimating cluster probability based on posterior samples becomes worse as a function of depth. Table compares posterior trees to the proposed baseline, implemented with various state-of-the-a aht samplers. The results suggest that our method yields comparable performance in both the PSNR along the optimal path and in sample log-likelihood while requiring a single forward pass (<2 14 with 1.3 GB memory footprint than the competition without considering our wide advantage in memory footprint. In case we do consider batched setting with sizable test set, our advantage becomes even more pronounced, where our method enables the inference of z 430 test images in a single second. See Appendices to I. for more details and visual comparisons.
A6000 GPU). This is 10 107 faster
5 Discussion and Conclusion
We demonstrated the wide applicability of posterior trees across diverse tasks and datasets. However, our method is not free of limitations. First, the optimal hyper-parameters K and d are task-dependent, with no rule of thumb for determining them a-priori (see Appendix H), Second, in this work, we focused on balanced trees with a fixed degree K which might sub-optimal for certain posteriors. Devising strategy for an input-adaptive tree layout is an exciting direction for future research. Third, our method is limited in the number of output leaves, as we amortize the entire tree inference to single forward pass. For predicting significantly deeper trees, an iterative inference procedure conditioning the model on the node index (in an analogous fashion to the timestep in diffusion models) is required. Finally, our method is tailored towards visualizing uncertainty and not sampling realistic-looking images. Although with increased depth our prototypes become increasingly sharper, they are nonetheless still cluster centers that average multiple plausible solutions and hence not expected to lic on the image data manifold for shallow trees. A possible solution is to apply posterior trees in the latent space of autoencoder such as VQ-VAE [60, 61], however, this is beyond the scope of this paper.
To conclude, in this work, we proposed technique to output a hierarchical quantization of the posterior a single forward pass. We discussed key design choices underlying our approach, including bottom-up tree construction, principled training scheme, and proper regularization techniques to prevent tree collapse. We further demonstrated the benefit of hierarchical clustering over flat trees and discussed the intuition behind it on toy example. In experiments, applied our method to highly iol-posed inverse problems and showed that diverse (prototypical) reconstructions are possible with a simple training scheme exposing uncertainty. Additionally, also proposed an appropriate baseline based on posterior sampling, and quantitatively compared our approach to several strong samplers. Our method demonstrated least comparable results while being orders of magnitude faster. Finally, we applied our method to the challenging task bioimage translation, demonstrating its practical relevance in real-world application.
10

---

Acknowledgments and Disclosure of Funding
This research was partially supported by the Israel Science Foundation (ISF) under Grant 2318/22 and by gift from Elbit Systems. This work was done while Rotem Mulayoff was at the Technion. The authors also thank Hila Manor and Noa Cohen for their help with baseline comparisons.
References
[1] Chawin Ounkomol, Sharmishtaa Seshamani, Mary M Maleckar, Forrest Collman, and Gregory R Johnson. Label-free prediction dirre-dimensional fluorescence images from transmitted-lgg -ight microscopy. Nature methods, 15(11):917-95) 2018. 1,
[2] Eric M Christiansen, Samuel Yang. D Michael Ando, Ashkan Javaherian, Gaia Skibinski, Scott Lipnick, Elliot Mount, Alison O'neil, Kevan Shah, Alicia Lee, et al. In silico labeling: predicting fluorescent labels in unlabeled images. Cell, 173(3):792-803 2018.
[3] Yair Rivenson, Hongda Wang, Zhensong Wei, Kevin de Haan, Yibo Zhang, Yichen Wu, Harun Günaydın, Jonathan E Zuckerman, Thomas Chong, Anthony E Sisk, al. Virtual histological staining of unlabelled tissue-autons autofluorescence images via deep learning. Nature biomedical engineering, 3(6):466-44 477, 2019.
[4] Thorsten Falk, Dominic Mai, Robert Bensch, Özgün Çigek, Ahmed Abdulkadir, Yassine Marrakchi, Anton Böhm, Jan Deubner, Zoe Jäckel, Katharina Seiwald, et U-net: deep learning for cell counting, detection, and morphometry. Nature methods, 16(1):67-70, 2019.
[5] Yang Song, Liyue Shen, Lei Xing, and Stefano Ermon. Solving inverse problems in medi- cal imaging with score-based generative models. In International Conference Learning Representations, 2021. 1,2
[6] Guy Ohayon, Theo Adrai, Gregory Vaksman, Michael Elad, and Peyman Milanfar. High perceptual quality image denoising with a posterior sampling cgan. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 1805-1813, 2021. 2
[7] Bahjat Kawar, Michael Elad, Stefano Ermon, and Jiaming Song. Denoising diffusion restoration models, In Advances in Neural Information Processing Systems, 2022. 9, 32
[8] Yinhuai Wang, Jiwen Yu, and Jian Zhang. Zero-shot image restoration using denoising diffusion null-space model. The Eleventh International Conference Learning Representations, 2023.
2, 32
[9] Hyungjin Chung, Jeongsol Kim, Michael Thompson Mccann, Mare Louis Klasky, and Jong Chul Ye. Diffusion posterior sampling for general noisy inverse problems. In The Eleventh Interna- tional Conference on Learning Representations, 2022.
[10] Wenbo Li. Zhe Lin, Kun Zhou, Lu Qi, Yi Wang, and Jiaya Jia. Mat: Mask-aware transformer for large hole image inpainting. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 10758 (0758-10768, 2022. 9, 32
[11) Andreas Lugmayr, Martin Danelljan, Andres Romero, Fisher Yu, Radu Timofte, and Luc Van Gool. Repaint: Inpainting using denoising diffusion probabilistic models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 11461-11471, 2022. 9,
[12] Chitwan Saharia, Jonathan Ho, William Chan, Tim Salimans, David Fleet, and Mohammad Norouzi. Image super-resolution via iterative refinement. IEEE Transactions on Patlern Analysis and Machine Intelligence, 2022.
[13] Chitwan Saharia, William Chan, Huiwen Chang, Chris Lee, Jonathan Ho, Tim Salimans, David Fleet, and Mohammad Norouzi. Palette: image-to-image diffusion models. In ACM SIGGRAPH 2022 Conference Proceedings, pages 1-10, 2022.

---

[14] Matthew Bendel, Rizwan Ahmad, and Philip Schniter. A regularized conditional gan for postenor sampling in image recovery problems. Advances in Neural Information Processing Systems, 36, 2024. 1,2
[15] Noa Cohen, Hila Manor, Yuval Bahat, and Tomer Michaeli. From posterior sampling to meaningful diversity image restoration. The Twelfth International Conference on Learning Representations, 2024. URL
[16] Vikash Sehwag, Caner Hazirbas, Albert Gordo, Firat Ozgenel, and Cristian Canton. Generating high fidelity data from low-density regions using diffusion models. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pages 11492 11501, 2022.
[17] Elias Nehme, Omer Yair, and Tomer Michaeli. Uncertainty quantification via neural posterior principal components. Advances Neural Information Processing Systems, 36:37128-37141, 2023. 2
[18] Omer Belhasin, Yaniv Romano, Daniel Freedman, Ehud Rivlin, and Michael Elad. Princi- pal uncertainty quantification with spatial correlation for image restoration problems. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2023. 2
[19] Hila Manor and Tomer Michaeli. On the posterior distribution in denoising: Application to uncertainty quantification. The Twelfth International Conference on Learning Representations, 2024. URL http:///eeeeeiie..///ooo/mdddadd
2
[20] Omer Yair, Elias Nehme, and Tomer Michaeli. Uncertainty visualization via low-dimensional posterior projections. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattem Recognition, pages 1001-11051, 2024. 2,
[21] Alex Kendall and Yarin Gal. What uncertainties do need in bayesian deep learning for computer vision? Advances in neural information processing systems, 30, 2017.
[22] Guotai Wang, Wenqi Li, Michael Aertsen, Jan Deprest, Sébastien Ourselin, and Tom Vercauteren. Aleatoric uncertainty estimation with test-time augmentation for medical image segmentation with convolutional neural networks. Neurocomputing, 338:34-55, 2019.
[23] Anastasios N Angelopoulos, Amit P Kohli, Stephen Bates, Michael I Jordan, Jitendra Malik, Thayer Alshaabi, Srigokul Upadhyayula, and Yaniv Romano. Image-to-image regression with distribution free uncertainty quantification and applications in imaging. In International Conference on Machine Learning, 2022.
[24] Andreas Lugmayr, Martin Danelljan, Lue Van Gool, and Radu Timofte. Srflow: Learning the super-resolution space with normalizing flow. In ECCV, 2020. 2
[25] Bahjat Kawar, Gregory Vaksman, and Michael Elad. Snips: Solving noisy inverse problems stochastically. Advances in Neural Information Processing Systems, 34:21757-21769, 2021.
[26] Hyungjin Chung, Suhyeon Lee, and Jong Chul Ye. Fast diffusion sampler for inverse problems by geometric decomposition. arXiv preprint arXiv:2303.055. 055554. 2023.
[27] Berthy T Feng. Jamie Smith, Michael Rubinstein, Huiwen Chang, Katherine L Bouman, and William T Freeman. Score-based diffusion models as principled priors for inverse imaging. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 10520- 10531, 2023. 2
[28] Aapo Hyvärinen and Peter Dayan. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research, 6(4), 2005.
[29] Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsuper- vised learning using nonequilibrium thermodynamics. In International conference on machine learning, pages 2256-2266 PMLR, 2015.
[30] Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution, Advances in neural information processing systems, 32, 2019.
12

---

[31] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-6851, 2020.
[32] Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models.
In International Conference on Learning Representations, 2022. URL https://ppenreviewnndd
Torum"rd=-TldiIlpphoo. 2
[33] Eric Luhman and Troy Luhman. Knowledge distillation in iterative generative models for improved sampling speed. arXiv preprint arXiv:2101.02388. 2021.
[34] Chenlin Meng, Robin Rombach, Ruiqi Gao, Diederik Kingma, Stefano Ermon, Jonathan Ho, and Tim Salimans. On distillation of guided diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 14297-14306, 2023.
[35] Yang Song. Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. In International Conference on Machine Learning, pages 32211 1-22222.... PMLR, 2023. 2
[36] Garoe Dorta, Sara Vicente, Lourdes Agapito, Neill DF Campbell, and Ivor Simpson. Structured uncertainty prediction networks. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 5477-5485, 2018. 2
[37] Miguel Monteiro, Loic Le Folgoc, Daniel Coelho de Castro, Nick Pawlowski, Bernardo Mar- ques, Konstantinos Kamnitsas, Mark van der Wilk, and Ben Glocker. Stochastic segmentation networks: Modelling spatially correlated aleatoric uncertainty. Advances in neural information processing systems, 13:11766-22767, 2020.
[38] Chenlin Meng, Yang Song, Wenzhe Li, and Stefano Ermon, Estimating high order gradients of the data distribution by denoising. Advances Neural Information Processing Systems, 34: 25359-25366 2021.
[39] Frank Nussbaum, Jakob Gawlikowski, and Julia Niebling. Structuring uncertainty for fine- grained sampling in stochastic segmentation networks. Advances Neural Information Pro- cessing Systems, 35:27678 -267911 2022. 2
[40] Abner Guzman Riveer, Dhruv Batra, and Pushmeet Kohli. Multiple choice learning: Learning produce multiple structured outputs. Advances in neural information processing systems, 25, 2012. 3.5
[41] Abner Guzman- Rivera, Pushmeet Kohli, Dhruv Batra, and Rob Rutenbar. Efficiently enforcing diversity in multi-output structured prediction. In Artificial Intelligence and Statistics, pages 282- PMLR, 2014.
[42] Stefan Lee, Senthil Purushwalkam, Michael Cogswell, David Crandall, and Dhruv Batra. Why heads are better than one: Training a diverse ensemble of deep networks. arXiv preprint 15 2015. 3,5
[43] Stefan Lee, Senthil Purushwalkam Shiva Prakash, Michael Cogswell, Viresh Ranjan, David Crandall, and Dhruv Batra. Stochastic multiple choice learning for training diverse deep ensembles. Advances in Neural Information Processing Systems, 29, 2016. 3,5
[44] Kimin Lee, Changho Hwang, KyoungSoo Park, and Jinwoo Shin. Confident multiple choice learning. In International Conference on Machine Learning, pages 2014-2023 PMLR, 2017.
[45] Christian Rupprecht, Iro Laina, Robert DiPietro, Maximilian Baust, Federico Tombari, Nassir Navab, and Gregory D Hager. Learning in uncertain world: Representing ambiguity through multiple hypotheses. In Proceedings of the IEEE international conference on computer vision, pages 3591-3600, 2017. 3, 4, 6
[46] Michael Firman, Neill DF Campbell, Lourdes Agapito, and Gabriel J Brostow. Diversenet: When one right answer is not enough. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 1598-5600 2018.
[47] Eddy 1lg, Ozgun Cicek, Silvio Galesso, Aaron Klein, Osama Makansi, Frank Hutter, and Thomas Brox. Uncertainty estimates and multi-hypotheses networks for optical flow. In Proceedings of the European Conference Computer Vision (ECCV), pages 652-667, 2018.
13

---

[48] Kai Tian, Yi Xu, Shuigeng Zhou, and Jihong Guan. Versatile multiple choice learning and its application to vision computing. Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition, pages 6349 644996657,,00 2019. 3
[49] Nitai Stein and Yaakov Oshman. Double-opportunity estimation via altruism. IEEE Transactions on Aerospace and Electronic Systems, 2022.
[50] Victor Letzelter, Mathieu Fontaine, Mickael Chen, Patrick Pérez, Slim Essid, and Gael Richard. Resilient multiple choice learning: A learned scoring scheme with application to audio scene analysis. In Teirty-seventh Conference on Neural Information Processing Systems, 2023. 3, 3,4,5
[51] Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In Medical Image Computing and Computer-Assisted Intervention-MICCAN 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18, pages 234-241. Springer, 2015. 6,6,1 16
[52] Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convolutional neural networks. Advances in neural information processing systems, 25, 2012.
[53] Mike Brodie. Chris Tensmeyer. Wes Ackerman, and Tony Martinez. Alpha model domination multiple choice learning. In 2018 17th IEEE International Conference on Machine Learning and Applications (ICMLA), pages 879-884. IEEE, 2018. 5
[54] Osama Makansi, Eddy 1lg. Ozgun Cicek, and Thomas Brox. Overcoming limitations of mixture density networks: A sampling and fitting framework for multimodal future prediction. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pages
7144-7153, 2019. 5,
[55] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. International Conference Learning Representations, 3, 2015. 6, 16
[56] Aron Yu and Kristen Grauman. Fine-grained visual comparisons with local learning. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 192- 199, 2014.7 7
[57] Phillip Isola, Jun-Yan Zhu, Tinghui Zhou, and Alexei A Efros. Image-to-image translation with conditional adversarial networks. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 1125-1134, 2017. 7, 16
[58] Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of International Conference Computer Vision (ICCV), December 2015.
7
[59] Lucas von Chamier, Romain F Laine, Johanna Jukkala, Christoph Spahn, Daniel Krentzel, Elias Nehme, Martina Lerche, Sara Hernández-PPrez, Pieta K Mattila, Eleni Karinou, et al. Democratising deep learning for microscopy with zerocostdl4mic. Nature communications, 12 (1):2276, 7
[60] Aaron Van Den Oord, Oriol Vinyals, et al. Neural discrete representation learning. Advances in neural information processing systems, 30, 2017. 10
[61] Ali Razavi, Aaron Van den Oord, and Oriol Vinyals. Generating diverse high-fidelity images with vq- Advances neural information processing systems, 32, 2019. 10
[62] David Perera, Victor Letzelter, Théo Mariotte, Adrien Cortés, Mickael Chen, Slim Essid, and Gaël Richard. Annealed multiple choice learning: Overcoming limitations of winner-takes-all with annealing. arXiv preprint arXiv:2407. 15580, 2024. 18
[63] Steven Diamond and Stephen Boyd. CVXPY: Python-embeddeed modeling language for convex optimization. Journal of Machine Learning Research, 17(83)11-5, 2016. 21
[64] Richard Zhang, Phillip Isola, Alexei A Efros, Eli Shechtman, and Oliver Wang. The unrea- sonable effectiveness of deep features as perceptual metrie. In Proceedings of the IEEE conference computer vision and pattern recognition, pages 586-595, 2018. 23
14

---

[65] Ji Lin, Richard Zhang, Frieder Ganz, Song Han, and Jun- Yan Zhu. Anycost gans for interactive image synthesis and editing. In Proceedings of the IEEE/CVE Conference on Computer Vision and Pattern Recognition, pages 14986-14996, 2021. 23
[66] Qiang Du, Vance Faber, and Max Gunzburger. Centroidal voronoi tessellations: Applications and algorithms. SIAM review, 41(4):637-676, 1999. 25
[67] Yunjey Choi, Youngjung Uh, Jaejun Yoo, and Jung- Woo Ha. Stargan v2: Diverse image synthesis for multiple domains. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pages 81 8188-8197, 2020. 37
15

---

Appendices
A Experimental Details
A.1 Architectures
As mentioned in Section 3.3, in our experiments we adopted the U-Net [51] architecture. Our architecture consisted using average pooling, and upsampling by 2x was implemented with a nearest-neighhor interpolation. In both cases, the feature maps at the updated spatial resolution were processed by 2 convolution blocks consisting of 2D convolution, group normalization, and LeakyReLU activation with negative slope of 0.2. The number of features at each spatial resolution was adapted to the predicted tree layout, Le., for predicting tree of degree K and depth d, the initial number of channels before downsampling was set to Cinit K" where Kd is the number of output leaves and Cinit E (4, 8). At each successive encoder block, the number of channels was doubled, reaching 16 K at the bottleneck. For example, assuming a tree with degree K = 3 and depth d = 2, and setting Cinit 4, the number of channels per level in the encoder is given by [36, 72, 144, 288, 576].
gevnnsampingguusaappingg blocks. Downsampling by 2x was performed
Similarly. in the decoder, the number of channels was halved at each upsampling step in symmetric fashion. However, as mentioned in Section 3.3, in some tasks (e.g., image inpainting). we noticed that fully sharing parameters between all leaves led to reduced diversity between predictions (see Fig. A1). On the other hand, learning a separate U-Net for each prediction leaf is computationally intensive, and highly inefficient as it is expected that the initial feature extraction stage for predicting the different leaves would be similar. Therefore, as a middle ground between these two extremes, we opted for architecture that shares the encoder between the different leaves, while having disjoint decoders, each with dedicated set of weights (<p,) m== learned separately from others. As for the skip connections from the encoder, the concatenated feature maps are interleaved equally per level such that each prediction has an equal share.
In addition to the output images, our architecture also predicts Kd scalars, which are the probabilities of the different leaves. This is achieved by global average pooling of all feature maps from the decoder, and feeding their concatenation to an additional lightweight MLP. This MLP has four linear layers with dimensions [dy, 256, 64, Kd), where d, is the dimension of the concatenated pooled features from the decoder. Each linear layer is followed by a ID batch normalization and SiLUU non-linearity, and the output is passed through softmax layer to produce valid probability vector. Figure A2 summarizes our architecture.
A.2 Per-task Details
In all tasks, we only learned the required residual from the input to produce the predictions. For MNIST experiments, the images were padded to 32 x 32 to enable proper downsampling in the encoder. For Fdges-SShess experiments, the images were kept the same in the pix2pix [57] paper, with a resolution of 256 x 256. For CelehA-HQ experiments, the images were resized to 256 x 256. Finally, for the bioimage translation dataset, we trained on 128 x 128 patches cropped from the full 1024 x 1024 images, as cell information tends to be local.
A.3 Optimization
We used the Adam optimizer [55] with 81 0.9, = 0.999 for all experiments. For the U-Net predicting the output leaves, we used an initial step size of 0.001. For the MLP predicting leaf probability, however, at the initial training phase, the predicted leaves are still not converged; hence, the MLP can easily "classify" which leaf is the most likely. This collapses the predicted probabilities to sparse vector, leading to leaves with zero probability. Hence, to avoid this instability, we used lower step size for the probabilities such that the leaves are first allowed converge, leading to the desired learning dynamics. For both components, the step size was reduced by a factor of 10 if the validation loss stagnated for more than 10 epochs, and the minimum step size was set to 5 10-6, We used a batch size of 32 for 70 epochs for all tasks. This resulted in training times of #40 mins, 10 hrs, 5 hrs, and 7 hrs for MNIST, Edges++Shor CelebA-HO and the bioimage datasets respectively.
found it beneficial to use smaller initial step size of 0.0002. This is because
16

---

Imput
MMSE
p=0.23
p=0.50
p=0 26
p=0.14
0.07
0.03
P=0.1 0.11
p=0.08
p=0.31
0=
p= 0.06
p= 0.00
(a) Fully shared
MMSE
0=0.80
P=0.10
0=0.10
p=
00001
m=0.60
p= 0.06
p=0.03
p=0.01
y=0.02
p=0.0
p= 0.05
(b) Encoder shared
Figure A1: Leaf weight sharing strategy. (a) Fully shared architecture, with all leaves predicted jointly. (b) Leaves only share encoder (see Fig. A2).
Input
5 encodier
Separate decoders
Predicted leaves
Predicted tree
#XMSE
Out Out 3
Global Average Pooling
Predicted probs. dg dry
Input features
Figure A2: Model architecture. Our model receives degraded image y and predicts the bottom K'd leaves and their probabilities (din....Auu ('k.....ka 1K k....kkk (illustrated here for K and = 1). The encoder shared between all leaves, however, the decoding stage, each leaf has a separate decoder (e.g., with parameters (2)49))9Pa) to enable diverse outputs. To predict leaf probability. all features from the decoder are passed through global average pooling layer, and the resulting concatenated vector is passed to lightweight MLP with softmax layer at the output. Afterward, as explained earlier, the output is constructed iteratively from the bottom up (right).
17

---

B Role of E, From Eq. (8)
As explained in Section 3.4, the convergence of MCL is strongly affected by initialization. This is because leaves that are better initialized are more likely to be chosen in the oracle loss, and hence will dominate the remaining leaves in training. This results output leaves that practically do not train, and therefore give meaningless results at test time. To remedy this, we scale the gradients of non-performing predictions with a constant that is annealed during training according to Eq. (8). For the first to = 5, our training degenerates to standard MSE minimization, bringing all leaves prediction
to reasonable starting point near the posterior mean (Fig. A3(c)). This results leaves having roughly the same initialization, with similar number of associated training samples. Afterward, is gradually decayed, and the leaves start to specialize in their respective posterior mode, converging the desired MCL behavior (Fig. A3(d)) In this research, we identified this strategy simple yet effective method of regularization. However, it is important to highlight that a parallel study by Perera et al. [62] presents a potentially improved solution that combines soft assignments with deterministic annealing. This new approach could be easily integrated with our method, and would be intriguing to explore this in future work.
18

---

08
0.22
0.0
0.20
6.18
381SE
0.2
0.16
0.14
10
20
30
15
Rjach
(a) Ee
(b) Distance of leaves [ from the GT a
GT
Inpat
MMSE
5
3 p=
0=0.60
p=o 20
3
p-0.30
0=0.21
pm
p=0.08
p=0.02
p=b.10
p=0.09
pmb.09
p.0.22
3
(c) Predicted tree at epoch f =
MMSE
5 P=0.20
0=0.60
p=0.20
3
5
3
0=0.30
p=0.2
p=0.09
p=0.08
p=0.02
p=o 02
3
3
3
S
S
5
3
3
3
(d) Predicted tree at convergence (epoch t 22)
Figure A3: Role of ex- (a) Et- (b) Distances of leaf predictions {æ::) from the GT throughout
epochs. For the first to = 5 epochs all predictions are brought to the vicinity of the MMSE estimator. Afterward, & decayed and each of the leaves is free to specialize leading to higher distances on average for the entire test set. (c) Predicted tree at epoch t = to = at convergence (epoch f 22)
subset of the posterior,
where all leaves near the MMSE estimator. (d) Predicted where each leaf specializes in a different posterior mode (e.g.. "3" or "5").
19

---

C Weighted Sampler
For imbalanced posteriors where the dominant mode is significantly more likely than the weakest mode (e.g., 10x), we found that an additional regularization complementary to the one provided by E, is needed. This is because even if we eliminate the dependence of training on leaf initialization, still, leaves associated with (much) less likely posterior modes will be chosen with lower frequency during training, resulting in transient gradients that highly affect the adaptive normalization in Adam. To test this hypothesis, we trained two models on the task of image inpainting: (i) one with Adam using an initial learning rate of 0.001, and (ii) one with SGD using a momentum of 0.9 and an initial learning rate of 0.1. Figure A4 demonstrates the results of this experiment. As evident in the resulting trees, Adam leads to highly implausible predictions in leaves with low likelihoods (e.g., p = 0.0 in Fig. A4(c)). SGD, On the other hand, does not lead to nonsensical predictions; however, it requires significantly longer training time and is more challenging to converge (Fig. A4(d))
Therefore, in our method, we proposed a novel non-stochastic weighted sampling scheme as simple fix (Fig. A4(e)). The proposed weighted sampler enables us to still enjoy the convergence speed of Adam while tackling the aforementioned optimization deficiency of MCL training for posteriors that are far from uniform. The goal of our sampler to ensure that on average the number of occurrences at each output leaf during training is roughly the same (Fig. A4(b)) This is done by undersampling training samples associated with leaves that represent probable posterior modes while oversampling training samples associated with leaves that represent rare posterior modes.
Formally, let c; denote the ith prediction leaf (i I..... Kd,, and 8; (zz, yj) denote the jth (paired) training sample. We assume we are given training set D { of N i.i.d. samples, such
that Pa (8,) = . The goal of the sampler is to manipulate ps ($,) via overssappling/undeersmpiiig such that the new sample probability in training (s))) leads to uniform marginal leaf distribution, i.e., ye (cx) R.. To estimate the marginal leaf probability during training, we keep track of an
association matrix A E IRKK xN where Ai counts the number of times sample 8j was associated with leaf q. Specifically, after each training batch b, of size B = [b,|, the association matrix is
updated with momentum j = 1 - 2 &, such that
EMA O,
A.MA UAA. + (1-x)A, Ar,
(AI)
where A, is a binary matrix that encodes the sample-leaf association for the batch b,. This smoothing is necessary accumulate statistics across epochs, making sample probability change gracefully while avoiding large changes to the probability of uncertain samples that switch Ieaf association between batches. Next, given the updated association matrix EMA, can normalize it to obtain an estimate of the current joint distribution Pe,, 88) as
EMA A
Pe,s(Cis 8j)
(A2)
Σp E
K2
N
AEMA
Manipulating the sample probability (s5;) will affect this statistic. Here we can assume that the conditional distribution over the leaves {c]] given a sample 8, stays the same under the change of distribution for the samples. This insight can then be leveraged as follows. First, we obtain the current conditional leaf probability via marginalization
Pe|s(c:(3)) = Pe.a(ci, 8j)
(A3)
ΣK4 (c))
Let P E IBK= denote the matrix where P = Pets (cils)). Recall that our goal is to dictate new sample probability (s,))) such that the induced marginal leaf probability (e(c)) is uniform. Let Qel= E x denote the conditional distribution matrix under probability q such that Q = (c]s()));)) In matrix form, are looking for a probability vector over the samples qa E RR. that satisfies Qe = qf = - I Since the conditional probability of the leaves assumed to be not affected by the change from ps to qs. we have that Q = P. Therefore, in simpler notations, we are searching for a probability vector q E N for which
1 Pq
14x11
(A4)
Kd
20

---

Naturally, since Kd < N, there are infinitely many different options that satisfy this constraint. Note that some of them can correspond an imbalanced distribution, where some samples associated with certain leaf have high probability, while other samples that are associated with the same leaf have low probability. Therefore, among all the different options that satisfy Eq. (A4), we want the one that has the flattest probability. The flattest distribution minimizes the &g loss, namely llqll.- However, minimizing this objective under the constraint Eq. (A4) can be problematic for ill-conditioned cases. Nevertheless, can achieve similar results by solving the following Tikhonov-reegulaiized relaxed problem:
min QER*
Kd 11 1
2
(A5)
s.t.
9j 2 0,
1=1,...,, N
1 =1
One can obtain an exact solution q' to Eq. (AS) using standard methods for convex optimization, e.g. algorithms in CVXPY [63]. However, in typical settings, this results in longer training times. Therefore, we opted to use an approximation. Specifically, we drop the positivity constraint qj 2 0, Vj, and solve Eq. (A5) accounting only for the constraint q 1Nx1 1, we obtain the following closed-form solution
i - (INNN-P' P (PPT + +MKaxK() P) 1NxI
(A6)
dº 1xx1
In practice, in most cases, the positivity constraints are inactive, and therefore dº =q" In the rare cases where Eq. (A6) resulted in negative values, we clipped entries below zero in qº and renormalized the result.
Given the desired sampling distribution (x8)) we pick for the next batch by+1 the distinct samples that their selection will minimize the maximal distance between the future statisties time + 1 and gs- Specifically, at each training step, the next batch of samples chosen such that the maximal
kd Pe.s Pe.a(c,: 88) and
difference between the next iteration estimated sample probability Ps())) the desired sample probability (s,)), Le., |lps gelloo. is minimized. In practice, fetching the samples for future batches is done in advance to speed up the training time. Therefore, to prevent the training from waiting for the successor samples, we use a buffer of m batches ahead and adjust the implementation accordingly. During training, we activated our weighted sampler one epoch after E started decaying to O.
Finally, note that switching from P. 9s during training effectively changes the dataset by repeti- tion/omission, and if not accounted for will distort the learned posterior probabilities. Therefore, to undo this undesired effect while still maintaining an increased number of optimization steps for weak posterior modes, we scaled the loss of sample by the factor % Nq.(,]]. Hence, we effectively only changed the number of optimization steps taken per output leaf, enabling all leaves to continuously receive gradients and train equally as well with Adam.
21

---

(0C0)
E000
Leak oxcurerre
41
41
Leak OcTuEnIre
6000-
8000-
3000
2000
0
Epesds
Ejeoh
(a) Standard sampling
(b) Weighted sampling
GT
Input
MMSE
p=0.08
p=0.52
po
p=0.01
p=0.00
p= 0.07
p=0.0n
p=0.19
p= 0.32
p=0.02
p=0.05
0.33
(c) Adam with standard sampling
MMSE
p 0.08
pim 0.4D
0.01
p-m..o 0.00
0.07
0.000.0
gr 199
pom 0.32
0.022
p- d.33
(d) SGD with standard sampling
MMSE
p= 0.80
p=0.10
p=o.11
p=0.29
p=0.01
p= 0.50
p=0.06
p=0.03
p= 0.01
p=0.02
p=0.03
==0.05
(e) Adam with weighted sampling
Figure A4: Weighted sampling effect. Optimization with Adam requires weighted sampling to train properly (see text in Appendix C).
22

---

D Tree Loss Function and Training Scheme
Our tree training algorithm is summarized Algorithm 1. Throughout this work, we used the MSE loss ((xz 2)) || - =I to both determine the closest tree node and as our optimization loss. However, general, distinction should be made between the measure f(::) used for the clustering (i.e. determining the associations of samples to clusters) and the loss used within each cluster to determine cluster representative.
theory, our approach can work without modification with any association measure (e.g. LPIPS [64], some domain-specifie classifier [65], etc.). However, changing the within-cluster loss requires some modifications. Specifically, the use the MSE loss what provides with the hierarchical decomposition of the posterior mean (sec Eqs. (2) to (7)). In particular, when using the MSE/L2 loss, each cluster representative becomes the posterior cluster mean, and weighted combination of those representatives gives the overall posterior mean (which is the tree root). This allows us to have the network output only the leaves of the tree, which implicitly defines the entire tree (as the nodes of each level are obtained as linear combinations of the nodes of its children).
One natural extension to investigate using approach is to take the association measure to be ((, x) = (I(()) [ where ʃ(.) some relevant domain-spect -ppecfc feature extractor. Specif- ically, we experimented with the deep features of the AnyCostGAN attribute predictor [65], and employed it in the task of eyes inpainting using the CelebAHQ dataset (Fig. A5). This preliminary result indicates that our method could potentially be used with other association metrics. However, further in-depth analysis necessary to ensure semantically meaningful results, which will be addressed in future work.
Algorithm 1 Tree Training
Require: Tree degree K, Tree depth d, Network architecture (:; Training D {(xi, y:) Distance measure E(:, -), Loss weights B E RR+1
1: repeat
/* Sample data and infer leaves +/
2:
(2., ~
D Sample paired training example
3: (á, 0), (y: 0)) 0)
> Infer tree leaves and their probabilities
/+ Bottom-up tree construction using Eqs. (6) and (7) */
4: Run BottomUp(($., ân (y:: 0) K= to get:
(AMsE(3:), ( (vi). ar, (y:) Rk= =
(y:), ax (yi)) kg=
/* Accumulate a along min path */
5:
X f
6:
Φx (y1) <- @MMSE(()) repeat dl times
Initialize closest tree node to root
7:
8:
L +- & U (v:)
D Add closest node to min path
9:
(&x ( + GetChilddeenéxy (y.))
> Get children closest node Update closest node to closest child
10:
K* +- arg min l( 24 (y.)) km I...K K
11:
end
12:
X' t X tg. (3))
> Include also the closest Ieaf
Take gradient step on weighted MSE loss to update */
ا-1
13:
((
= Σ 1 1 1=
14:
0 +- Optimizer O0 VoL)
until converged
23

---

GT
Inpot
P-0.50
p-00.0
P-
0=0.22
p=0.02
p=0.18
(a) Pixel space training ((r, .)
||
P=0.86
p=s0.03
0=0.73
p=0.00
p=0.02
p=0.01
p=0.03
p=0.01
(b) Feature space training E(x, .))
[] (p)
f(t;; 3
Figure A5: Different association metrics ((..).. (a)-(b) Predicted tree in the task of eye inpainting with an MSE association loss in pixel-spaceffeauur-sspae respectively.
24

---

E Centroidal Voronoi Tessellations
Consider domain X' S 12 and p : R+ be weight function over this domain. Additionally, let { K denote any set of K points belonging to X' and let (V)) denote any tessellation of X into K regions. Using these notations, we can define the following loss function,
C (f(2, VI)
Σ i=1 mEV;
p(z))
-x,11 x21/dxx.
(A7)
This formalism arises in various applications such as data compression, optimal quantization, clus- tering, and other applications [66]. In all of these applications, our goal is the same; we want to minimize Eq. (A7) to obtain optimal performance.
A fundamental result in this setting is that for Eq. (A7) to be minimized, it is necessary that {Vi) are the Voronoi regions corresponding to {a; and, simultancously, {x; } are the centroids of the corresponding {= Namely, {(x,, Vi) form a special tessellation space, commonly referred to as Centroidal Voronoi Tessellation (CVT). Specifically, it can be shown that the optimal centroids the resulting CVT are given by
JaEV. xp(ar)da
2; Sæev, p(x)
(A8)
where (V form a Voronoi tessellation of X w.r.t. {xx and are defined
V2 = fx || 11,12 < 2| - E,l12,
Vj *}.
(A9)
In general, the solution to Eq. (A7) is not unique, and no existing algorithm is guaranteed to converge optimal solution. Nevertheless, efficient algorithms like Lloyd's algorithm (Voronoi Iteration)
are known to converge to a local minimum, which frequently produces satisfactory results in practice.
25

---

F Unconditional Hierarchical K K-means
In Section 3.1 we mentioned that in discrete setting, our approach equivalent to a hierarchical K K-means clustering. Here, we describe the hierarchical K-means algorithm and review its properties in the standard (unconditional) setting.
Given a set of N data points { where æ; E R the goal of K K-means is to partition the data points into K clusters/sets (Ci,-.., CK)). such that each data point belongs to the cluster with the nearest mean/centroid serving as a prototype of the cluster. This results in partitioninggtessellatioo of the data space into centroidal Voronoi cells, where for each cluster the within-cluster variances (squared Euclidean distances) are minimized. Formally, the objective function for finding the clusters given by
arg min CI....CK
Σ
ΣER
x 1xl12
(A10)
where 24 is the mean/centroid of cluster Cs. given by PAk T.. Earec, 3;. While in general Problem (A10) is NP-hard, given some initial means/centroidss, the K-means algorithm (also known as Lloyd's method) finds local minimum by alternating between (i) assigning each data point to the cluster with the nearest mean, and (ii) updating the cluster means based on the new assignment.
Extending K-means hierarchically is rather straightforward. Starting from a single cluster comprised of all data points {æ,). at each level of the hierarchy, the data points belonging the kth cluster Cs are split further into K sub-clusters 1C.. Cx,K) } by applying K-means. We refer to this extended algorithm as hierarchical K K-means. The result of applying hierarchical K-means d times is a dendrograam/top-doww balanced tree T of degree K, depth d, and breadth (final number leaves) K'd Note that this successive process is different from applying K-means with Kd clusters once. This is because the hierarchical memberships are enforced, restricting data points in subclusters (C.,1....,,c... from being assigned subclusters .C., Cj.k] if j.
Hierarchical K K-means enjoys several advantages compared to the standard ("Flat") K-means al- gorithm. One of these advantages is increased robustness to initialization demonstrated in Ap- pendix G.3.
26

---

G
Toy Example
G.1 Model Architecture
The model used to produce our results in Fig. 3 consisted of two 5-layer MLPs with 256 hidden units, one for predicting the leaves and one for predicting the likelihoods. Both MLPs consisted of linear layers interleaved with the SiLU activation. In addition, the output of the second MLP was passed through a softmax represent a valid probability distribution.
G.2 Analytical Posterior
In Fig. 3, we plotted the posterior trees on the (unknown) ground truth posterior. Here we provide the closed-form expression for the analytically derived posterior for completeness. The denoising task we assumed was = x + n, where x comes from mixture of L = 4 Gaussians, px(x) = El= (eN(x; p4g, Σt), and
0, is a white Gaussian noise. Specifically, in our toy example we used equally probable spherical Gaussians (i.e., We means
1. Ee I), with the following
(-600 +2.5)
+1.00) (2.5) Ms = (6.0)
-2.
225
11
H2 =
M4 = (-1.5)
(A11)
Let e be an auxiliary random variable taking values in (1....,I, L}} with probabilities [ T Then we can write the posterior by invoking the law of total probability conditioned on the event
c = L
p( y)
Σ ox y. Peeey(()))
8=1
Σ e (x) (z/y, y. (p) 8=1
Py(
Σ Nax; pic. Ee)
qE
(A12)
Et T
t=1
where the first step by Bayes rule, and the result we denoted
qe (3y p4z Σ, 0f I),
μ. - pr E3 Σ - Σ, E. + 1) 'E., 1.
? -1 (g p42)
(A13)
In Fig. A6, in a similar fashion to Fig. from the main text, we visualize posterior trees for additional test inputs yı- As clearly evident in all cases, our method recovers the (approximated) ground truth posterior tree with high accuracy.
G.3 Stability "Flat" vs. Hierarchical K K-means
As mentioned earlier, the hierarchical K-means algorithm has several advantages over the classical "Flat" K-means algorithm. Here, we examine the effect of the random initialization on the resulting clusters found by "Harrarchhhl K-means. In both cases, we applied the respective method 10K posterior samples avoid errors resulting from an insufficient sample size. Figure A7 demonstrates the resulting clusters for different seeds. As can be seen by the result, despite using the widely adopted K + initialization for both algorithms, "Flat" K-means was less resilient to bad initialization compared to its hierarchical counterpart. This is yet another advantage of working with trees of depth bigger than d =
27

---

po (sk)
Aa (A))
Pholati samplo
"Fhi Keaasua
7iM)
V(m)
Tipl
p(a))
p(gn
MM 0.00, 0.23, 0.63, 0.12
(337,035.6334000 0.31
05.4.. 0,08]
Hirnecezl Kemens
piniol
Peeeens angles
The" Kebeen
Tis
Tise
Tiod
M
Mul 0.37, 1.50, JUST,
Mal D106, 0.07, 1.45, D.46
Iieenroooi meas Ours
pielot
p (n)
piela)
ja mm)
pielo)
Kimple
Timl
Tim.
Tiad
pla) pdo 0.255.0200..00
[u23, 0.9 0.27, 0.24]
|a20 0.230.000 0.20 0.25)
(a) "Flat" K K-means
(b) Hier. K-menns
(c) Ours
(d) Clusters
Figure A6: Additional test points y:- Each row above shows the results for different point y: (red dot) when clustering with (a) "Flat" K-means, (b) Hierarchical K-means, (c) Posterior trees (ours). (d) we show the resulting partition/clustringg induced by each method by coloring 10K samples from the posterior Pxly (x(y)) according to their nearest cluster. Our method consistently recovers the result of hierarchical K Kmmenns in all cases. Moreover, other than trivial cases such as the bottom row, the hierarchy better represents weaker modes compared to "Flat" K -eans.
28

---

Seed= 10
Seed=20 20
p(rlpa
pielo
pielos
pirlo)
Pestensc simple
Sed
Tiall
Tisl
Tias
p(gn)
plan)
[L10, 0.200.
pa1 0286,220224 0.31]
Seel
See
"Flat" K-means
Seed=0
Scod=10
Seed=20 20
(A))
pi
MA (A)
pi
M (s)
Pholaii sumple
Saal
(b) Hierarchical K-means
Figure A7: Resilience to initialization. Here we inspect the effect the seed used for initialization, on the resulting clusters recovered by (a) "Flat" K-means and (b) Hierarchical K-means. The hierarchy results a more resilient algorithm that recovers the same result regardless the seed used for initialization.
29

---

H Tree Width vs. Depth
Throughout Section 4 we presented results with trees of degree K 3 and depth d This was done only for the convenience of the exposition and our method is not restricted to this specific layout. Here we examine the effect of tree width compared to tree depth. Figures A8 and presents the resulting trees for four different model configurations on the task of digit inpainting. For fair comparison between model pairs, we fix the number of output leaves which dictates the number of model parameters, and only change the hierarchical structure. Figure A8(a) compares a layout of K 2, d 2. As evident in the result, the hierarchy better organizes the different options, presenting the "4" cluster more faithfully than the flat tree which presents cluster midway between "4" and a "1". Figure A9 repeats this experiment for deeper trees comparing K = 4, d to K = 2, 4. A couple of remarks in order regarding the results. First, in both cases, the "4" cluster has similar probability relative to the "I" cluster. This indicates that our method is consistent in the predicted posterior modes and their likelihoods, with the only change between different layouts being the chosen tessellation the output space. Second, the degree K controls the seliaasisoover-rppeeeeeteee
K = 4,
devoted to weaker posterior modes. A smaller degree leads more emphasis on weaker modes as tree depth grows. Lastly, is important to note that the optimal layout K and is task and input-dependant, and setting these adaptively is an interesting direction for future research.
GT
Input
p=0.69
4
p=0.07
p=0.19
0=0029
p=ó
p=
p50.15
p=
p=0.40
p=b
4
4
4
4
(a) Tree with 4,
(b) Tree with 2,
Figure A8: Flat vs. hierarchical trees. (a) Flat tree with K = 4, = 1. (b) Binary tree with
K 2, 2, with overall 4 leaves. The binary tree categorizes the leaves and better represents the
"4" mode.
30

---

GT
Input
p=o.18
p= 22
p=b.35
4
0.02
0.b6
0.04
0.07
0.11
0.b3
0.10
0.04
0.03
0.05
0.
0.b6
0.24
0.07
0.01
0.03
4
4
4
4
(a) Tree with 4, d
MMSE
0=2.33
0=0.77
4
11
p=0.12
0=0.38 0.38
p= 0.39
4
4
p=0.00
P=0.07
p=0.00 0.06
p=0.06
p= 0.24
p=0.14
p=0.14 0.14
0=
4
4
4
4
0.03
0.02
0.04
0.03
0.
0.b4
0.b4
0.10
0.19
0.08
4
4
4
4
4
4
4
4
(b) Tree with K = 2, = 4
Figure A9: Tree depth vs. width. (a) Tree with 4, d = (b) Binary tree with K = 2, d also resulting in 16 leaves. The binary tree emphasizes more the "4" mode relative to the "I" mode, although both cases the probability mass associated with the "4" 20%.
31

---

I GAN-based Posterior Samplers
In Section 4 we compared our method mainly to diffusion-based posterior samplers (e.g., 11D. While these samplers often lead to state-of-the-ar sample quality, they are known to be computationally intensive. theory, GAN-based samplers such as MAT [10], can significantly speed up the proposed two-step baseline of sampling followed by hierarchical K K-means; however, as evident Fig. A10, even top-perforning GAN-based samplers often collapse to single mode the posterior, producing samples with little variability that do not faithfully reflect the full distribution.
GT
Inpot
0.17
p=0..4
p-u.
0=0.44
p=0.06 06
p=0007
p=0.16
p=0.08
p=0.17
0=0.15
p= 12
1=0.16
(a) Mouth inpainting
GT
Impul
MMSE
27
p=0.58 0.58
P=Ù.1 15
0=0.04
0=0.11
m=0.12
P=0.16
p=0.18
P=0.26
m=00.02
0=0..4 04
p=0.00 .08
(b) Eyes inpainting
Figure A10: GAN-based posterior sampler. (a)-(b) Posterior trees constructed with MAT, by applying hierarchical K-means to 100 samples in the tasks of mouth/eyes inpainting respectively. As evident in both cases, the resulting tree exhibits little to no variance due to mode collapse.
32

---

Visual Comparison to Baselines
In Table we reported quantitative comparisons the proposed baseline implemented with various samplers. In Figs. A11 and A12, we include two sample visual comparisons of the resulting trees with our method and the baselines for the tasks of face colorization and mouth inpainting.
GT
Inpat
MMSE
0..49
pmo.41
p=0.10
p= 0.23
p=0.00
0.22
p=
0.10
p=b.17
gm0.06
pm0.03
p=0.01
(a) DDNM
MMSE
p=0.41
p=0.46
P=0.15
P=0.19
p=0.12
p=
p=0.12
p=0.22
p=0.11
p=
p=o
p=0.01
(b) DDRM
p=0.48
P=0.24 24
p=0.19
p=
p=0.11
p=0.20 2020
p=0.06
p=0.02
p=0.14
p=o.05 0.0
p=0.05 0.05
(c) Ours
Figure All: Tree comparison in colorization.
33

---

GT
Input
p=0.47
p=0.20
p= 33
p=b..7
p-b.17
p-b.
p=b..07
p-m0.03
p-b..
p=d.11
p=d.11
(a) DDNM
p-0030
p=0.22
p=d.41
p=0.09
p=0.17
p=0.04
p=o.12
p=0.12
p=0.05
p=0.09
p=o. 20
p=0.12
RePaint
MMSE
0.28
p=0.41
p=
p=
p=0.12
p=0.06
p=0.15
p=0.12
p=m.14
p=0.05
p=d.14
p=d.13
(c) MAT
MMSE
p=2.27
p=0.41
p=0.33
p=0.07
p=0.14
p=0.06
p=0.10
p=0.12
p=0.19
p=
p=o.16
p=0.08
(d) Ours
Figure A12: Tree comparison in mouth inpainting.
34

---

K Runtime Comparison as a Function of Batch Size
In Table I reported the runtime and memory footprint assuming a batch of size However, in GPU computing, memory footprint and hardware parallelization ultimately determine the overall runtime. We tried refrain from factoring these in our calculation in Table as these numbers are hardware-specific depending on the available GPU. Nonetheless, it is important to note that our method is not only faster than the competing approaches with batch size I, it also has lower memory footprint and is therefore more amenable to parallelization. Table 2 benchmarks the speed of the forward pass and the GPU memory usage as function of the batch size on an A6000 GPU with 48 GB. For cach method, we tested three different batch sizes: 10, and the maximal number of samples that in memory. The forward pass speed and the memory usage each setting were tested 100 times using CUDA events (reported as mean + std). Our method is extremely fast and parallclizable, enabling the inference of 430 test images with a single forward pass of 1.15 seconds.
Even in the case of single test image, our method is far superior. For simplicity let us assume the cost of running hierarchical K-means negligible, and that we can squeeze batch size of 100 samples for the diffusion-based samplers, and ">33.33" samples for MAT. Sampling N, = 100 posterior samples with DDRM (fastest diffusion baseline, requires only 20 denoising steps) lasts 46.3 seconds, and with MAT lasts 7.8 seconds. In comparison, our method requires only 0.014 seconds, leading to a 557 speedup compared the fastest baseline.
Runtime comparison as a function of batch size on an Nvidia A6000 GPU. Note that unlike Table here we report the speed of a forward pass separately from the number of required forward passes to infer posterior tree for a single test image.
Table 2:
Batch
Forward pass
Memory (GB)
Forward passes
Total time per tree (sec)
Method
Size
speed (ms)
per tree
34:0.4
1.85
68-15,504
2,000- 456,000
Diffusion
10
24.510.1
7.5
49-11,172
97
23.9++0.6
46.8
48-10,89
1504:0.8
3.0
15
MAT
10
86.7+1.2
16.0
100
8.67
30
86.6:2.55
47.2
8.66
14:00.2
1.3
0.014
Ours
10
3.9+0.1
2.1
1
0.0039
430
2.6710.1
46.9
0.00267
35

---

L Baselines Performance as a Function of Compute
As mentioned in Section 4.1, baseline trees were computed using a two-step procedure: (1) Sampling N, images from the posterior, and (2) performing hierarchical K- --means d times to build a tree of degree K and depth d. Therefore, saving computation in this process requires sampling fewer images N, per test imput. However, in our experiments with K = 3, d = 2 (i.e. a tree with 9 leaves), we noticed that using less than N,- 100 images often led to degenerate trees with one or more leaves having I sample or less. For example, consider the case of using N. 9 samples. Even if the posterior is perfectly balanced (often not the case), each leaf will have only sample to "average" at depth 2. Figure A13 plots the success probability of building tree with Kd 9 leaves function of the number of sampled images N, in the different tested tasks. A tree is considered to be successfully constructed if each of its leaves has at least 2 samples average. As evident in the results (e.g., Fig. Al3()), it is likely that to achieve the optimal performance with the baselines more than N, 100 samples are needed, which would be even more computationally demanding. Moreover, the number of samples needed to successfully construct trees with more than 9 leaves is expected to be significantly larger rendering this baseline impractical, especially when considering sizable test set.
100
DDRM
DDNM
Surcess rate [%]]
60
40
20
0
20
40
60
80
100
Number of posterior samules (N.)
(a) Colorization
100
80
Sueness rate [%]]
60-
40
DDNM
20
RePaint
MAT
0:
20
80
100
Number posterior samples (N.)
(b) Eyes inpainting
Figure A13: Tree success rate vs. number of posterior samples N.. (a)-(b) Success rate of building a (posterior) tree by applying hierarchical K K-means to N, samples the tasks of colorization and eyes inpainting respectively.
36

---

M More Results
Here we include more results for each of the tasks presented in Section from the main text. Additionally, we also include colorization results on the AFHQ(v2) dataset [67], with similar performance that obtained on CelebA-HQ.
GT
Inpat
MMSE
5
5
p= 0.60
P=0.20
p=0.20
S
5
3
p=0.30
p=0.21
p
p=0.08
p=0.02
P=0.10
y=0.09
p=t.09
p=o 0.02
S
S
5
5
3
5
3
3
3
GT
Input
MMSE
4
9
0=0.27
P=0.30
0=0.43
9
7
4
p=0.13
p=0.06
p=d.08
p=d.11
p=0.16
p=b.04
p=0.08
p=ú.13
p=o
9
9
9
7
7
?
4
4
GT
Input
4
4
0=0.35
4
4
p-0.15
p=t.11
ped..0
p=o.11
p=b.us
p-ú.12
p=d.16
4
4
4
4
4
4
Figure A14: More digit inpainting results.
37

---

GT
Inpot
MMSE
p=
p=0.94
p=o 02
p=0.04
p=0.00
p=d.00
p=d.99 ..
p=0.00
p=b.00
p=
p=0.00
p=0.01
GT
Input
p=0.21
p=0.75
p=d 04
0.20
p=0.01
pmm0.0
d.09
p=0.0G
p=b.00
p=ú.00
p=0.01
GT
Inpat
p=0.13
p=0.00
m=0.00
p=0.08
p=0.04
p=0.00
p=0.09
p=0.04
p=0.08
p=0.66
Figure A15: More Edges- +Shoes results.
38

---

GT
Input
MMSE
p=0.335
0.33
p=o
p=0.07
p=0.20
p=d.08
pad
p=0.09
p=b.27
p=ú.11
p=0.00
GT
Input
p=0.41
p=o
p.00.3
gmm0.04
p=d.32
p=0.25
p=0.4
0.00.01
pm0.05
p=0.01 0.01
GT
Inpot
p=0.43
0.04
0=0.15
p=00.14
p=0.06
p=0.36
p=0.01
0=0.02
p=0.16
p=0.06
Figure A16: AFHQ(y2) colorization results.
39

---

GT
Input
MMSE
p=0.884
p=0.08
0=0.07
p=0.06
p=b.02
p=o
GT
Inpot
p=0.79
P=0.11
p=0
0.20
p=d.04
p-0.06
p=b.01
p-0.07
p=0.02 0.01
GT
Inpat
p
p=0.32
p=0.19
0=0.17
p=0.19
p=0.23
p=0.06
p=0.03
y=0.08
p=0.03
p=0.02
P
Figure A17: More CelebA-HQ colorization results.
40

---

GT
Input
MMSE
p= 0.00
p=0.50
0=0.40
p=0.07
padd
p=0.15
b=b.29
p=ú.17
p=0.10
GT
Input
p=0.8
P=0.28
P=0.35
0.14
p=0.00
p=d.08
p=0.09
p-0.12
0.11
GT
Input
p=0.21
0.20
0=0.000
p=0.22
p=0.10
p=0.08
y==0.16
n=0.07
p=0.16
Figure A18: More mouth inpainting results.
41

---

GT
Input
MMSE
p=0.50
p=0.09
P=
p=0.18
p=0.11
p=0.02
p=b.05
p=0.02
p=0.18
GT
Input
P=0.32
P=O.32
ped.08
p=0.11
p=0.13
p=0.112
p=t.11
p=d
GT
Inpot
p=0.70
0.05
0=0.05
0=0..3
p=0.18
p=0.30
p=0.22
0=00.08
p=0.03
0=0.07
Figure A19: More eyes inpainting results.
42

---

GT
Input
MMSE
P=0.10
p=o 02
p=0.42
p=0.00
p=d.48
p=d.00
p=s0.05
p=b.04
p=0.001
p=u..0
p=d.00
GT
Input
p=0.02
p=0.90
p=o 00
p=0.00
p=0.000
p.00.00
ped.01
p=o.00
p=b.98
go00.00
p=ú.00
p=d.00
GT
Inpot
p=0.16
p=00.06
0.03
0=0.00
p=00.75
p=0.03
p=0.11
p=0.02
y==0.00
p=0.04
p=0.02
Figure A20: More Bioimage translation results.
43

---

NeurIPS Paper Checklist
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The claims made in the abstract and introduction are reflected in our theoretical and experimental results found in Sections and
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We discuss the limitations of method in Section 5.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
44

---

Our results in Section 3 are followed by proper theoretical justification where applicable.
Justification:
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: Both the main text (Section 3) and the appendix (Appendices A, and G) provide all essential details to replicate our experimental results.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
45

---

Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: All datasets used in the paper publicly available and can be retrieved from the cited papers. Our code is made publicly available on GitHub.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please the NeurlPS code and data submission guidelines (Nttps://nips.cc/mbblic guidessooddSubiissioooollyy for more details.
While encourage the release of code and data, understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce results. See the NeurlPS code and data submission guidelines (https: httpsc/cpbbbccgguiss/CCoooo Sawpuddse/CooSStbmbiiioooiiioioo for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Our settings are discussed in detail the main text (Section 4) and Ap- pendix A.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: In Table we report the standard error of the mean over the test images for both MSE and NLL.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
46

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
Question:
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Answer: [Yes]
The per-task training times are detailed in Appendix A.3. terms inference times and memory usage, discussed these extensively in Tables and 2 and Appendices K and L..
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification: the paper fully conforms with the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g.. if there is special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive socictal impacts and negative socictal impacts of the work performed?
Answer: [Yes]
Justification: In Section we discussed the potential positive societal impact of our method, and in Section we demonstrated practical use for properly visualizing reconstruction uncertainty in -afity-critical tasks such bioimage translation.
47

---

Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considcrations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning conrectly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: This paper poses such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model. for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: All dataset papers are properly cited. Licenses are as follow: MNIST (CC BY-SA 3.0). Edges (CC BY-NC 4), AFHQ(v2) (CC BY-NC 4), CelebA-HQ (CC BY-NC 4), Bioimage Translation (CC BY 4).
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
48

---

The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswitheode. com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: The paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
49

---

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
50


Reviews:
[{"review_id": "yTgem9AFEu", "reviewer": "Reviewer_4WsD", "summary": "This work proposes to solve the problem of quantifying and visualising the uncertainty in the solutions of ill-posed inverse problems like image-to-image translation, image re-construction, inpainting, etc. The paper proposes to do so by using 'Posterior-trees' where the authors make use of the result that optimizers of Eq. (1) form CVT of the posterior. Further application of simple Bayes rule and simple probability allows for construction of a hierarchal tree to quantify as well as visualize the uncertainty associated with the predicted solution of the inverse problem at hand.", "questions": "See weaknesses", "limitations": "See weaknesses", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is well written with clear ideas and rationale.\n2. The idea of using hierarchal trees to quantify and visualize the uncertainty associated with the posterior of inverse problem is indeed novel. Although the core idea on which the paper builds upon is already presented in [45].\n3. Nonetheless, extending the results of [45] to construct a heirarchal version of it is notable.\n4. The results and visualizations are satisfying.", "weaknesses": "1. I think the authors should provide more elaborate background on CVT and results of [45], either in main text or in supplementary for ease of the reader.\n2. One thing that I did not understand is - what is the trade-off between breadth and depth of the constructed tree? For example - what is the difference between having more children nodes with low depth and having less children nodes with high depth. This would help in further understanding the advantages and limitations of the proposed method.\n3. Another point that I am a bit skeptical about is that the method operates directly in the image space (or the space of the input of the problem at hand). In case of images, this space is very high-dimensional. In such high-dimensional spaces, capturing all the variations of the solution specially through a discrete tree-like structure is problematic. I don't know how the proposed method is able to handle it.\n\nOverall I enjoyed reading the paper, and the paper certainly seems to be novel, albeit, incrementally. Hence, I lean for weak acceptance.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Providing more background on CVT and the results of \\[45\\]**\n\nWe thank the reviewer for bringing this to our attention. In the camera-ready version, we will include an additional appendix summarizing the main results of \\[45\\] upon which we build.\n\n**Tradeoff between breadth and depth of the constructed trees**\n\nKindly note that we discussed and visualized this tradeoff for the task of digit inpainting in Appendix F and Figs. A7-A8. We apologize for not referring to this appendix from the main text. We’ll add a reference, thanks for catching this.\n\nIn short, the choice of $K$ and $d$ determines the layout of the output space tessellation. The degree $K$ controls the emphasis/over-representation devoted to weaker posterior modes. A smaller degree leads to more emphasis on weaker modes of the posterior as tree depth $d$ grows. However, the optimal layout $K$ and $d$ is task and input-dependant, and setting these adaptively is an interesting direction for future research.\n\n**Number of prototypes and operating in input/pixel space**\n\nThis is a very good point. Please note that ultimately the goal is to present the user with only a **few** representative prototypes **summarizing** the different options. Otherwise, navigating the underlying possibilities becomes tedious and time-consuming as for every input the user has to skim through many many images. In our case, we chose these prototypes to be the cluster centers in pixel space at different levels of granularity. However, as we mentioned in Section 5 (Discussion and Conclusion), averaging in pixel space might lead to off-manifold “blurry” reconstructions, which could (possibly) be tackled by applying posterior trees in the latent space of an autoencoder such as VQ-VAE \\- an intriguing direction for future work.\n\n**Novelty**\n\nWhile our approach is conceptually indeed a hierarchical extension of \\[45\\], as the reviewer noted, our work has several non-trivial contributions compared to \\[45\\], some of which are unique to the setting of working with tree structures:\n\n* Proposing an efficient architecture that enables predicting **diverse** solutions (and their likelihoods) compared to the fully shared (likelihood-free) architecture of \\[45\\], as we verify in Appendix A (Figs. A1-A2).  \n* Preventing tree collapse and stabilizing the optimization steps through a novel adaptive regularization scheme and a weighted sampler that ensures proper leaf optimization with Adam (Appendices B-C, Figs. A3-A4).  \n* Demonstrating, to the best of our knowledge, the first successful application of MCL with diverse results for high-resolution image-to-image regression tasks.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "WhvWhTadDO", "reviewer": "Reviewer_7Dwv", "summary": "This work proposes a technique to predict a tree-structured hierarchical summarization of a posterior distribution using a single forward pass of a neural network. The technique is an amortized hierarchical version of the oracle loss in multiple choice learning. Experiments show the method is effective at hierarchically representing the posterior both qualitatively and quantitatively.", "questions": "1. Can you apply this method when the reconstruction loss is not an $L_2$ loss but a generic function?\n2. How does the runtimes compared in Table 1? \n3. Can you show how the methods compare as a function of runtime in Table 1? It's possible that you don't need to run that many function evals for the baselines to get similar performance.", "limitations": "The authors discussed limitations of the paper.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "1. The method is simple, efficient, and sound. \n2. It addresses a fundamental and practically relevant problem (hierarchically visualizing complex distributions).\n3. The presentation of the method and the experiments is excellent.", "weaknesses": "1. Table 1 reports number of function evaluations, but it's not clear how that translate to runtime since the time per function call and extent of parallelization varies .\n2. The qualitative evaluation in Figure 1 and Figure 4 are not very informative. It's often hard to tell the differences between many of the nodes or judge whether there is an unambiguous underlying hierarchical structure.\n3. The clustering / hierarchy produced the method is based on $L_2$ distance, which is not a very useful metric in many applications including images. Though as the authors discussed, this can be potentially addressed e.g. by first embedding the inputs with an autoencoder. Empirically demonstrating that this limitation can be addressed would be important for adoption of this method in many applications.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Comparing runtime in neural function evaluations vs GPU seconds**  \nPlease see Table 1 in the rebuttal PDF where runtime is reported in seconds. As mentioned in checklist item 8 (Experiments Compute Resources), we reported neural function evaluations (NFEs) as an architecture-agnostic measure. However, our architecture is significantly lighter than the U-net used in the compared posterior sampling baselines (both in terms of compute and memory footprint). Therefore, the numbers in Table 1 are actually underestimating the speedup factor introduced by our method. Putting aside the lower memory footprint, a single forward pass with a batch of 1 on an A6000 GPU with our architecture lasts 5$\\\\pm$0.1 ms, compared to 20$\\\\pm$0.4 ms for the U-Net used by DDRM/DDNM/RePaint, and 140$\\\\pm$8 ms for MAT. To ensure this is clear enough, in the camera-ready version we will include another column in Table 1 translating NFEs to runtime in seconds further emphasizing the speed advantage of posterior trees.\n\n**Baselines performance as a function of compute**  \nNote that to compute the baseline trees we need to perform a two-step procedure: (1) Sampling $N\\_s=100$ images from the posterior, and (2) performing hierarchical $K$-means $d$ times to build a tree of degree $K$ and depth $d$. Therefore, saving computation in this process requires sampling fewer images $N\\_s$ per test input. However, in our experiments with $K=3,d=2$ (i.e. a tree with 9 leaves), we noticed that using less than $N\\_s=100$ images often led to degenerate trees with one or more leaves having 1 sample or less. For example, consider the case of using $N\\_s=9$ samples. Even if the posterior is perfectly balanced (often not the case), each leaf will have only 1 sample to “average” at depth 2\\. In fact, it is likely that to achieve the optimal performance with the baselines more than $N\\_s=100$ samples are needed, which would be even more computationally demanding. To better clarify this point, we will include an additional appendix discussing the success probability of building a tree with $K^d=9$ leaves as a function of the number of sampled images $N\\_s$.\n\n**Learning posterior trees with other loss functions**  \nThis is a great point. A distinction should be made between the measure used for the clustering (i.e. determining the associations of samples to clusters) and the loss used within each cluster to determine the cluster representative. Our approach can work without modification with any association measure (e.g. LPIPS, some domain-specific classifier, etc.). However, changing the within-cluster loss requires some modifications. Specifically, the use of the $L\\_2$ loss is what provides us with the hierarchical decomposition of the posterior mean (see Eqs. (2)-(7)). In particular, when using the $L\\_2$ loss, each cluster representative becomes the posterior cluster mean, and a weighted combination of those representatives gives the overall posterior mean (which is the tree root). This allows us to have the network output only the leaves of the tree, which implicitly define the entire tree (as the nodes of each level are obtained as linear combinations of the nodes of its children).\n\nWe will add an appendix discussing the distinction between association losses and within-cluster losses, and will include an illustration of using an association loss that is not $L\\_2$. As for generalizing the method to work with other within-cluster losses, this is a great avenue, which we leave for future research.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "fKMfOEzsp0", "reviewer": "Reviewer_VHHJ", "summary": "The paper proposes a method to obtain a hierarchical representation of samples from the posterior distribution of inverse problems. This method can be integrated into any existing approach to learn and sample from the posterior distribution. It involves learning a tree structure where each node represents a prototype reconstruction, with child nodes refining their parent node. This allows for the exploration of different modes and levels of detail in the solution space.", "questions": "* When augmenting existing methods for learning conditional generative models (for learning the posterior), how does the loss associated with the tree structure affect the learning of the generative model? How do the authors prevent the proposed loss from negatively biasing the learning of the posterior distribution?\n\n* What is the computational overhead of the proposed method compared to simply training a generative model to learn the posterior distribution? What is the cost of post-hoc hierarchical clustering? How many inverse problems need to be solved to justify the extra training time?", "limitations": "Authors provide a discussion on the limitations of their\nmethod.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "**Originality.** While the method has connections to post-hoc hierarchical clustering of samples, it is novel in that it allows for faster visualization of the posterior distribution samples at test time.\n\n**Quality and clarity.** Overall, the paper is well-structured, easy to follow, and sufficiently motivates the problem. However, the clarity of the presentation of the proposed loss function could be improved.\n\n**Significance.** The problem of informatively visualizing the variability in samples from the posterior distribution of inverse problems is impactful.", "weaknesses": "* While I truly appreciate the importance of informatively visualizing samples from the posterior distribution, the proposed method seems to require several heuristics for training to yield desirable results (e.g., discussions in sections 3.3 and 3.4). Are the computational complexity gains of the proposed method, compared to performing post-hoc hierarchical clustering, worth the need for tuning the parameters in these heuristics?\n\n* It seems like the architecture of the network required to predict the tree depends on the tree depth and width as $K^d$, where $d$ is the tree depth and $K$ is the branching factor. This appears to be a very limiting factor in the depth and width of the tree that can be learned. How does the cost of learning the tree structure as we scale $K$ and $d$ compare to the cost of post-hoc hierarchical clustering?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Tree loss function and training scheme**   \nThanks for pointing out this is not clear enough. We will include an appendix with the explicit training algorithm to better clarify the loss function and our overall training scheme. Kindly note that we trained our models from scratch on a standard dataset of a **single** posterior sample per input, and did not augment existing conditional generative models. After training with a hierarchical MCL loss, for every input, our model learns to predict a tree that hierarchically discretizes the posterior distribution. In our experiments, the various posterior sampling baselines are only used in the comparisons to benchmark our results and quantify the quality of the predicted trees.\n\n**Runtime comparisons**  \nTable 1 compares the speed of our method and the baselines at test time, as ultimately this is what matters to the user. Our method is at least 2 orders of magnitude faster than the baseline in neural function evaluations (NFEs). However, as we detail next, our method is also orders of magnitude faster than the baselines in training time. \n\n* **Training time**: The diffusion-based baselines we compared against are all zero-shot methods that rely on a pre-trained unconditional diffusion prior (DDPM). Hence, the “effective” training time of these methods is the time it took to learn the DDPM prior, which for CelebA-HQ $256\\\\times 256$, Ho et al. stated it was $\\\\approx$ 63 hours on **8** V100 GPUs. For comparison, training our method (from scratch) on a CelebA-HQ $256\\\\times 256$ restoration task took $\\\\approx$5 hours on a **single** A6000 GPU (see Appendix A.3).\n\n* **Testing time**: As reported in Table 1, our method is at least 2 orders of magnitude faster than the baselines.  As mentioned in checklist item 8 (Experiments Compute Resources), we reported neural function evaluations (NFEs) as an architecture-agnostic measure. However, our architecture is significantly lighter than the U-net used in the compared posterior sampling baselines (both in terms of compute and memory footprint). Therefore, the numbers in Table 1 are actually underestimating the speedup factor introduced by our method. Putting aside the lower memory footprint, a single forward pass with a batch of 1 on an A6000 GPU with our architecture lasts 5$\\\\pm$0.1 ms, compared to 20$\\\\pm$0.4 ms for the U-Net used by DDRM/DDNM/RePaint, and 140$\\\\pm$8 ms for MAT. To ensure this is clear enough, in the camera-ready version we will include another column in Table 1 translating NFEs to runtime in seconds further emphasizing the speed advantage of posterior trees. (The rebuttal PDF includes Table 1 with runtime reported in GPU seconds).\n\n**Required hyper-parameter tuning to stabilize training**  \nAs mentioned in Section 3.4 L179, we observed that fixing the hyperparameters $\\\\varepsilon\\_0=1,t\\_0=5$ worked well for all experiments, and did not need to optimize these further. In general, it might be that for some tasks this strategy is suboptimal, requiring more careful tuning. Nonetheless, given that our training time is relatively short, this added computational burden is manageable, and it is still beneficial to use our method due to its much faster inference at test time. \n\n**Scaling to wider/deeper trees**  \nIndeed, as mentioned in Section 5 (Discussion and Conclusion), our method is limited by the number of output leaves as we amortize the inference of the entire tree to a single forward pass. However, note that this limitation also exists for the baselines. To compute the baseline trees, we need to perform a two-step procedure: (1) Sampling $N\\_s=100$ images from the posterior, and (2) performing hierarchical $K$-means $d$ times to build a tree of degree $K$ and depth $d$. In our experiments with $K=3,d=2$ (i.e. a tree with 9 leaves), we noticed that using less than $N\\_s=100$ images often led to degenerate trees with one or more leaves having 1 sample or less. For example, consider the case of using $N\\_s=9$ samples. Even if the posterior is perfectly balanced (often not the case), each leaf will have only 1 sample to “average” at depth 2\\. Therefore, if we want to use the baselines for trees with significantly more leaves, we need to sample enough images ($N\\_s\\\\gt 100$) per test input which scales poorly as well. \n\nUltimately, the goal is to present the user with only a **few** representative prototypes **summarizing** the different options. Otherwise, navigating the underlying possibilities becomes tedious and time-consuming as the user has to skim through many images for every input.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
Updated runtime table in seconds

Meta-review: The paper presents a novel method for quantifying and visualizing uncertainty in solutions to ill-posed inverse problems using a tree-valued hierarchical summarization of the posterior distribution, achieved in a single forward pass of a neural network, allowing for efficient and granular exploration of solution spaces. Reviewers appreciated the paper's excellent presentation, the novelty and practical relevance of the hierarchical approach to visualizing complex distributions, and the method is qualitatively and quantitatively effective. Reviewers noted the need for more background on foundational concepts, concerns about the method's handling of high-dimensional data, potential limitations in scaling the tree depth and width due to architectural constraints, and the dependency on several heuristics for optimal performance. The reviewers responded positively to the authors' replies and unanimously recommended that the paper be accepted. Therefore, we accept the paper.
Decision: Accept (poster)