Optimizing Loss Function for Uni-modal and Multi-modal Medical Registration

Abstract

Recent learning-based methods render fast registration by leveraging deep networks to directly learn the spatial transformation fields between the source and target images. However, manually designing and tuning loss functions for multiple types of medical data need intensive labor and extensive experience, and automatic design of loss functions remains under-investigated. In this paper, we introduce a unified formulation of the loss function and raise automated techniques to search hyperparameters of losses to obtain the optimal loss function. Specifically, we take into consideration of the multifaceted properties of image pairs and propose a unified loss to constraint similarity from different aspects. Then, we propose a bilevel self-tuning training strategy, allowing the efficient search of hyperparameters of the loss function. Based on the adaptive degrees, the proposed unified loss would be applicable to the registration of arbitrary modalities and multiple tasks. Moreover, this training strategy also reduces computational and human burdens. We conduct uni-modal and multi-modal registration experiments on seven 3D MRI datasets, the networks trained with the searched loss functions deliver accuracy on par or even superior to those with the handcrafted losses. Extensive results demonstrate our advantages over state-of-the-art registration techniques in terms of accuracy with efficiency.