An Expectation-Maximization Algorithm for Training Clean Diffusion Models from Corrupted Observations

TL;DR

This paper introduces EMDiffusion, an EM-based method that trains diffusion models from corrupted data, enabling effective image reconstruction in inverse problems without requiring large clean datasets.

Contribution

The paper presents a novel EM algorithm for training diffusion models directly from corrupted observations, bypassing the need for clean training data.

Findings

Achieves state-of-the-art results on inpainting, denoising, and deblurring tasks.

Effectively learns clean data distribution from corrupted images.

Demonstrates robustness across diverse imaging inverse problems.

Abstract

Diffusion models excel in solving imaging inverse problems due to their ability to model complex image priors. However, their reliance on large, clean datasets for training limits their practical use where clean data is scarce. In this paper, we propose EMDiffusion, an expectation-maximization (EM) approach to train diffusion models from corrupted observations. Our method alternates between reconstructing clean images from corrupted data using a known diffusion model (E-step) and refining diffusion model weights based on these reconstructions (M-step). This iterative process leads the learned diffusion model to gradually converge to the true clean data distribution. We validate our method through extensive experiments on diverse computational imaging tasks, including random inpainting, denoising, and deblurring, achieving new state-of-the-art performance.