Learning Diffusion Model from Noisy Measurement using Principled Expectation-Maximization Method

Diffusion models (DMs) [2, 1] capture data distributions by approximating the score function, i.e., the gradient of the log-likelihood $\nabla_{\boldsymbol{x}}\log p(\boldsymbol{x})$, using a deep neural network $\boldsymbol{s}_{\theta}$ parameterized by $\theta$. To achieve this, DMs define a forward stochastic differential equation (SDE) that progressively injects noise, while its solution is a reverse-time SDE that progressively removes the noise:

where $\beta_{t}\in(0,1)$ is the noise schedule, $t\in[0,T]$ represents the time index, and $\boldsymbol{w}$ and $\overline{\boldsymbol{w}}$ are the forward and backward Wiener processes, respectively. Once a DM is trained on large-scale clean data, diverse samples can be generated by replacing $\nabla_{\boldsymbol{x}_{t}}\log p_{t}(\boldsymbol{x}_{t})$ in Eq. 1 with the learned score function $\boldsymbol{s}_{\theta}(\boldsymbol{x}_{t},\sigma(t))$, where $\sigma(t)$ is the pre-defined noise strength corresponding to $\beta_{t}$.

By Bayes’ rule, $p(\boldsymbol{x}|\boldsymbol{y})\propto p(\boldsymbol{y}|\boldsymbol{x})p(\boldsymbol{x})$, the reverse-time SDE can also be adapted to sample from the posterior distribution of images $\boldsymbol{x}$ conditioned on measurements $\boldsymbol{y}$. In this case, $p(\boldsymbol{y}|\boldsymbol{x})$ acts as a data-fidelity term to ensure the reconstructed images match the measurements. The score function from Eq. 1 is then modified as a conditional version:

where $p_{t}\left(\boldsymbol{y}|\boldsymbol{x}_{t}\right)$ is a time-dependent data-fidelity term distinct from $p(\boldsymbol{y}|\boldsymbol{x})$ and is generally intractable. The challenge in posterior sampling via diffusion models is bridging $p_{t}\left(\boldsymbol{y}|\boldsymbol{x}_{t}\right)$ with $p(\boldsymbol{y}|\boldsymbol{x})$. In the widely used diffusion posterior sampling (DPS) algorithm [11], we empirically assume that:

However, this approximation lacks theoretical guarantees for convergence to the true posterior distribution, which can result in inaccurate sampling, even in simple cases.

In this paper, we propose a principled expectation-maximization (EM) framework for learning clean diffusion models from noisy measurements. This approach inherits basic procedure of the standard EM algorithm [20, 21] and EMDiffusion [19], alternating between two steps: 1) the E-step, where we sample the posterior distribution of clean images from noisy measurements using a known diffusion model, and 2) the M-step, where we refine the DM using these recovered samples. This iterative process guides the DM towards a local minimum, progressively improving both the model and the quality of the recovered images.

Datasets. We evaluate our method on two datasets with distinct characteristics: CIFAR10 [26] at a resolution of $32\times 32$ and CelebA [27] at $64\times 64$. Noisy measurements $\boldsymbol{y}$ are generated by corrupting clean images $\boldsymbol{x}$ from the training sets. In the first six iterations, 5,000 measurements are randomly selected for posterior sampling and diffusion model refinement. In the last three iterations, all measurements are used for posterior sampling and training diffusion model from scratch. We randomly choose 250 clean images from test sets for evaluation of image reconstruction and original training sets for evaluation of image generation.

Baselines. We compare our method with four related baselines: Ambient Diffusion [15], SURE-Score [18], EMDiffusion [19], and DPS with a clean prior [11]. Ambient Diffusion, SURE-Score, and EMDiffusion have similar settings to our method, as they do not rely on large-scale clean data to train DMs, but instead train DMs directly from noisy measurements. In contrast, DPS [11] leverages a pre-trained clean diffusion prior for posterior sampling, representing the performance upper bound for our approach.

Metrics. We evaluate image reconstruction quality using two metrics: peak signal-to-noise ratio (PSNR) and learned perceptual image patch similarity (LPIPS). Additionally, we compute the Fréchet Inception Distance (FID) between samples generated by the trained DMs and the reserved clean training set to assess the quality of image generation.

We evaluate our principled expectation-maximization (EM) framework on two datasets across three representative imaging inverse problems: image denoising, random inpainting, and Gaussian deblurring.

Image denoising. For image denoising, we add Gaussian noise $\boldsymbol{n}\sim\mathcal{N}(0,\sigma^{2}\boldsymbol{I})$ with $\sigma=0.2$. As shown in Table I and Fig. 1, our method significantly outperforms all baselines. The reconstructed images exhibit notably fewer artifacts compared to the current state-of-the-art, EMDiffusion, demonstrating the advantage of moving beyond empirical approximations for time-dependent data fidelity.

Image inpainting. For image inpainting, we randomly mask 60% of the pixels in the images. As shown in Table I and Fig. 1, the posterior samples closely match the ground truths, highlighting the robustness of our method. Although our method achieves the second-best performance in the random inpainting task, this is due to EMDiffusion’s simple approximation being particularly effective for this problem [28]. In EMDiffusion, the measurement $\boldsymbol{y}$ guides the updates of all pixels, through $\nabla_{\boldsymbol{x}_{t}}\log p_{t}\left(\boldsymbol{y}|\hat{\boldsymbol{x}}_{0}\right)$, not just the unmasked ones.

Image deblurring. For image deblurring, we apply a $9\times 9$ Gaussian blur kernel with a standard deviation of 2. As shown in Table. I and Fig. 2, our method consistently outperforms all baselines. The posterior samples progressively approximate the ground truths throughout the iterations, verifying the effectiveness of the EM framework.

In Table I and Fig. 3, we compare the quantitative and qualitative results of the diffusion priors learned by our method against the baselines. The FID scores in Table I demonstrate that our method outperforms the baselines in all three tasks: image inpainting and denoising on CIFAR10, and image deblurring on CelebA. The progressively improving quality of the posterior samples in Fig. 2 suggests that the diffusion model is effectively approximating the true prior distribution, validating the effectiveness of our iterative learning framework. As shown in uncurated samples in Fig. 3, our method significantly reduces sharp artifacts compared to EMDiffusion, further confirming its superior performance.