New Insights into Learning-Based Image Reconstruction under Distribution Shifts
Figure 1: Evolution of MRI image reconstruction under a true and mismatched CNN prior, trained on MRI and natural images, respectively. The numbers on the top-left corner of the images are the SNR (dB) and SSIM values.
There is a growing interest in deep model-based architectures (DMBAs) for solving imaging inverse problems by combining physical measurement models and learned image priors specified using convolutional neural nets (CNNs). For example, well-known frameworks for designing DMBAs include plug-and-play priors (PnP), deep unfolding (DU), and deep equilibrium models (DEQ). While the empirical performance and theoretical properties of DMBAs have been widely investigated, the existing work has primarily focused on their performance when the desired image prior is known exactly. This work addresses this gap by providing new theoretical and numerical insights into DMBAs under mismatched CNN priors. Mismatched priors arise naturally when there is a distribution shift between training and testing data, for example, due to test images being from a different distribution than images used for training the CNN prior (see Figure 1 above). They also arise when the CNN prior used for inference is an approximation of some desired statistical estimator (MAP or MMSE). Our theoretical analysis provides explicit error bounds on the solution due to the mismatched CNN priors under a set of clearly specified assumptions. Our numerical results compare the empirical performance of DMBAs under realistic distribution shifts and approximate statistical estimators.
Figure 2: Reconstruction of the same MRI brain image using CNN priors trained on CT, Natural, and MRI images. Note how all the CNN priors improve over the traditional TV prior. The CT prior is the best mismatched prior for this image.
Figure 3: Reconstruction of a brain MRI image for two subsampling rates (10% and 30%) using several priors. MRI (AWGN) refers to a CNN prior trained as an image denoiser for MRI images. MRI (AR), CT (AR), and Natural (AR) refer to CNN priors trained as artifact-removing (AR) operators within implicit DMBAs using deep equilibrium models (DEQ).
Figure 4: Parallel MRI reconstruction of a brain MRI image from the AXT2 dataset using CNN priors trained on MRI images from other datasets and anatomical regions.
Figure 5: Reconstruction of Parrot using the exact TV prior and its approximation as a pre-trained CNN. The performance of the approximate TV prior is similar that of the true TV prior.