GSURE-Based Diffusion Model Training with Corrupted Data

Published: 15 Apr 2024, Last Modified: 15 Apr 2024Accepted by TMLREveryoneRevisionsBibTeX
Abstract: Diffusion models have demonstrated impressive results in both data generation and downstream tasks such as inverse problems, text-based editing, classification, and more. However, training such models usually requires large amounts of clean signals which are often difficult or impossible to obtain. In this work, we propose a novel training technique for generative diffusion models based only on corrupted data. We introduce a loss function based on the Generalized Stein’s Unbiased Risk Estimator (GSURE), and prove that under some conditions, it is equivalent to the training objective used in fully supervised diffusion models. We demonstrate our technique on face images as well as Magnetic Resonance Imaging (MRI), where the use of undersampled data significantly alleviates data collection costs. Our approach achieves generative performance comparable to its fully supervised counterpart without training on any clean signals. In addition, we deploy the resulting diffusion model in various downstream tasks beyond the degradation present in the training set, showcasing promising results.
Submission Length: Regular submission (no more than 12 pages of main content)
Changes Since Last Submission: As requested by AE
Code: https://github.com/bahjat-kawar/gsure-diffusion
Supplementary Material: pdf
Assigned Action Editor: ~Bertrand_Thirion1
Submission Number: 1752
Loading