Skip to main content
Ido AvitalAI Engineer · Data Scientist
Menu

AI systems

Diffusion Lab

From DDPM to zero-shot image restoration

The DDPM algorithm written from its equations and trained to generate digits, then one projection added to a pretrained face model's sampler to inpaint, super-resolve and colourise without training.

Key result The from-scratch DDPM turns pure noise into recognisable digits after 50 epochs on MNIST. The restoration results are qualitative: no restoration metric was computed.

Research study Status: Completed Academic · team of 3

Pure noise Predictclean image Projectonto data One DDPM step

Workflow diagram

Research question

Two questions: can a short, exact implementation of DDPM turn noise into handwritten digits, and can a model trained only to generate faces restore damaged images it never saw, with no further training?

My Role & Leadership

Project Owner · Technical Lead · Developer. Owned and implemented the diffusion algorithm and the DDNM sampler of a course assignment: the noise schedule, the one-step forward process, the noise-prediction loss and the reverse sampler, then the DDNM reverse loop with its range and null-space correction. The assignment was done in a course team of three; this implementation is my hands-on part.

  • Led the implementation of the projection that keeps the measured part of the image fixed and leaves the rest to the diffusion prior, for three degradation operators.

Key findings

With no further training, the same DDNM sampler fills a masked region, super-resolves an 8x thumbnail and colourises a grayscale image: only the degradation operator and its pseudo-inverse change.

Method

One sampler, three operators. The projection keeps the range space of the degradation fixed to the measurement and leaves the null space to the diffusion prior.

  1. Write the forward process: a noise schedule and the closed-form jump from a clean image to any noise level.

  2. Train a network to predict the added noise, then sample by running the reverse process for 1,000 steps.

  3. For restoration, take a pretrained face model and, after each predicted clean image, project it onto the measurement.

  4. Use the same sampler for three degradations by swapping the operator and its pseudo-inverse.

What the team built

  • A linear noise schedule (1e-4 to 0.02 over 1,000 steps) and the closed-form jump to any noise level.
  • The noise-prediction training objective and the 1,000-step ancestral sampler, trained for 50 epochs on MNIST.
  • The DDNM sampler: after each predicted clean image, x0 is corrected to x0 + A+(y - A x0), so the measured part matches the data and the model fills only what the measurement cannot see.
  • Three degradation operators with simple pseudo-inverses: a mask (inpainting), 8x average pooling (super-resolution) and channel averaging (colourisation).

Methodological choices

  • Restoration without training: the projection uses only the degradation operator and its pseudo-inverse, so one generative model serves every task.
  • Operators chosen to have simple, exact pseudo-inverses, so the projection is cheap and correct.

Built on

  • The U-Net, the training loop and the data loader came with the assignment.
  • The restoration model (google/ddpm-celebahq-256) and its posterior step are pretrained third-party components, not trained here.
  • The DDPM and DDNM methods are published work by others; this is an implementation of them.

Limitations

The restoration model is pretrained by Google, not trained here, and the face being restored is itself a sample from that model. No quantitative evaluation. Faces from a model trained on CelebA-HQ are not shown.

Key technologies

  • PyTorch
  • DDPM
  • DDNM (null-space projection)
  • Hugging Face diffusers

Evidence

  • Two assignment notebooks with their saved outputs, and a theory write-up.
  • The images on this page are the notebook's own outputs; no face image is shown.
All projects
  • AI systems

    RefundIt

    Software system Status: Prototype Status: Completed

  • AI systems

    Breaking SAM 2

    Research study Status: Completed

Questions about Diffusion Lab