EMAC-Diffusion: Learning Clean Distributions from Unknown Corruptions via Expectation Maximization

We introduce EMAC, an expectation-maximization (EM) algorithm for learning a clean data distribution from a small clean dataset and abundant observations corrupted by an unknown channel. Our method introduces an auxiliary-corruption channel that enables us to train a conditional diffusion model using EM without access to the original corruption channel. We apply this channel to both clean and already-corrupted images in the dataset and impose specific constraints on its design. We theoretically and experimentally identify limitations of prior work in handling low-frequency corruptions and we establish convergence guarantees for our method under appropriate assumptions. Experimentally, we first validate our method in controlled experiments, show how it can be used in the context of a privacy-related application, and finally test it at scale by improving imperfect data that naturally occur in ImageNet at $512\times 512$ resolution. Qualitative and quantitative analyses show that our method improves over previously proposed baselines on several image-quality metrics.