Cross-Resolution Distillation for Efficient 3D Medical Image Registration

Abstract

Images captured in clinic such as MRI scans are usually in 3D formats with high spatial resolutions. Existing learning-based models for medical image registration consume large GPU memories and long inference time, which is difficult to be deployed in resource-limited diagnosis scenarios. To address this problem, instead of shrinking the model size as in previous works, we turn to reducing the input resolution of existing registration models and boosting their performance through knowledge distillation. Specifically, we propose a cross-resolution distillation (CRD) scheme, which is designed to train low-resolution models under the guidance of corresponding high-resolution models. Nevertheless, due to the resolution gap between features in high/low-resolution models, straightforward distillation is difficult to apply. To overcome this challenge, we first introduce a feature-shifted teacher (FST) to shift and fuse features of high/low-resolution models. Then, we exploit this teacher model to guide the learning of the low-resolution student model with distillation losses on both features and deformation fields. Finally, we only need to use the distilled student model during inference. Experimental results on four 3D medical image datasets demonstrate that the low-resolution models trained through our CRD scheme use fewer than 20% GPU memories and less than 20% inference time while achieving competitive performance compared with corresponding high-resolution models.