Research map / Image generation / Image restoration
An Expectation-Maximization Algorithm for Training Clean Diffusion Models from Corrupted Observations
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes EMDiffusion, an expectation-maximization loop that trains a clean diffusion prior from corrupted observations (inpainting, denoising, deblurring), starting from a small set of clean images.
Problem
Training diffusion models for imaging inverse problems normally requires large clean datasets, which are often unavailable; the paper asks whether a diffusion prior can be learned primarily from corrupted observations with known forward models.
Approach
Treat clean images as latent variables: alternate E-steps that reconstruct clean images from corrupted data via diffusion posterior sampling with an adaptively chosen likelihood scale λ, and M-steps that retrain the score network on those reconstructions, initialized from a diffusion model trained on 50 clean images.
Key results
- On CIFAR-10 inpainting (60% masked), FID of the learned prior is 21.08 versus 28.88 for AmbientDiffusion, with PSNR 24.70 vs 20.57 for AmbientDiffusion.
- On CIFAR-10 denoising (sigma=0.2), PSNR 23.16 and LPIPS 0.022, outperforming SURE-Score (22.42 PSNR) but trailing DPS with clean prior (25.91).
- On CelebA deblurring, PSNR 23.74 vs 22.07 for SURE-Score; the learned prior's FID (91.89) is worse than the blurry observations' FID (72.83).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Train from corrupted observations with EMDiffusion instead of AmbientDiffusion | CIFAR-10 inpainting · CIFAR-10 inpainting | FID: 28.88 -> 21.08 | Helped |
| Train from corrupted observations with EMDiffusion instead of SURE-Score | CIFAR-10 denoising · CIFAR-10 denoising | PSNR: 22.42 -> 23.16 | Helped |
| Use EMDiffusion instead of DPS with a clean prior | CIFAR-10 denoising · CIFAR-10 denoising | PSNR: 25.91 -> 23.16 | Hurt |
| Train from corrupted observations with EMDiffusion instead of SURE-Score | CelebA deblurring · CelebA deblurring | PSNR: 22.07 -> 23.74 | Helped |
| Use EMDiffusion instead of DPS with a clean prior | CelebA deblurring · CelebA deblurring | PSNR: 29.05 -> 23.74 | Hurt |
| Initialize with 10 clean CIFAR-10 images instead of 500 (ablation) | EMDiffusion · CIFAR-10 inpainting | PSNR: 20.58 -> 20.21 | No clear effect |
| Initialize with 10 clean CIFAR-10 images instead of 500 (ablation) | EMDiffusion · CIFAR-10 denoising | PSNR: 19.89 -> 19.57 | No clear effect |
| Initialize with 10 clean CelebA images instead of 500 (ablation) | EMDiffusion · CelebA deblurring | PSNR: 25.56 -> 21.19 | Hurt |
| Initialize with 50 out-of-distribution CelebA images instead of 50 in-distribution CIFAR-10 images (ablation) | EMDiffusion · CIFAR-10 inpainting | PSNR: 20.93 -> 20.10 | No clear effect |
| Initialize with 50 out-of-distribution CelebA images instead of 50 in-distribution CIFAR-10 images (ablation) | EMDiffusion · CIFAR-10 denoising | PSNR: 20.47 -> 19.57 | No clear effect |
| Initialize with 50 out-of-distribution CelebA images instead of 50 in-distribution CelebA images (ablation) | EMDiffusion · CelebA deblurring | PSNR: 21.19 -> 25.50 | Helped |
| Increase the inpainting masking ratio from 0.4 to 0.8 (ablation) | EMDiffusion · CIFAR-10 inpainting | FID: 13.75 -> 45.24 | Hurt |
Surprising findings
- A diffusion model initialized from only 50 clean images from an out-of-distribution dataset (CelebA) performs comparably to in-distribution initialization on CIFAR-10 inpainting and denoising.
- The EM-learned CIFAR-10 prior from 50,000 corrupted images performs similarly to a prior trained on roughly 15,000-20,000 clean images.
Limitations
- Relies on a small initial set of clean images and known forward operator and noise statistics.
- Posterior sampling (DPS) can introduce artifacts that hurt FID, especially for deblurring, and the E-step quality limits the M-step.
- Reported results are single runs with no error bars; deblurring FID is worse than the observations.
Experimental setup
Datasets: CIFAR-10 CelebA 50 randomly selected clean images per dataset for initialization
Benchmarks: CIFAR-10 inpainting CIFAR-10 denoising CelebA deblurring
Metrics: PSNR LPIPS FID
Baselines: AmbientDiffusion SURE-Score DPS with clean prior Noise2Self
Builds on: DDPM score-based SDE Diffusion Posterior Sampling Ambient Diffusion SURE-Score
Closest papers by meaning
- Learning Diffusion Priors from Observations by Expectation Maximization (NeurIPS 2024 · similarity 0.91)
- A Variational Perspective on Solving Inverse Problems with Diffusion Models (ICLR 2024 · similarity 0.90)
- A Mixture-Based Framework for Guiding Diffusion Models (ICML 2025 · similarity 0.89)
- A Diffusion Model with State Estimation for Degradation-Blind Inverse Imaging (AAAI 2024 · similarity 0.89)
- Unleashing the Denoising Capability of Diffusion Prior for Solving Inverse Problems (NeurIPS 2024 · similarity 0.89)
- Ambient Diffusion Posterior Sampling: Solving Inverse Problems with Diffusion Models Trained on Corrupted Data (ICLR 2025 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.