Transferring Models Trained on Natural Images to 3D MRI Via Position Encoded Slice Models

Transfer learning from large 2D-image datasets can be helpful for neuroimaging tasks for which training set sizes are often small. However, a mismatch in the input space (2D images vs. 3D MRIs) restricts the direct transfer of models. To this end, we leverage the 2D-Slice-CNN architecture of Gupta et al. (2021), which embeds MRI slices with 2D encoders (neural networks that take 2D image input) and processes these slice embeddings via permutation-invariant layers. The key insight is that a pretrained model can be used as the 2D encoder. Our experiments with brain age prediction on the UK Biobank dataset and Alzheimer's disease detection on the ADNI dataset show that the 2D encoders initialized with ImageNet pretrained weights outperform those initialized and trained from scratch. Further, we improve the modeling capabilities of 2D-Slice models by incorporating spatial information through position embeddings. Positional information improves the performance of 2D slice architecture in some cases.