Diffusion Models for Inverse Problems in the Exponential Family
This work addresses a key bottleneck for applying diffusion models to real-world inverse problems with non-Gaussian noise, such as in medical imaging or epidemiology, though it is incremental in extending existing methods to new data types.
The authors tackled the limitation of diffusion models to Gaussian noise by extending them to handle inverse problems with exponential family distributions, such as Poisson or Binomial, using an evidence trick for tractable likelihood scores. They demonstrated effectiveness on complex Poisson processes and competitive performance in predicting malaria prevalence in Sub-Saharan Africa.
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.