Casting Light on Large Generative Networks: Taming Epistemic Uncertainty in Diffusion Models

Lucas Berry, Axel Brando, David Meger

Primary Area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Keywords: Calibration & Uncertainty Quantification, Ensemble Methods, Diffusion Models

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2024/AuthorGuide.

TL;DR: DECU is an innovative framework for efficiently estimating epistemic uncertainty in large generative diffusion models, offering valuable insights into these complex models, particularly for under-sampled image classes.

Abstract: Epistemic uncertainty plays a pivotal role in contemporary machine learning, serving as a fundamental element that underlies decision-making processes, risk evaluations, and the overall generalizability of models. In this work, we introduce an innovative framework, diffusion ensembles for capturing uncertainty (DECU), designed for estimating epistemic uncertainty within the realm of large high-performing generative diffusion models. These models typically encompass over 100 million parameters and generate outputs within a high-dimensional image space. Consequently, applying conventional methods for estimating epistemic uncertainty is unrealistic without vast computing resources. To address this gap, this paper first presents a novel method for training ensembles of conditional diffusion models in a computationally efficient manner. This is achieved by fitting an ensemble within the conditional networks while using a static set of pre-trained parameters for the remainder of the model. As a result, we significantly reduce the computational load, enabling us to train only a fraction (one thousandth) of the entire network. Furthermore, this substantial reduction in the number of parameters to be trained leads to a marked decrease (87%) in the required training steps compared to a full model on the same dataset. Second, we employ Pairwise-Distance Estimators (PaiDEs) to accurately capture epistemic uncertainty with these ensembles. PaiDEs efficiently gauge the mutual information between model outputs and weights in high-dimensional output space. To validate the effectiveness of our framework, we conducted experiments on the Imagenet dataset. The results demonstrate our ability to capture epistemic uncertainty, particularly for under-sampled image classes. This study represents a significant advancement in detecting epistemic uncertainty for conditional diffusion models, thereby casting new light on the $\textit{black box}$ of these models.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors' identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 7545