DiffEM: Learning from Corrupted Data with Diffusion Models via Expectation Maximization

TMLR Paper9242 Authors

Abstract: Diffusion models have emerged as powerful generative priors for high-dimensional inverse problems, yet learning them when observations are only available through a corruption channel remains challenging. In this work, we propose DiffEM, a new method for training diffusion models with Expectation-Maximization (EM) from corrupted data that does not rely on any approximations or heuristics. DiffEM utilizes conditional diffusion models to reconstruct clean data from observations in the E-step, and then uses the reconstructed data to refine the conditional diffusion model in the M-step. Theoretically, we provide monotonic convergence guarantees for the DiffEM iteration, assuming appropriate statistical conditions. We demonstrate the effectiveness of our approach through experiments on various image reconstruction tasks.

Submission Type: Regular submission (no more than 12 pages of main content)

Assigned Action Editor: ~Valentin_De_Bortoli1

Submission Number: 9242