Learning Diffusion Priors from Observations by Expectation Maximization

2024-05-22

Neural Information Processing Systems

AI Summary

To address the challenge of scarce clean training data in Bayesian inverse problems, this paper proposes the first framework for learning theoretically grounded diffusion priors solely from incomplete and noisy observations. Methodologically, we embed diffusion probabilistic modeling into an Expectation-Maximization (EM) algorithm: the E-step estimates the latent variable posterior via iterative denoising sampling, while the M-step updates diffusion model parameters by maximizing the marginal likelihood. Our key contributions are: (1) the first provably consistent learning of diffusion priors directly from noisy and/or missing observations; and (2) an unconditional posterior sampling strategy that eliminates reliance on assumptions about the forward process. Experiments demonstrate that the learned prior achieves performance on par with fully supervised models in downstream inverse tasks—including denoising and inpainting—while ensuring rigorous generative consistency and theoretical guarantees.

Technology Category

• Machine Learning: Bayesian Learning
• Reasoning under Uncertainty: Relational Probabilistic 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: User privacy protection in personalized systems

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.

Problem

• Training diffusion models without clean data
• Learning priors from noisy observations only
• Improving posterior sampling for diffusion models

Innovation

• Trains diffusion models from noisy observations
• Uses expectation-maximization algorithm for training
• Proposes improved posterior sampling scheme