Risk-Aware Image Generation by Estimating and Propagating Uncertainty

While generative AI models have revolutionized content creation across various modalities, they have yet to be deployed in safety-critical scenarios. This is in part due to limited understanding of their underlying uncertainty, as general-purpose frameworks for estimating uncertainty in large-scale generative models are lacking. Here we analyze the effects of uncertainty and risk estimation methods on generative AI systems and their applications to two critical domains of deployment -- identification of failures, and fast optimization of input prompts. As a case study, we apply our approach to create an uncertainty-aware variant of the Stable Diffusion text-to-image model, allowing us to estimate and propagate uncertainty over inputs, latent representations, and outputs. We demonstrate that our method enables the identification of uncertain output regions and the optimization of input prompts to minimize output uncertainty. We envision that our framework will enable the deployment of more robust and auditable generative AI systems.