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

TL;DR

This paper demonstrates that smaller, simpler diffusion models trained on natural images can perform effectively in medical image reconstruction, offering robustness and generality over larger, task-specific models.

Contribution

It introduces the effectiveness of small, natural image-trained diffusion models for medical reconstruction and highlights the importance of model complexity adaptation.

Findings

Small models perform nearly as well as large models in in-distribution reconstruction.

Natural image-trained models outperform medical image-trained models in out-of-distribution cases.

Smaller models are more robust to distribution shifts.

Abstract

Diffusion model have been successfully applied to many inverse problems, including MRI and CT reconstruction. Researchers typically re-purpose models originally designed for unconditional sampling without modifications. Using two different posterior sampling algorithms, we show empirically that such large networks are not necessary. Our smallest model, effectively a ResNet, performs almost as good as an attention U-Net on in-distribution 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, out-performing model trained on medical images in out-of-distribution 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. Moreover, we argue that a key step towards a general diffusion-based prior is training on natural images.