💻 computer science

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

This paper introduces EMDiffusion, an expectation-maximization framework that enables the training of high-quality diffusion models directly from corrupted observations by iteratively reconstructing clean images and refining model weights, thereby achieving state-of-the-art performance in various computational imaging tasks without requiring clean training data.

Original authors: Weimin Bai, Yifei Wang, Wenzheng Chen, He Sun

Published 2026-06-29

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Weimin Bai, Yifei Wang, Wenzheng Chen, He Sun

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: The "Chicken and Egg" Dilemma

Imagine you are trying to teach a robot how to draw perfect, high-quality portraits.

The Catch: To teach the robot, you need a huge library of perfect, clean portraits.
The Reality: In many real-world situations (like medical scans or old photos), you don't have perfect pictures. You only have corrupted ones: blurry, noisy, or with big chunks missing (like a puzzle with missing pieces).

This creates a frustrating loop:

To fix the bad pictures, you need a smart robot (a "Diffusion Model") that knows what a clean picture looks like.
But to teach the robot to be smart, you need clean pictures to start with.

If you don't have clean pictures, you can't train the robot. If you don't have the robot, you can't fix the pictures. It's a classic chicken-and-egg problem.

The Solution: EMDiffusion (The "Guess and Refine" Loop)

The authors propose a clever new method called EMDiffusion to break this loop. They use a strategy called Expectation-Maximization (EM), which is essentially a "Guess and Refine" cycle.

Think of it like a detective trying to solve a crime using only blurry security footage.

Step 1: The "E-Step" (The Guess)

The Setup: The detective starts with a tiny, vague idea of what a person looks like (trained on just a few clean photos, maybe 50).
The Action: The detective looks at the blurry security footage and tries to imagine what the person probably looked like. Because the initial idea is weak, the detective might guess wrong at first.
The Trick: To stop the detective from just guessing the same few faces they know, they add a "reality check." They force the guess to match the blurry evidence (the corrupted data) more closely.
Result: They produce a "reconstructed" image. It's not perfect yet, but it's cleaner than the original blurry photo.

Step 2: The "M-Step" (The Refine)

The Action: Now, the detective takes those "reconstructed" images (which are better than the blurry ones) and uses them to re-train their mental model of what a person looks like.
The Update: The robot learns from these new, slightly cleaner guesses. It updates its internal "rulebook" to be more accurate.
Result: The robot is now smarter than it was before.

The Cycle

The process repeats:

Guess: Use the smarter robot to clean the blurry photos again.
Refine: Use the new, even cleaner photos to make the robot even smarter.

With every loop, the robot gets better at guessing, and the guesses get better at training the robot. Eventually, the robot learns to see the "true" image hidden inside the noise, even though it never saw a single perfect photo during the main training phase.

Why This is Special

Usually, if you try to train an AI on bad data, it learns bad habits (like thinking all cars look like blurry blobs).

The Secret Sauce: The authors found that starting with a tiny amount of clean data (like 50 images) acts as a "seed." It prevents the robot from going off the rails.
Adaptive Balancing: The method automatically adjusts how much it trusts the "guess" versus how much it trusts the "blurry evidence." In the beginning, it trusts the evidence more to avoid hallucinations. As the robot gets smarter, it trusts its own knowledge more.

The Results

The team tested this on three types of "corrupted" data:

Inpainting: Filling in missing parts of an image (like a puzzle).
Denoising: Removing static or grain from a photo.
Deblurring: Fixing a photo that was taken while the camera was shaking.

In all cases, their method worked significantly better than previous attempts that tried to learn from bad data. They managed to create a robot that could generate high-quality, clean images, effectively solving the chicken-and-egg problem without needing a massive library of perfect photos to start with.

In a Nutshell

EMDiffusion is a self-improving system. It starts with a small hint of what "good" looks like, uses that hint to clean up "bad" data, and then uses the cleaned-up data to learn how to be even better at cleaning. It's like a student who starts with a rough sketch, uses it to practice drawing, and gradually becomes a master artist, all without ever having a perfect textbook to copy from.

Problem Statement
Diffusion models (DMs) have demonstrated exceptional capability in solving computational imaging inverse problems by modeling complex image priors. However, their practical application is severely constrained by a reliance on large-scale, high-quality, clean datasets for pre-training. In many real-world scenarios, particularly in scientific and biomedical imaging (e.g., X-ray CT, cryo-EM), acquiring such clean data is infeasible or prohibitively expensive. Conversely, corrupted observations (noisy, blurred, or masked) with differentiable forward models are abundant. This creates a "chicken-egg" dilemma: training an accurate diffusion model requires clean images, yet reconstructing clean images from corrupted observations requires a well-trained diffusion model. Existing methods for learning from corrupted data, such as AmbientDiffusion or SURE-Score, often struggle with specific corruption types (e.g., large-scale inpainting or significant noise) or fail to generalize across diverse inverse problems.

