Bigger Isn’t Always Better: Towards a General Prior for Medical Image Reconstruction

Diffusion models play an important role in many state-of-the-art algorithms to solve inverse problems in imaging. In particular, reconstructing MRI and CT images from undersampled data can be reformulated as a canonical inverse problem that benefits from strong priors. Researchers typically re-purpose models originally designed for unconditional sampling without modifications and achieve remarkable accuracy on in-distribution (ID) reconstruction. However, due to the scarce availability of training data and the large number of different imaging setups, increasing the generalization capabilities of diffusion-based priors is key to clinical adoption. To do so, we propose two solutions: (i) using smaller models and (ii) training on natural images. Using three different posterior sampling algorithms, we evaluate the influence of network size and training data. Our smallest model, effectively a ResNet, performs almost as good as an attention U-Net on ID reconstruction, while being significantly more robust towards distribution shifts. Furthermore, we introduce models trained on natural images and demonstrate that they can be used in both MRI and CT reconstruction, outperforming models trained on medical images in OOD cases. As a result of our findings, we strongly caution against simply re-using very large networks and encourage researchers to adapt the model complexity to the respective task.