Hierarchical Uncertainty Exploration via Feedforward Posterior Trees

Tree loss function and training scheme
Thanks 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.

Runtime comparisons
Table 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.

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

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

Required hyper-parameter tuning to stabilize training
As 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.

Scaling to wider/deeper trees
Indeed, 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.

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 the user has to skim through many images for every input.

Comparing runtime in neural function evaluations vs GPU seconds
Please 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.

Baselines performance as a function of compute
Note 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$.

Learning posterior trees with other loss functions
This 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).

We 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.

**Providing more background on CVT and the results of $45$ **

We 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.

Tradeoff between breadth and depth of the constructed trees

Kindly 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.

In 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.

Number of prototypes and operating in input/pixel space

This 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.

Novelty

While 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:

• 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).
• 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).
• Demonstrating, to the best of our knowledge, the first successful application of MCL with diverse results for high-resolution image-to-image regression tasks.