U-Net vs Transformer: Is U-Net Outdated in Medical Image Registration?

TL;DR

This study demonstrates that a modified U-Net with large kernels can match or surpass transformer-based models in 3D medical image registration, challenging the notion that transformers are necessary for long-range modeling.

Contribution

The paper introduces LKU-Net, a large kernel U-Net that enhances receptive field, showing it can outperform transformer-based networks in medical image registration tasks.

Findings

LKU-Net outperforms TransMorph on the IXI dataset.

U-Net's receptive field is sufficient for accurate registration.

LKU-Net ranks first on the MICCAI Learn2Reg leaderboard.

Abstract

Due to their extreme long-range modeling capability, vision transformer-based networks have become increasingly popular in deformable image registration. We believe, however, that the receptive field of a 5-layer convolutional U-Net is sufficient to capture accurate deformations without needing long-range dependencies. The purpose of this study is therefore to investigate whether U-Net-based methods are outdated compared to modern transformer-based approaches when applied to medical image registration. For this, we propose a large kernel U-Net (LKU-Net) by embedding a parallel convolutional block to a vanilla U-Net in order to enhance the effective receptive field. On the public 3D IXI brain dataset for atlas-based registration, we show that the performance of the vanilla U-Net is already comparable with that of state-of-the-art transformer-based networks (such as TransMorph), and that the proposed LKU-Net outperforms TransMorph by using only 1.12% of its parameters and 10.8% of its mult-adds operations. We further evaluate LKU-Net on a MICCAI Learn2Reg 2021 challenge dataset for inter-subject registration, our LKU-Net also outperforms TransMorph on this dataset and ranks first on the public leaderboard as of the submission of this work. With only modest modifications to the vanilla U-Net, we show that U-Net can outperform transformer-based architectures on inter-subject and atlas-based 3D medical image registration. Code is available at https://github.com/xi-jia/LKU-Net.