EMDiffusion-PMC

Learning Diffusion Model from Noisy Measurement using Principled Expectation–Maximization Method (ICASSP 2025)

This repository is a PMC-oriented extension of EMDiffusion, built on top of the official implementation
ai4imaging/EMDiffusoin, and tailored for our principled EM framework for learning diffusion models directly from noisy measurements (arXiv:2410.11241).

The key change compared with the original EMDiffusion codebase is:

• E-step: performs the Monte Carlo estimation of clean images from noisy measurements under our PMC setting.
• M-step: we keep the original m-step.sh training script (gradient-based refinement of diffusion model parameters).
• EM driver: alternates between e-step.sh and m-step.sh for multiple outer EM iterations.

---

1. Environment

The environment requirement is essentially the same as the original EMDiffusion implementation and DPS:

• A working PyTorch + GPU environment.
• A conda environment compatible with DPS; in our examples we use a conda env named DPS.
• For multi-GPU training in the M-step we rely on HuggingFace Accelerate.

Please refer to:

• Original DPS repo https://github.com/DPS2022/diffusion-posterior-sampling for low-level E-step environment.
• Original EMDiffusion repo https://github.com/ai4imaging/EMDiffusoin for additional training details.

---

2. E-step: recovering signals from noisy measurements

• It activates the DPS conda environment.
• It calls:

bash e-step.sh

• The key experiment-specific hyperparameters you should manually tune in e-step.sh include:
  • --model_path: path to the conditional diffusion checkpoint used as prior.
  • --task_config: measurement / inpainting configuration YAML.
  • --batch_size, --batch_num, --start, --seed.
  • --save_dir: where reconstructed “clean” samples (posterior draws) are saved and later used by the M-step.

after adjusting the scheduler directives and paths.

---

3. M-step: diffusion model update

• Activates the DPS conda environment.
• Launches denoiser_diffusion.py with Accelerate, e.g.:

bash m-step.sh

Here you should manually tune:

• --dataset_path: path to the reconstructed “clean” images produced by the latest E-step (or a dataset built from them).
• --model_path: path to the diffusion model checkpoint you are updating.
• Any other training hyperparameters inside denoiser_diffusion.py / config files (learning rate, total steps, etc.).

depending on whether you want to submit a job or run interactively.

---

4. EM loop driver: run_em.sh

We provide a minimal EM loop implementation in run_em.sh:

• It defines:

  NUM_EM_ITERS=5
  E_STEP_SCRIPT="./e-step.sh"
  M_STEP_SCRIPT="./m-step.sh"

• Then runs:

  for iter in 1..NUM_EM_ITERS:
      bash e-step.sh
      bash m-step.sh

Usage:

1. Edit e-step.sh and m-step.sh to set all experiment-specific hyperparameters and paths.

2. Optionally edit NUM_EM_ITERS in run_em.sh to control how many outer EM iterations you run.

3. From the repo root:

   bash run_em.sh

This script is intentionally simple so that you can easily adapt it to:

• Submit E-step and M-step as separate jobs.
• Insert monitoring, logging, or intermediate evaluation between steps.

---