Diffusion Models for Inverse Problems in the Exponential Family

2025-02-09

AI Summary

This work addresses the limitation of diffusion models in non-Gaussian inverse problems by extending denoising diffusion probabilistic models to exponential-family observations—such as Poisson and binomial distributions—thereby relaxing the conventional Gaussian noise assumption. We introduce the “evidence trick,” leveraging exponential-family conjugacy to derive an analytic approximation of the likelihood score, enabling efficient Bayesian inversion for non-Gaussian, sparse, and count-valued data within the diffusion framework. Our method integrates exponential-family conjugate prior modeling, variational inference, and score matching. Experiments demonstrate high-fidelity uncertainty quantification in Poisson intensity field reconstruction at ImageNet scale, and achieve state-of-the-art performance in malaria prevalence prediction across sub-Saharan Africa.

Technology Category

• Reasoning under Uncertainty: Relational Probabilistic Models
• Machine Learning: Probabilistic Circuits and Graphical Models
• Computer Vision: Diffusion Models for Vision

Application Category

• Graph Algorithms and Modeling for the Web: Algorithms and analysis for incomplete, noisy, or partially observed Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Explainable and interpretable methods for personalization

Abstract

Diffusion models have emerged as powerful tools for solving inverse problems, yet prior work has primarily focused on observations with Gaussian measurement noise, restricting their use in real-world scenarios. This limitation persists due to the intractability of the likelihood score, which until now has only been approximated in the simpler case of Gaussian likelihoods. In this work, we extend diffusion models to handle inverse problems where the observations follow a distribution from the exponential family, such as a Poisson or a Binomial distribution. By leveraging the conjugacy properties of exponential family distributions, we introduce the evidence trick, a method that provides a tractable approximation to the likelihood score. In our experiments, we demonstrate that our methodology effectively performs Bayesian inference on spatially inhomogeneous Poisson processes with intensities as intricate as ImageNet images. Furthermore, we demonstrate the real-world impact of our methodology by showing that it performs competitively with the current state-of-the-art in predicting malaria prevalence estimates in Sub-Saharan Africa.

Problem

• Extend diffusion models
• Handle non-Gaussian observations
• Bayesian inference on Poisson processes

Innovation

• Extends diffusion models
• Handles exponential family distributions
• Introduces evidence trick