Test-Time Adaptation Improves Inverse Problem Solving With Patch-Based Diffusion Models

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. In practice, the size of the available training dataset can range from nonexistent to very large. In some cases, conventional diffusion model training from limited data can lead to poor reconstruction results due to poorly learned priors. One potential improvement is to start with a diffusion model trained from available training data having a possibly mismatched distribution, and then refine the network at reconstruction time to account for the distribution mismatch. In this work, we investigate the effect of this network refining process on diffusion models trained from varying degrees of out-of-distribution data. Specifically, we use a self-supervised loss to adapt the learned diffusion network to the testing data while helping the network output maintain consistency with the measurements. We show that, both theoretically and experimentally, test-time adaptation of a patch-based diffusion prior leads to higher quality reconstructions than test-time refinement of traditional whole-image diffusion models. Extensive experiments show that across a wide range of inverse problems, test-time adaptation significantly improves image reconstruction quality when there are significant domain shifts between training and testing distributions. Interestingly, even for the in-distribution case, test-time adaptation also significantly improves reconstruction quality.