Intensity-Based Wasserstein Distance As A Loss Measure For Unsupervised Deformable Deep Registration

摘要

Traditional pairwise medical image registration techniques are based on computationally intensive frameworks due to numerical optimization procedures. While there is increasing adoption of deep neural networks to improve deformable image registration, achieving a clinically suitable solution remains scarce. One of the primary difficulties lies in the choice of tractable distance functions to assess image similarity. Recent works have explored the Wasserstein distance as a loss function in generative deep neural networks. In this work, we evaluate a fast approximation variant - the sliced Wasserstein distance - for deep image registration of brain MRI datasets. Based on a VoxelMorph backbone architecture, which includes a combination of UNet and spatial transformer networks (STN) for deformable registration, we propose three implementation variants to compare the model's performance: the standard sliced Wasserstein, the Radon transform performing a low dimensional embedding, and a novel patch-based method that allows fine-grained deformation comparison. Experiments performed on public datasets of brain images from the Learn2Reg open challenge demonstrate the Wasserstein methods converge faster than the baseline mean square error method, with the proposed patch-based method yielding similar performance to baseline methods, and improved overall accuracy compared with other implementations. This makes the sliced Wasserstein a valuable metric for deep mono-modal and multi-modal deformable medical image registration problems with our proposed implementation.