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

NeurIPS 2024

Weimin Bai1,2,3, Yifei Wang4, Wenzheng Chen5,6, He Sun1,2,3
1AAIS, Peking University 2College of Future Technology, Peking University 3Beijing National Research Center for Information Science and Technology 4Wangxuan Institute of Computer Technology, Peking University 5Academy for Advanced Interdisciplinary Studies, Peking University 6State Key Lab of Multimedia Information Processing, Peking University
Paper Code Citation

Overview

EMDiffusion is designed for settings where clean training data are scarce but corrupted observations are abundant. It alternates posterior reconstruction and model updating, which gradually improves both the clean prior and inverse-task quality.

Links: Paper PDF | arXiv | Code

EMDiffusion overview diagram

Fig. 1. EM loop: sample clean posteriors in E-step, then update the score prior in M-step.

Abstract

Diffusion models are powerful priors for imaging inverse problems, but training clean diffusion models typically requires large clean datasets that are often unavailable in real-world settings. EMDiffusion introduces an expectation-maximization framework to train diffusion priors directly from corrupted observations. In each E-step, the method reconstructs posterior clean samples from corrupted data; in each M-step, it updates the diffusion model on those samples. This alternating process progressively improves both reconstruction quality and the learned diffusion prior, and enables strong performance on inpainting, denoising, and deblurring.

More: Method details | Main results | BibTeX

Method Summary

The algorithm uses adaptive posterior sampling to recover latent clean images from corrupted inputs, then fits the score model to these reconstructions. This coupling allows robust training even from highly degraded observations.

References: NeurIPS paper | Implementation

Adaptive posterior sampling with different lambda values

Fig. 2. Adaptive posterior sampling: balanced scaling improves posterior quality over naive settings.

Algorithm 1 of EMDiffusion

Algorithm 1. Full EMDiffusion training pipeline with alternating E-step/M-step updates.

Qualitative Results

Qualitative comparisons show consistent gains across inpainting, denoising, and deblurring. Reconstructions become cleaner over EM iterations and approach methods trained with clean priors.

References: Figure 3 & Figure 4 in paper | Reproduce with code

CIFAR-10 inpainting results over EM iterations

Fig. 3. CIFAR-10 inpainting: output quality improves step-by-step along EM iterations.

Denoising and deblurring qualitative comparisons

Fig. 4. Denoising and deblurring: EMDiffusion recovers sharper structures and more faithful details.

  • No large clean dataset required: learn useful clean priors directly from corrupted data.
  • General framework: supports random inpainting, Gaussian denoising, and Gaussian deblurring.
  • Iterative improvement: both reconstruction quality and prior generation quality improve across EM iterations.

Quantitative Results

EMDiffusion outperforms corrupted-data baselines and narrows the gap to methods that rely on external clean priors. The gains are consistent across metrics and tasks.

See also: Table 1 in paper | Evaluation scripts

24.70 / 0.009 / 21.08

CIFAR-10 Inpainting
(PSNR / LPIPS / FID)

23.16 / 0.022 / 86.47

CIFAR-10 Denoising
(PSNR / LPIPS / FID)

23.74 / 0.103 / 91.89

CelebA Deblurring
(PSNR / LPIPS / FID)

Main quantitative table across inpainting denoising and deblurring

Table 1. Best quantitative trade-off on all three tasks under corrupted-data training.

Ablation Studies

Ablations analyze initialization quality, EM iteration effects, and scaling-factor selection. These results explain why iterative updates and adaptive balancing are critical to performance.

References: Figure 5 in paper

Ablation studies on EM iterations and scaling factor lambda

Fig. 5. EM iterations improve priors, and adaptive λ selection stabilizes reconstruction quality.

Citation

@inproceedings{bai2024emdiffusion,
  title     = {An Expectation-Maximization Algorithm for Training Clean Diffusion Models from Corrupted Observations},
  author    = {Bai, Weimin and Wang, Yifei and Chen, Wenzheng and Sun, He},
  booktitle = {Advances in Neural Information Processing Systems},
  year      = {2024}
}