Estimating Epistemic and Aleatoric Uncertainty with a Single Model

Summary

The work applies hypernetworks to estimate both aleatoric and epistemic uncertainty in the context of diffusion models. The authors leverage the distribution on weights created by the hypernetwork to estimate uncertainty via the Total Uncertainty (TU) decomposition into Aleatoric Uncertainty (AU) and Epistemic Uncertainty (EU). They validate their method on a toy dataset, lung scan dataset, CT scan dataset and weather predictions.

Strengths

1. The paper is clear and well-written. 2. The authors provide a solution to an important topic, given the prevalence of generative models. 3. The authors apply their method to both a straightforward toy problem and real-life datasets.

Weaknesses

1. The authors seem to be missing a citation where previous researchers have estimated epistemic uncertainty for diffusion models. - Berry, Lucas, Axel Brando, and David Meger. "Shedding Light on Large Generative Networks: Estimating Epistemic Uncertainty in Diffusion Models." The 40th Conference on Uncertainty in Artificial Intelligence. 2. There are claims regarding ensembles that appear incorrect. It is not necessary to train $M$ distinct networks; many modern methods work by ensembling certain parts of the network while sharing weights across the rest. Thus, training ensembles is not as expensive as stated. - Osband, Ian, et al. "Deep exploration via bootstrapped DQN." Advances in Neural Information Processing Systems 29 (2016). - Berry, Lucas, and David Meger. "Normalizing flow ensembles for rich aleatoric and epistemic uncertainty modeling." Proceedings of the AAAI Conference on Artificial Intelligence. Vol. 37. No. 6. 2023. 3. Minor writing comments: - The abstract is a bit unusual with two paragraphs and could be shortened. - The notation in section 4.2 seems overloaded. Specifically, the subscript on $x$ refers to both the denoising steps and the samples in the training set.

Questions

How does the scalability of hypernetworks compare to that of MC Dropout? In both instances, there are challenges with scaling the processes to networks that require a lot of parameters.

Rating

6

Confidence

4

Soundness

3

Presentation

2

Contribution

2

Limitations

The authors address their limitations in the appendix.

Summary

The paper proposes a approach, HyperDDPM, to estimate both aleatoric and epistemic uncertainty with an ability to disentangle them. This result is achieved by combining diffusion networks (DDPM) to sample predictions from single set of model weights, and Bayesian Hyper Network to generate sets of weights for main DDPM model. Using synthetic data the authors show, that proposed approach can estimate aleatoric and epistemic components of predictive uncertainty, and provide evidence for real-world applicability on medical imaging and weather forecasting tasks.

Strengths

1. The paper presents a new and potentially useful way of producing ensembled predictions. 2. On synthetic data it shows that this approach is sensitive to both epistemic and aleatoric uncertainty components when ground truth influence of such components is isolated. 3. It shows on real world tasks, like medical image reconstruction and meteorological forecasting, that proposed architecture can outperform some, and perform on par with the other baseline methods. 4. The authors provide codebase used to train and evaluate models, which allows for better transparency and reproducibility. [After rebuttal comment] I increased my score from 3 to 5. I am still not convinced that uncertainty is treated in coherent way as the paper jumps from the Bayesian predictive distribution to the law of total variance without any connection betweeen them. Additionally, real world experiments do not evaluate uncertainty numerically.

Weaknesses

Here I summarize weaknesses of the paper, some details on the weaknesses are clarified by the Questions section of the review. 0. The paper lacks a Contributions section, clearly and concisely stating which claims and results within the paper are new. 1. The experimental section does not provide sufficient evidence of the ability of the proposed approach to disentangle epistemic and aleatoric components of uncertainty. 2. Theoretic justifications do not fully align nicely with the engineering approach taken for experiments, and contain questionable claims. 3. Some details on training procedure are unclear. 4. The paper lacks a straightforward comparison of computational overhead of baselines and proposed method, during training and evaluation. Some of these are present scattered within paper text, but the paper would benefit from having them collected in one place, given that both baselines and proposed method incur significant overhead.

Questions