Methodology: EMDiffusion
The authors propose EMDiffusion, a novel Expectation-Maximization (EM) framework designed to jointly solve imaging inverse problems and train a diffusion model directly from corrupted observations. The method operates through an iterative cycle of two steps:

Initialization: The process begins with a "vague" diffusion prior trained on a very small set of clean images (as few as 10–50 images). This initialization is crucial to prevent the model from collapsing into a distribution defined solely by corrupted inputs. The authors note that these initial clean images need not be from the same distribution as the target task (out-of-distribution data can serve as effective initializers).
E-Step (Expectation): Using the current diffusion model as a prior, the algorithm performs Adaptive Diffusion Posterior Sampling to reconstruct clean images from the corrupted observations.
- Standard posterior sampling often fails in early iterations because the prior is weak, leading to "mode collapse" where samples ignore the observation and revert to the training distribution.
- To mitigate this, EMDiffusion introduces a hyper-parameter $\lambda$ that rescales the data likelihood term in the reverse-time stochastic differential equation (SDE).
- $\lambda$ is adaptively tuned in each E-step to minimize the data loss ( $\|y - A(\hat{x}_0)\|_2^2$ ), balancing the influence of the imperfect prior against the observed data. This allows the generation of cleaner reconstructions even when the prior is initially poor.
M-Step (Maximization): The algorithm updates the weights of the diffusion model using the reconstructed images (posterior samples) generated in the E-step.
- This step resembles standard denoising score matching, where the model learns to approximate the score function of the reconstructed "clean" data.
- To accelerate training and avoid overfitting to noisy early reconstructions, the method employs a two-phase training strategy: initially, it fine-tunes the inherited weights from the previous iteration; once the reconstruction quality improves (indicated by the optimal $\lambda$ dropping below 1), the model is reinitialized and retrained from scratch on the full set of posterior samples.

Key Contributions

EM Framework for Corrupted Data: The paper introduces the first general framework to train score-based diffusion models from arbitrarily corrupted observations using an EM approach, effectively breaking the dependency on large clean datasets.
Adaptive Posterior Sampling: The proposal of an adaptive scaling factor ( $\lambda$ ) for the likelihood term in the reverse SDE addresses the mode collapse issue inherent in early-stage EM iterations, enabling stable reconstruction even with weak priors.
Iterative Refinement: The method demonstrates that a small amount of clean data can act as an implicit regularizer, allowing the model to progressively converge toward the true clean data distribution through alternating reconstruction and training steps.
Generalizability: Unlike prior works tailored to specific corruptions (e.g., masking), EMDiffusion provides a unified approach applicable to random inpainting, denoising, and deblurring.

Experimental Results
The authors validate EMDiffusion on CIFAR-10 and CelebA datasets across three tasks: random inpainting (60% pixel masking), Gaussian denoising, and Gaussian deblurring.

Inpainting: EMDiffusion significantly outperforms baselines like AmbientDiffusion and SURE-Score, achieving reconstruction quality (PSNR/LPIPS) comparable to Diffusion Posterior Sampling (DPS) that uses a prior trained on the full clean dataset.
Denoising & Deblurring: The method surpasses SURE-Score and self-supervised baselines (Noise2Self), recovering fine details (e.g., car wheels, facial features) that other methods smooth over.
Generative Quality: The learned diffusion models achieve state-of-the-art FID scores compared to other methods trained on corrupted data. Notably, a model learned from 50,000 corrupted images via EM performs better than a model trained on approximately 15,000 clean images.
Efficiency: The method achieves these results using a standard DDPM architecture (35.7M parameters), which is significantly smaller than the architectures used by competing baselines (e.g., AmbientDiffusion's 113M+ parameters).

Significance and Claims
The paper claims that EMDiffusion demonstrates it is information-theoretically possible to learn the underlying clean data distribution primarily from corrupted measurements. The significance lies in:

Overcoming Data Scarcity: It offers a practical solution for domains where clean data is unavailable, leveraging the abundance of corrupted observations.
Robust Initialization: It establishes that a minimal set of clean images (even out-of-distribution) is sufficient to initialize the manifold, preventing the model from collapsing into a biased distribution of corrupted inputs.
State-of-the-Art Performance: The method achieves new state-of-the-art performance in inverse imaging tasks without requiring pre-trained models on large clean datasets, bridging the gap between theoretical priors and practical data constraints.

The authors maintain modesty regarding the scope, acknowledging that the method relies on the assumption that the forward model is known and differentiable, and that future work could extend to learning unknown forward models or noise statistics.