Diffusion Models for Inverse Problems in the Exponential Family

TL;DR

This paper extends diffusion models to inverse problems with exponential family observations, enabling Bayesian inference in complex real-world scenarios like Poisson processes and disease prevalence prediction.

Contribution

It introduces the evidence trick, a novel method leveraging exponential family conjugacy to approximate likelihood scores in diffusion models.

Findings

Effective Bayesian inference on complex Poisson processes.

Competitive performance in malaria prevalence prediction.

Demonstrates applicability to real-world spatial data.

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.