1. In explanation of eq. 1 the authors claim that likelihood under the integration contains information about aleatoric uncertainty, which can be argued as true, but it seems that marginal distribution on the left side of the equation should provide a better estimate of aleatoric component, given that it's marginalized over model weights distribution, and thus epistemic component is eliminated. This can be further confirmed by the way the authors approach estimating aleatoric uncertainty in practice, described in eq. 10 - here the authors take expectation over the model weights in the outer operator. 2. In Appendix A authors describe training of BHN as: - compute an L2 loss, and back-propagate the gradients to update the hyper-network weights. But in the next paragraph the following is said: - The BHN is trained with the same L2 loss and an additional Kullback-Leibler (KL) divergence loss How exactly was BHN training performed? It is also quite interesting why BHN training did not collapse to an optimal mode with very low variance over produced model weights. How often was input noise vector for BHN changed during training? On per-batch level, per-input, or using some other strategy? 3. Toy experiment on synthetic data is only performed to estimate AU and EU when corresponding uncertainty component is changing. It seems to not fully support claim that proposed approach can disentangle AU from EU. It would be much more demonstrative of claim, if the experiment was done by independently varying noise level and train dataset size, and plotting results on AU and EU axes. 4. Real world experiments only show base performance of proposed architecture, but do not provide reasonably in-depth analysis of UE capabilities on such task, apart from figures with uncertainty maps over generated images. The paper would benefit greatly from some form of analysis of how uncertainty estimates improve prediction/reconstruction quality when filtering out low-confidence inputs (i.e. rejection verification). 5. Was training of base model of MC-Dropout baseline performed using active dropout masks in the same way as during inference? Some degree of performance hit that MC-Dropout model suffered could be explained by difference between training and evaluation dropout regimes.

Rating

5

Confidence

4

Soundness

2

Presentation

3

Contribution

2

Limitations

The authors have adequately addressed limitations of their work.

Summary

This paper proposes HyperDDPM, a novel uncertainty quantification method for generative tasks that uses a hypernetwork and a denoising diffusion probabilistic model (DDPM) and outputs both aleatoric and epistemic uncertainty estimates. HyperDDPM is evaluated on a toy task with available ground-truth uncertainties as well as two real-life tasks, CT reconstruction and weather forecasting. The authors find that HyperDDPM gives accurate uncertainty estimates that are on par with deep ensembles.

Strengths

- Uncertainty quantification in generative models is a high-impact research area. - The paper is easy to read. - HyperDDPM is a single-model uncertainty estimator that performs on par with ensembles that are SotA in many uncertainty application domains. - The use of hypernetworks for uncertainty quantification is an interesting avenue. - The experiments are extensive and clearly explained. The toy experiment allows for fine-grained control and comparison with ground-truth quantities, whereas the real-world tasks of CT reconstruction and weather forecasting are important application areas. The authors provide both quantitative and qualitative evaluation.

Weaknesses

- The paragraph starting at L29 is (i) misleading and (ii) vague. (i) A useful uncertainty estimate does not necessarily have to differentiate between aleatoric and epistemic uncertainty. A total/predictive uncertainty estimate subsumes both of these sources of uncertainty and can be excellent at predicting the underlying model's correctness on an input $x$ (e.g., see Fig. 6 in (Mukhoti et al., 2023; CVPR)). (ii) The "valuable insights into the strengths and weaknesses of a predictive model" are not elaborated, neither are the pathways these insights offer "towards improving [the model's] performance". I _do_ agree that uncertainty disentanglement is an important and interesting research direction but this paragraph doesn't quite answer _why_. Consider providing concrete use cases for the separate aleatoric and epistemic estimates and how they can improve the model's performance. - A more complex model class generally does not lead to a decrease of epistemic uncertainty (Section 3.2). Larger models introduce more degrees of freedom into the learning problem, leading to a larger set of models consistent with the training data (as captured by the Bayesian posterior $p(\phi \mid \mathcal{D})$). - The claim that "This estimate converges to the true AU of the inverse problem as $N \to \infty$" is imprecise and unproven. The statement assumes that (i) the score estimator $s_\phi$ is expressive enough to model the true score of the generative process perfectly, (ii) $s_\phi$ is trained on the entire data distribution (i.e., the expectation in Eq. (5) is optimized), and (iii) the global minimum is reached during optimization. Only then can one argue that the true AU of the inverse problem is reached at an _in-distribution_ input $x$ when one samples infinitely many samples from the implicit distribution. **Minor Issues:** - Consider making the figures vector-graphical. The current PNGs are quite pixelated. - The precise formulation of Eq. (7) would be $\mathbb{E}_{(x, y) \sim \mathcal{D}}\left[\mathbb{E}_{z \sim \mathcal{N}(0, \sigma^2)}\left[\mathcal{L}(f(y \mid \phi(z)), x)\right]\right]$. The sampling is not captured by the original formula. - Consider moving Appendix C into the main paper in the final version.

