Hierarchical Uncertainty Exploration via Feedforward Posterior Trees

Authors: elias  nehme
rotem  mulayoff
tomer  michaeli
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 13:54:59

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/).

Main paper: 
 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 a given input can typically be associated with a multitude of plausible solutions. In such cases, it is advantageous to equip users with efficient tools for exploring and visualizing the set of admissible solutions. Such tools may be of high value especially in safety-critical domains, such as scientific and medical image analysis [1] [2] [3] [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, high-dimensional posteriors are challenging to visualize. One popular approach for communicating uncertainty is to generate samples from the posterior [5] [6] [7] [8] [9] [10] [11] [12] [13] [14] . Yet, in complex domains with high uncertainty levels, users may need to 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 set of samples that highlight posterior diversity [15, 16] . However, choosing this set to be small is often insufficient for summarizing the posterior, while choosing it to be large hinders the user's ability to properly 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] [18] [19] [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 a hypothesis about the unknown image. Specifically, our model receives a 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 of posterior uncertainty (Fig. 1 ).Our method, which we coin posterior trees, is a generic technique for uncertainty visualization that is seamlessly transferable across tasks and datasets. Compared to the popular practice of training point estimators with mean square error (MSE) minimization, our method is an ideal drop-in replacement. This is because it adds a minor additional computational burden, while simultaneously accompanying the point estimate with its decomposition to the constituting clusters, revealing the underlying modes of variation and disclosing posterior uncertainty.We showcase our method on multiple inverse problems in imaging, demonstrating the practical benefit of tree-valued predictions. We also quantitatively compare our learned posterior trees to a two-step baseline that generates samples using a diffusion-based 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. 
related work: 
 The common avenue for probing data uncertainty (also known as aleatoric uncertainty [6] [7] [8] [21] [22] [23] ) involves posterior sampling using conditional generative models like [5] [6] [7] [8] [9] [10] [11] [12] [13] [14] [24] [25] [26] [27] , with the dominance of score-based/diffusion models [28] [29] [30] [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-art models. Despite attempts to enhance speed [32] [33] [34] [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] [37] [38] [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] [41] [42] [43] [44] [45] [46] [47] [48] [49] [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-Rivera 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-all strategy encourages each head to specialize in a 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-hypothesis 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. 
method: 
 Our goal is to predict a clean signal x ∈ X ⊆ R dx , given some degraded measurements y ∈ Y ⊆ R dy . We assume that x and y are realizations of random vectors x and y with an unknown joint distribution p x,y (x, y), from which we have a training set of i.i.d. samples D = {(x i , y i )} N i=1 . The objective of our predicted trees is therefore to visualize the uncertainty in the posterior distribution p x|y (x|y) via a hierarchical set of a few prototypes.Many image restoration methods output a single prediction x for any given input y. A common choice is to aim for the posterior mean x = E[x|y = 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 { x1 , . . . , xK } that portray the different options. In MCL, this is achieved by minimizing the so-called oracle/winner-takes-all loss given byEQUATIONwhere ℓ(•, •) is a loss function that measures prediction distance. Note that L MCL is minimized as long as one restoration in the predicted set is close to x. This encourages prediction diversity, such that each restoration specializes in a different mode of the posterior. Rupprecht et al. [45] proposed a probabilistic interpretation of MCL, showing that the optimal minimizers of Eq. ( 1) form a centroidal Voronoi tessellation (CVT) of the posterior, where { x⋆ i } K i=1 are given by the conditional means/centroids of the resulting K Voronoi cells (see Appendix E). In a discrete setting where we are given samples from p x|y (x|y), CVT with the Euclidean distance is equivalent to K-means clustering.Our goal here is to generalize the concept of predicting a set of restorations. Given an input y, instead of predicting an unordered set { x1 (y), . . . , xK (y)}, we propose to output an input-adaptive tree T (y) representing a hierarchical clustering of the posterior p x|y (x|y). Specifically, we want the tree nodes to correspond to a hierarchical CVT of the posterior (analogous to a hierarchical K-means process in the discrete setting), where the children of each node constitute a 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.As discussed above, training a model to predict a tree of degree K and depth d = 1 with the oracle loss, results in a CVT of the posterior p x|y (x|y) for every y. Let {A k (y)} K k=1 denote the resulting Voronoi cells at d = 1 for some input y = y, such that ∪ K k=1 A k (y) = X , with X ∈ R dx being the Next, the tree is iteratively constructed from the bottom up using weighted averaging, until we reach the root node which is the minimum MSE predictor xMMSE . During training, starting from the root, the ground truth x is propagated through the tree until it reaches the leaves (dashed red lines). At tree level d, x is compared to its immediate K children nodes, and the MSE loss to the nearest child is added to the loss trajectory. output space. Training with ℓ 2 loss, i.e. ℓ( x, x) = ∥ x -x∥ 2 2 , the optimal K predicted leaves and probabilities are given by [45, 50] EQUATIONWe now want to extend the tessellation to the children of each node at the next level. Let {A k,q (y)} K q=1 denote the Voronoi cells at d = 2, partitioning/tessellating the Voronoi cell A k (y) from level 1 into K sub cells such that ∪ K q=1 A k,q (y) = A k (y). We want the K 2 leaves and their associated probabilities outputted by our model to be given byEQUATIONFor simplicity of notation, in the following, we omit the dependence of the Voronoi cells on y, and write A k /A k,q instead. Equations ( 2) and ( 3) expose an interesting connection between tree levels. First, recall that since {A k,q } k,q is a tessellation, it satisfies ∪ K q=1 A k,q = A k , and A k1,q1 ∩A k2,q2 = ∅ for (k 1 , q 1 ) ̸ = (k 2 , q 2 ). Invoking the law of total expectation we can writeE [x|y = y, x ∈ A k ] = K q=1 E [x|y = y, x ∈ A k , x ∈ A k,q ] P (x ∈ A k,q |y = y, x ∈ A k ) . (4)Equation ( 4) can be simplified by noting that A k ∩ A k,q = A k,q , hence the expectation in the summand is given by E[x|y = y, x ∈ A k,q ]. Moreover, we further simplify the probability in Eq. ( 4) by writingEQUATIONwhere the first equality follows from Bayes' rule and the second equality is by the law of total probability and the fact that A k,q ⊆ A k , which implies that P(x∈ A k |y = y, x ∈ A k,q ) = 1.Substituting Eq. ( 5) and the simplified expectation in Eq. ( 4) with the notations of Eqs. ( 2) and (3), we obtain thatEQUATIONEQUATIONwhere Eq. ( 7) is by the law of total probability. Note that Eqs. ( 6) and ( 7) expose the redundancy in predicting both { x⋆ k (y), α⋆ k (y)} and { x⋆ k,q (y), α⋆ k,q (y)}, 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 a tree of depth d. Given the leaves and their probabilities at level d, { x⋆ k1,...,k d (y), α⋆ k1,...,k d (y)} K k1,...,k d =1 , we can compose their parents at the upper d -1 levels, by employing Eqs. ( 6) and (7) iteratively from the bottom up.Here, this is precisely the property we exploit to output an input-adaptive tree T (y) of degree K and depth d, using a single model. Specifically, we train a model T (y; θ) ≜ ( Xd (y; θ), αd (y; θ)) that outputs K d prediction leaves Xd (y; θ) = { xn (y; θ)} K d n=1 and their accompanying probabilities αd (y; θ) = {α n (y; θ)} K d n=1 (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 i , the model T (y; θ) outputs are used as explained earlier to construct the full tree, where the root of the tree approximates the posterior mean E[x|y = y i ]. Next, the label x i is first compared to the tree root with an ℓ 2 loss to ensure our decomposition recovers the MMSE estimator. At each successive level, the label x i is 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, it 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.The approach presented in Section 3.2 is general, and can be used to augment any architecture outputting K d images and probabilities. Here, we choose to adapt the U-Net architecture [51] as a generic established choice for image-to-image regression tasks. The number of output images is K d to accommodate a tree of degree K and depth d. In case all features are shared in 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, we use group convolutions [52] in 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 K d 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 a softmax at the output.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 in 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 ε, such that predictions with non-minimal loss get updated with ε-scaled gradients. Moreover, Makansi et al. [54] proposed an evolving version of the oracle loss, where the top-k predictions are updated in each step, with k 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, we 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 ε that is annealed during training according toEQUATIONwhere t is the epoch number, and we fixed ε 0 = 1 and t 0 = 5 for all experiments. At the beginning of training, this has the benefit of bringing all predictions to a reasonable starting point, alleviating sensitivity to initialization and top prediction domination. As training progresses, this regularization is decayed, and only the relevant prediction is chosen for each sample, hence converging to the desired MCL behavior (see Appendix B).Albeit the regularization mentioned in Section 3.4, for tasks with highly imbalanced posteriors, training with Adam [55] leads to trees with a 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 a weighted sampler during training. The purpose of this sampler is to roughly balance out the overall number of occurrences at the bottom leaves during training. This is achieved by keeping track of an association matrix A ∈ R K d ×N , where A i,j counts the number of times sample x j was associated with leaf c i over some time window. This allows us to estimate the conditional probability p c|x (c i |x j ), and subsequently derive an optimized sample probability to achieve a 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. 
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 K d predictions according to the desired tree layout. Moreover, as 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 in Appendix A.Toy Example. As a warm-up, we demonstrate posterior trees on a 2D denoising task. Here, x is sampled from a mixture of four Gaussians (arranged in a rhombus-like layout), and y is a noisy version of x. The prior distribution p x (x) and exemplar samples from p x,y (x, y) are presented in Fig. 3 (a). For this simple case, the posterior distribution p x|y (x|y) 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, we train a model T (y; θ) outputting a tree of depth d = 2 and degree K = 2. The results are compared against an approximate "ground truth" in Fig. 3(c ). Note that even though we have a closed-form 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 was computed with 10K samples. Moreover, it is important to note that besides trivial cases, our result could not have been achieved with a "flat" K-means clustering with K = 4 (Fig. 3(b) ). 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 G 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 x, being either a "7" or a "9". In addition, at 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-shoes 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 on the task of image inpainting. Figures 1 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/closing, eyebrow-raising/lowering, 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 to as the task of image-to-image translation [57] . In the realm of bioimaging, such transformations were utilized to predict fluorescence from bright-field images [1] , "virtually stain" unstained tissue [3] , and transfer the images of one fluorescent dye to appear as if they were imaged with another [59] . The common scenario in image-to-image translation is the task being highly ill-posed with a wide range of plausible transformations satisfying the desired statistics.For general image editing/style transfer, this is a 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-disk microscope, simultaneously 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 × 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. Memory (GB ↓) d = 0 d = 1 d = 2 d = 1 d = 2Color.DDNM Baseline. Building on the notion of clustering from Section 3.1, we propose a 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 on samples from p x|y (x|y). Therefore, given a method to sample from the posterior (e.g., using [7, 8, 10, 11] ), a natural baseline in our case is a two-step procedure: (i) Generate N s samples from the posterior, and (ii) Apply hierarchical K-means d times to generate a tree of degree K and depth d. Note that due to computational reasons, the comparison is performed over a random subset of 100 test images from FFHQ, where for each image we generate N s = 100 samples, and apply hierarchical K-means with K = 3 and d = 2. To ensure a fair comparison, at each clustering step, K-means was run 5 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 is 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) as a function of (optimal) node depth. The second metric we adopt is the sample negative log-likelihood (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. This 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 1 compares posterior trees to the proposed baseline, implemented with various state-of-the-art 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 (≈ 14 ms with a 1.3 GB memory footprint on a A6000 GPU). This is 10 3 -10 7 × faster than the competition without considering our wide advantage in memory footprint. In case we do consider a batched setting with a sizable test set, our advantage becomes even more pronounced, where our method enables the inference of ≈ 430 test images in a single second. See Appendices I to L for more details and visual comparisons. 
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 be sub-optimal for certain posteriors. Devising a 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 a 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 are not expected to lie on the image data manifold for shallow trees. A possible solution is to apply posterior trees in the latent space of an autoencoder such as VQ-VAE [60, 61] , however, this is beyond the scope of this paper.To conclude, in this work, we proposed a technique to output a hierarchical quantization of the posterior in a single forward pass. We discussed key design choices underlying our approach, including bottom-up tree construction, a 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 a toy example. In our experiments, we applied our method to highly ill-posed inverse problems and showed that diverse (prototypical) reconstructions are possible with a simple training scheme exposing uncertainty. Additionally, we also proposed an appropriate baseline based on posterior sampling, and quantitatively compared our approach to several strong samplers. Our method demonstrated at least comparable results while being orders of magnitude faster. Finally, we applied our method to the challenging task of bioimage translation, demonstrating its practical relevance in a real-world application. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes]Justification: The per-task training times are detailed in Appendix A.3. In terms of inference times and memory usage, we discussed these extensively in Tables 1 and 2 and Appendices K and L.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, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as 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). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?Answer: [Yes]Justification: the paper fully conforms with the NeurIPS Code of Ethics.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [Yes] Justification: In Section 1 we discussed the potential positive societal impact of our method, and in Section 4 we demonstrated its practical use for properly visualizing reconstruction uncertainty in safety-critical tasks such as bioimage translation.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or 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 considerations. • The conference expects that many papers will be foundational research and not 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 correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or 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 a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA] Justification: This paper poses no such risks. Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a 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 a best faith effort. 12. Licenses for existing assets Question: 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? Answer: [Yes] Justification: All dataset papers are properly cited. Licenses are as follow: MNIST (CC BY-SA 3.0), Edges→Shoes (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 use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL.• 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. 
 
 Appendix: As mentioned in Section 3.3, in our experiments we adopted the U-Net [51] architecture. Our architecture consisted of 4 downsampling/upsampling blocks. Downsampling by 2× was performed using average pooling, and upsampling by 2× was implemented with a nearest-neighbor 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 a negative slope of 0.2. The number of features at each spatial resolution was adapted to the predicted tree layout, i.e., for predicting a tree of degree K and depth d, the initial number of channels before downsampling was set to c init K d , where K d is the number of output leaves and c init ∈ {4, 8}. At each successive encoder block, the number of channels was doubled, reaching 16c init K d at the bottleneck.For example, assuming a tree with degree K = 3 and depth d = 2, and setting c init = 4, the number of channels per level in the encoder is given by [36, 72, 144, 288, 576] .Similarly, in the decoder, the number of channels was halved at each upsampling step in a 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 an architecture that shares the encoder between the different leaves, while having disjoint decoders, each with a dedicated set of weights {φ n } K d n=1 , 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 K d 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 [d f , 256, 64, K d ], where d f is the dimension of the concatenated pooled features from the decoder. Each linear layer is followed by a 1D batch normalization and SiLU non-linearity, and the output is passed through a softmax layer to produce a valid probability vector. Figure A2 summarizes our architecture.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 × 32 to enable proper downsampling in the encoder. For Edges→Shoes experiments, the images were kept the same as in the pix2pix [57] paper, with a resolution of 256 × 256. For CelebA-HQ experiments, the images were resized to 256 × 256. Finally, for the bioimage translation dataset, we trained on 128 × 128 patches cropped from the full 1024 × 1024 images, as cell information tends to be local.We used the Adam optimizer [55] with β 1 = 0.9, β 2 = 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, we found it beneficial to use a smaller initial step size of 0.0002. This is because 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 a sparse vector, leading to leaves with zero probability. Hence, to avoid this instability, we used a lower step size for the probabilities such that the leaves are first allowed to 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→Shoes, CelebA-HQ, and the bioimage datasets respectively. ). The encoder is shared between all leaves, however, in the decoding stage, each leaf has a separate decoder (e.g., with parameters φ 1 ,φ 2 ,φ 3 ) to enable diverse outputs. To predict leaf probability, all features from the decoder are passed through a global average pooling layer, and the resulting concatenated vector is passed to a lightweight MLP with a softmax layer at the output. Afterward, as explained earlier, the output tree is constructed iteratively from the bottom up (right).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 in 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 t 0 = 5, our training degenerates to standard MSE minimization, bringing all leaves to a reasonable starting point near the posterior mean (Fig. A3(c) ). This results in all prediction leaves having roughly the same initialization, with a similar number of associated training samples. Afterward, ε is gradually decayed, and the leaves start to specialize in their respective posterior mode, converging to the desired MCL behavior (Fig. A3(d) ). In this research, we identified this strategy as a 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 it would be intriguing to explore this in future work. where each leaf specializes in a different posterior mode (e.g., a "3" or a "5").Naturally, since K d ≪ N , there are infinitely many different options q that satisfy this constraint. Note that some of them can correspond to an imbalanced distribution, where some samples associated with a 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 ℓ 2 loss, namely ∥q∥ 2 2 . However, minimizing this objective under the constraint in Eq. (A4) can be problematic for ill-conditioned cases. Nevertheless, we can achieve similar results by solving the following Tikhonov-regularized relaxed problem:min q∈R N 1 2 ∥P q - 1 K d 1∥ 2 2 + λ 2 ∥q∥ 2 2 s.t. q j ≥ 0, j = 1, . . . , N q ⊤ 1 N ×1 = 1 (A5)One can obtain an exact solution q ⋆ to Eq. (A5) 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, if we drop the positivity constraint q j ≥ 0, ∀j, and solve Eq. (A5) accounting only for the constraint q ⊤ 1 N ×1 = 1, we obtain the following closed-form solutionq = I N ×N -P ⊤ P P ⊤ + λI K d ×K d -1 P 1 N ×1 q ⋄ = 1 q⊤ 1 N ×1 q.(A6)In practice, in most cases, the positivity constraints are inactive, and therefore q ⋄ = 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 q s (s j ) we pick for the next batch b τ +1 the B distinct samples that their selection will minimize the maximal distance between the future statistics at time τ + 1 and q s . Specifically, at each training step, the next batch of samples is chosen such that the maximal difference between the next iteration estimated sample probability p s (s j ) = K d i=1 p c,s (c i , s j ) and the desired sample probability q s (s j ), i.e., ∥p sq s ∥ ∞ , 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 ε t started decaying to 0.Finally, note that switching from p s to q s during training effectively changes the dataset by repetition/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 s j by the factor γ j = Our tree training algorithm is summarized in Algorithm 1. Throughout this work, we used the MSE loss ℓ(x, x) = ∥x -x∥ 2 2 to both determine the closest tree node and as our optimization loss. However, in general, 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. In theory, our approach can work without modification with any association measure (e.g. LPIPS [64] , some domain-specific classifier [65] , etc.). However, changing the within-cluster loss requires some modifications. Specifically, the use of the MSE loss is what provides us with the hierarchical decomposition of the posterior mean (see Eqs. ( 2) to ( 7)). In particular, when using the MSE/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 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 our approach is to take the association measure to be ℓ( x, x) = ∥f ( x) -f (x)∥ 2 2, where f (•) is some relevant domain-specific feature extractor. Specifically, 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 is necessary to ensure semantically meaningful results, which will be addressed in future work. { xn (y i ; θ), αn (y i ; θ)} K d n=1 ← T (y i ; θ) ▷ Infer tree leaves and their probabilities /* Bottom-up tree construction using Eqs. ( 6) and ( 7) */ { xk (y i )} K k=1 ← GetChildren( xk ⋆ (y i )) ▷ Get K children ofL(θ) = d+1 j=1 β j ∥ Xj -x i ∥ 2In Section 3.1 we mentioned that in a discrete setting, our approach is equivalent to a hierarchical 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 {x i } where x i ∈ R dx , the goal of K-means is to partition the data points into K clusters/sets {C 1 , . . . , C K }, such that each data point belongs to the cluster with the nearest mean/centroid serving as a prototype of the cluster. This results in a partitioning/tessellation 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 is given by arg minEQUATIONwhere µ k is the mean/centroid of cluster C k , given by µ k = 1|C k | xi∈C k x i .While in general Problem (A10) is NP-hard, given some initial means/centroids, the K-means algorithm (also known as Lloyd's method) finds a 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 {x i }, at each level of the hierarchy, the data points belonging to the kth cluster C k are split further into K sub-clusters {C k,1 , . . . , C k,K } by applying K-means. We refer to this extended algorithm as hierarchical K-means. The result of applying hierarchical K-means d times is a dendrogram/top-down balanced tree T of degree K, depth d, and breadth (final number of leaves) K d . Note that this successive process is different from applying K-means with K d clusters once. This is because the hierarchical memberships are enforced, restricting data points in subclusters {C i,1 , . . . , C i,K } from being assigned to subclusters {C j,1 , . . . , C j,K } if i ̸ = j.Hierarchical K-means enjoys several advantages compared to the standard ("Flat") K-means algorithm. One of these advantages is increased robustness to initialization as demonstrated in Appendix G.3. Throughout Section 4 we presented results with trees of degree K = 3 and depth d = 2. 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 A9 presents the resulting trees for four different model configurations on the task of digit inpainting. For a 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 ofK = 4, d = 1 to 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 a cluster midway between a "4" and a "1". Figure A9 repeats this experiment for deeper trees comparing K = 4, d = 2 to K = 2, d = 4. A couple of remarks are in order regarding the results. First, in both cases, the "4" cluster has a similar probability relative to the "1" 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 of the output space. Second, the degree K controls the emphasis/over-representation devoted to weaker posterior modes. A smaller degree leads to more emphasis on weaker modes as tree depth grows. Lastly, it is important to note that the optimal layout K and d is task and input-dependant, and setting these adaptively is an interesting direction for future research. In Section 4 we compared our method mainly to diffusion-based posterior samplers (e.g., [7, 8, 11] ). While these samplers often lead to state-of-the-art sample quality, they are known to be computationally intensive. In theory, GAN-based samplers such as MAT [10] , can significantly speed up the proposed two-step baseline of sampling followed by hierarchical K-means; however, as evident in Fig. A10 , even top-performing GAN-based samplers often collapse to a single mode of the posterior, producing samples with little variability that do not faithfully reflect the full distribution. In Table 1 we reported quantitative comparisons to 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. In Table 1 we reported the runtime and memory footprint assuming a batch of size 1. However, in GPU computing, memory footprint and hardware parallelization ultimately determine the overall runtime. We tried to refrain from factoring these in our calculation in Table 1 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 a batch of size 1, it also has a 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 a function of the batch size on an A6000 GPU with 48 GB. For each method, we tested three different batch sizes: 1, 10, and the maximal number of samples that fit in memory. The forward pass speed and the memory usage in each setting were tested 100 times using CUDA events (reported as mean ± std). Our method is extremely fast and parallelizable, enabling the inference of 430 test images with a single forward pass of 1.15 seconds.Even in the case of a single test image, our method is far superior. For simplicity let us assume the cost of running hierarchical K-means is negligible, and that we can squeeze a batch size of ≈ 100 samples for the diffusion-based samplers, and ≈"33. As mentioned in Section 4.1, baseline trees were computed using a two-step procedure: (1) Sampling N s 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. Figure A13 plots the success probability of building a tree with K d = 9 leaves as a function of the number of sampled images N s in the different tested tasks. A tree is considered to be successfully constructed if each of its leaves has at least 2 samples to average. As evident in the results (e.g., Fig. A13(b )), 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. 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 a sizable test set. Here we include more results for each of the tasks presented in Section 4 from the main text. Additionally, we also include colorization results on the AFHQ(v2) dataset [67] , with a similar performance to that obtained on CelebA-HQ. • 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 or 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 a 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 or figures symmetric error bars that would yield results that are out of 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.N qs(sj ) . 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. 


Decision: 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.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: 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.
2. 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.
3. 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.

Overall I enjoyed reading the paper, and the paper certainly seems to be novel, albeit, incrementally. Hence, I lean for weak acceptance.
questions: questions: See weaknesses
limitations: limitations: See weaknesses
paper_summary: paper_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.
strength_weakness: strength_weakness: 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 .
2. 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.
3. 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.
questions: questions: 1. Can you apply this method when the reconstruction loss is not an $L_2$ loss but a generic function?
2. How does the runtimes compared in Table 1? 
3. 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: limitations: The authors discussed limitations of the paper.
paper_summary: paper_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.
strength_weakness: strength_weakness: * 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?

* 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?
questions: 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?

* 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: limitations: Authors provide a discussion on the limitations of their
method.

Comments: title: Response to Rebuttal
Runtime detailed comparison
comment: Thank you for your response. I am convinced with the response to my questions, further the response provided by the authors to Reviewer VHHJ helped me in further clarification. I am increasing my score accordingly!
I appreciate the thorough response from the authors. My concerns have been addressed. I will increase my score.
Indeed, memory footprint and hardware parallelization ultimately determine the overall runtime. We tried to refrain from factoring these in our calculation as these numbers are hardware-specific depending on the available GPU. Nonetheless, our method also has a lower memory footprint and is just as amenable to parallelization. We apologize if our previous answer needed to be clearer. The table below benchmarks the speed of the forward pass and the GPU memory usage as a function of the batch size on an A6000 GPU with 48 GB:

|                  |       | Ours  |       |&#124;       | DDRM/DDNM/RePaint |       |&#124;        |  MAT  |       |
| :---:            | :---: | :---: |  :---: |  :---: |       :---:       | :---: |  :---: | :---: | :---: |
| Batch size       |   1   |  10   |  430  |   1   |         10        |   97  |   1   |   10  |   30  |
| forward pass (ms)|14$\pm$0.2|39$\pm$0.2|1150$\pm$2|34$\pm$0.4|245$\pm$0.4|2315$\pm$6.5|150$\pm$8|867$\pm$3.8|2599$\pm$13.6|
| GPU memory (GB)  |  1.3  |  2.1  | 46.9  |  1.85 |        7.5        |  46.8 | 3.0   | 16.0  | 47.2  |

For each method, we tested 3 different batch sizes: 1, 10, and the maximal number of samples that fit in memory. The forward pass speed and the memory usage in each setting were tested 100 times using CUDA events (reported as mean$\\pm$std). Our method is extremely fast and parallelizable, enabling the inference of 430 test images with a single forward pass of $\approx$1.15 seconds. 

**How does this compare to the baselines for a single test image?**  
Our method is far superior even for a single test image. For simplicity let us assume the cost of running hierarchical $K$-means is negligible, and that we can squeeze a batch size of $\approx$100 samples for the diffusion-based samplers, and $\approx$“33.33” samples for MAT. Sampling $N\_s=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\\times$ speedup compared to the fastest baseline.
Thank you for the clarifications. I don't see the architecture-agnostic nature of NFE necessarily as a strength at all since what matters in practice is the runtime (and memory considerations), which depends not only on architecture, but also hardware utilization of each method. Regarding the provided runtime table, why did you choose to use a batch size of 1? Using a larger batch size should lead to significant speed up for the baselines since sampling is parallelized.