Patch-Based Diffusion Models Beat Whole-Image Models for Mismatched Distribution Inverse Problems
This addresses a critical limitation in inverse problem solving for image processing, enabling more robust reconstructions in out-of-distribution scenarios, though it is incremental as it builds on existing diffusion model frameworks.
The paper tackles the problem of artifacts and hallucinations in diffusion model-based image reconstruction when training and test distributions are mismatched, showing that a patch-based diffusion prior with self-supervised loss achieves high-quality reconstructions, outperforming whole-image models and competing with methods using large in-distribution datasets.
Diffusion models have achieved excellent success in solving inverse problems due to their ability to learn strong image priors, but existing approaches require a large training dataset of images that should come from the same distribution as the test dataset. When the training and test distributions are mismatched, artifacts and hallucinations can occur in reconstructed images due to the incorrect priors. In this work, we systematically study out of distribution (OOD) problems where a known training distribution is first provided. We first study the setting where only a single measurement obtained from the unknown test distribution is available. Next we study the setting where a very small sample of data belonging to the test distribution is available, and our goal is still to reconstruct an image from a measurement that came from the test distribution. In both settings, we use a patch-based diffusion prior that learns the image distribution solely from patches. Furthermore, in the first setting, we include a self-supervised loss that helps the network output maintain consistency with the measurement. Extensive experiments show that in both settings, the patch-based method can obtain high quality image reconstructions that can outperform whole-image models and can compete with methods that have access to large in-distribution training datasets. Furthermore, we show how whole-image models are prone to memorization and overfitting, leading to artifacts in the reconstructions, while a patch-based model can resolve these issues.