UTNet: A Hybrid Transformer Architecture for Medical Image Segmentation

TL;DR

UTNet introduces a hybrid Transformer architecture that combines self-attention with convolutional neural networks to improve medical image segmentation, achieving superior performance without pre-training on a cardiac MRI dataset.

Contribution

The paper presents a novel hybrid Transformer architecture with an efficient self-attention mechanism and a new decoder, enabling effective medical image segmentation without pre-training.

Findings

UTNet outperforms state-of-the-art methods in cardiac MRI segmentation.

The model demonstrates robustness across multi-vendor datasets.

Self-attention modules improve long-range dependency capture at multiple scales.

Abstract

Transformer architecture has emerged to be successful in a number of natural language processing tasks. However, its applications to medical vision remain largely unexplored. In this study, we present UTNet, a simple yet powerful hybrid Transformer architecture that integrates self-attention into a convolutional neural network for enhancing medical image segmentation. UTNet applies self-attention modules in both encoder and decoder for capturing long-range dependency at different scales with minimal overhead. To this end, we propose an efficient self-attention mechanism along with relative position encoding that reduces the complexity of self-attention operation significantly from $O(n^2)$ to approximate $O(n)$. A new self-attention decoder is also proposed to recover fine-grained details from the skipped connections in the encoder. Our approach addresses the dilemma that Transformer requires huge amounts of data to learn vision inductive bias. Our hybrid layer design allows the initialization of Transformer into convolutional networks without a need of pre-training. We have evaluated UTNet on the multi-label, multi-vendor cardiac magnetic resonance imaging cohort. UTNet demonstrates superior segmentation performance and robustness against the state-of-the-art approaches, holding the promise to generalize well on other medical image segmentations.