Questions

- Why should one trust their aleatoric uncertainty estimates on OOD data (Fig. 2)? Usually, a two-step procedure is carried out (see Fig. 3 in Mukhoti et al., 2023; CVPR) where the AU estimates are only considered if the EU estimates are below a certain threshold. It also seems that at OOD regions, the AU and EU estimates are highly correlated.

Rating

6

Confidence

4

Soundness

3

Presentation

4

Contribution

3

Limitations

The authors address the limitations of their work adequately.

Summary

Diffusion has been applied in many domains beyond image generation, e.g. weather forecasting and CT reconstruction. Many of these applications are safety-critical as such the model should be capable of expressing uncertainty. As such the submission proposes a method based on the idea of hypernet for uncertainty quantification: The method would train a network that maps from noise to the a set of weights of the U-net of the diffusion, then generated weight is then treated a sample from the "posterior" and later used to estimate the epistemic uncertainty. In order to estimate aleatoric uncertainty, the method uses MC estimation to estimate var(y | x, phi) for a fixed model weight and then average the variance across all weights generated by the hypernet. The paper then evaluates the method on weather forecasting problem and CT image reconstruction problem: In both setting, the method demonstrates a nice decomposition of aleatoric and epistemic uncertainty. Overall I find the application interesting and paper is nicely written and presentation is clear. However I find the technical contribution a little limited although the empirical results seem to be promising.

Strengths

- The problem studied, uncertainty quantification for diffusion, is important. - The method proposed is technical sound and it works well in empirical evaluation. - The experiment and hyper-parameter setting are clearly presented in the appendix.

Weaknesses

- Error bar / standard deviation is not provided. - It would be nice if a summarization of runtime comparison can be provided, i.e. total training time, total inference time. - The technical novelty is not very outstanding (though the application seems to be novel), see the bullet point below for more detail. - The contribution of the paper is not described in a very precise way: From my perspective, the paper applies existing methods, i.e. Bayesian hyper-network, in a new setting: **diffusion models**. However this is not highlighted in the abstract. The abstract gives readers the impression that the paper is proposing uncertainty quantification methods for generic problems. - (Minor) Lack of baseline methods: Some strong BNN baselines such cyclical SGLD or linearized Laplace, SWAG are not included. MC-dropout used in the paper is a pretty out-of-date technique for uncertainty quantification problem in my opinion. - (Minor) In Fig. 1, hyper-network is written as hypernetwork in the figure body (a), it would be better if the term can be spelled in a consistent way.

Questions

- What's the training overhead for the hyper-network? - How should one set the complexity of the hyper network? If the network is too large / small, would the uncertainty vanish? - How does the number of samples, i.e. M and N, affect the performance? - What does BNN mean in the main text? Does it refer to MC-dropout or DPS-UQ or both?

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

2

Limitations

- To estimate epistemic uncertainty, although the method only needs a single model (in contrast to BNNs that requires multiple copies/samples of a model), the method still needs **multiple forward propagation** to estimate the uncertainty, so the part of the computational cost at test time still exists

Rebuttal well received

Dear authors, your rebuttal was well received and you partially addressed my concerns. I will decide on the score changes after the discussion with other reviewers.