Learning Diffusion Priors from Observations by Expectation Maximization (2405.13712v4)

Abstract: Diffusion models recently proved to be remarkable priors for Bayesian inverse problems. However, training these models typically requires access to large amounts of clean data, which could prove difficult in some settings. In this work, we present a novel method based on the expectation-maximization algorithm for training diffusion models from incomplete and noisy observations only. Unlike previous works, our method leads to proper diffusion models, which is crucial for downstream tasks. As part of our method, we propose and motivate an improved posterior sampling scheme for unconditional diffusion models. We present empirical evidence supporting the effectiveness of our method.

• The paper demonstrates that EM can iteratively refine diffusion priors using noisy observations, eliminating the need for clean latent samples.
• It introduces Moment Matching Posterior Sampling to enhance sampling accuracy by aligning model outputs with the observed data's statistical moments.
• Empirical experiments on corrupted CIFAR-10 and MRI data validate the approach, broadening diffusion model applications in image reconstruction and data assimilation.

Summary of "Learning Diffusion Priors from Observations by Expectation Maximization"

The paper "Learning Diffusion Priors from Observations by Expectation Maximization" introduces a novel framework for training diffusion models in contexts where clean training data is scarce or noisy, typical in domains such as Earth and space sciences. The authors address the challenge of developing diffusion models for Bayesian inference using the Expectation-Maximization (EM) algorithm, presenting a method that circumvents the need for large datasets of clean latent samples.

Core Contributions

1. Expectation-Maximization for Prior Learning: The paper repurposes the EM algorithm to iteratively refine a diffusion prior, catering to scenarios where only corrupted observations are available. Unlike prior work that requires either explicit latent observations or access to clean and complete datasets, this framework leverages incomplete data (in the form of observed noisy measurements) in combination with the EM paradigm.
2. Moment Matching Posterior Sampling: A key innovation presented is the Moment Matching Posterior Sampling (MMPS) technique, which improves posterior sampling accuracy for unconditional diffusion models. By accounting for both the mean and covariance within the posterior sampling, MMPS advances beyond traditional heuristics, producing samples that prove critical for the accurate alignment of diffusion models with observed data.
3. Empirical Bayes Contextualization: The paper places this EM approach within the empirical Bayes context, wherein the parameters of the prior model are optimized to best match the empirical distribution of observations. This setup is particularly useful in high-dimensional scenarios, where classic empirical Bayes methods fall short.

Numerical Results and Implications

The authors provide substantial empirical evidence demonstrating the efficacy of their approach. Noteworthy experiments include:

• Studies on low-dimensional manifolds showcasing the convergence of estimated distributions to the true data distribution, despite the conditional observations.
• Performance on corrupted CIFAR-10 and MRI data, where the new method yields higher fidelity samples compared to existing techniques like AmbientDiffusion.

These results extend the applicability of diffusion models to a broader range of practical problems, underpinning the method’s potential to enhance tasks in fields like image reconstruction and data assimilation, where access to clean training data is limited.

Theoretical and Practical Impacts

From a theoretical standpoint, this work bridges a gap in Bayesian inference, effectively leveraging diffusion processes in conjunction with the EM algorithm to optimize prior models directly from noisy observations. This development enriches the existing statistical frameworks for dealing with ill-posed problems, whereby traditional methods cannot exhaustively capture the underlying distributions.

Practically, the pipeline offers a significant leap in training efficient diffusion models where previously infeasible due to data cleanliness requirements. In healthcare and remote sensing, for instance, this methodology can promote accurate image recovery from incomplete data, enhancing both diagnosis and scientific exploration.

Future Directions

The paper opens several avenues for future research:

• Exploration of more complex observation models, including non-linear and non-Gaussian processes, that would broaden the scope of potential applications.
• Algorithmic improvements to mitigate the computational demand inherent in high-dimensional settings, which is crucial for real-time or resource-constrained applications.

In conclusion, the presented work delineates a strategic and methodological advancement in leveraging diffusion models for practical Bayesian inference challenges. The approach fortifies the robustness and versatility of diffusion-based priors, thus paving a pathway for future explorations into scalable algorithms in data-scarce environments.