Source-linked AI summary
Convolution-Free Medical Image Segmentation using Transformers
Davood Karimi, Serge Vasylechko, Ali Gholipour
TL;DR
Medical image segmentation has been dominated by convolution-based deep learning, while the potential of convolution-free self-attention remains underexplored. The paper proposes a 3D patch-based self-attention network and finds competitive or better segmentation than a state-of-the-art CNN across three datasets, with stronger advantages when labeled data are scarce and pre-training is used.
Problem
Medical image segmentation research has largely relied on convolutional architectures, leaving the effectiveness of a convolution-free self-attention approach insufficiently explored.
Method
The paper predicts the center-patch segmentation of a 3D image block from self-attention among linear embeddings of neighboring 3D patches, without convolution operations.
Results
The proposed network achieves segmentation accuracy better than or at least on par with a state-of-the-art CNN on three datasets, and pre-training improves performance with fewer labeled images.
Takeaways & Limitations
Convolution-free self-attention is a viable approach for 3D medical image segmentation, particularly when labeled training data are limited and unlabeled images are available for pre-training.
Abstract
from arXiv · showhide
Like other applications in computer vision, medical image segmentation has been most successfully addressed using deep learning models that rely on the convolution operation as their main building block. Convolutions enjoy important properties such as sparse interactions, weight sharing, and translation equivariance. These properties give convolutional neural networks (CNNs) a strong and useful inductive bias for vision tasks. In this work we show that a different method, based entirely on self-attention between neighboring image patches and without any convolution operations, can achieve competitive or better results. Given a 3D image block, our network divides it into $n^3$ 3D patches, where $n=3 \text{ or } 5$ and computes a 1D embedding for each patch. The network predicts the segmentation map for the center patch of the block based on the self-attention between these patch embeddings. We show that the proposed model can achieve segmentation accuracies that are better than the state of the art CNNs on three datasets. We also propose methods for pre-training this model on large corpora of unlabeled images. Our experiments show that with pre-training the advantage of our proposed network over CNNs can be significant when labeled training data is small.
I. INTRODUCTION
Medical image segmentation is important but has become dominated by convolution-based deep learning, motivating exploration of a convolution-free self-attention alternative. The proposed network operates on embeddings of neighboring 3D image patches and is evaluated against a state-of-the-art CNN across three datasets.
- I. INTRODUCTION: Medical image segmentation supports clinical and research uses including organ quantification, disease measurement, treatment planning, and computer-aided intervention.Automatic methods are presented as faster, cheaper, and more reproducible than expert manual segmentation.
- I. INTRODUCTION: Deep learning has largely displaced classical segmentation methods, but contemporary medical segmentation architectures still use convolution as their main building block.Attention and recurrent approaches discussed in prior work also continued to rely on convolution operations.
- I. INTRODUCTION: Convolutional networks benefit from local sparse connections, parameter sharing, and translation equivariance, which provide useful inductive biases for vision.These properties help explain why convolution has remained central across segmentation, classification, and detection.
- I. INTRODUCTION: The proposed architecture uses self-attention between linear embeddings of neighboring 3D image patches without convolution operations.The work explicitly frames this as a convolution-free network for 3D medical image segmentation.
- I. INTRODUCTION: The paper evaluates the convolution-free network on three medical segmentation datasets and compares it with a state-of-the-art CNN.Its contributions include studying whether self-attention can replace convolution in this setting.
A. Proposed network
The proposed network segments the center patch of a 3D image block using self-attention over embedded neighboring patches, without convolution operations.
- A. Proposed network: The network partitions a 3D block into n^3 non-overlapping patches and uses all patches to predict the center patch's segmentation.Experiments use n=3 or 5, corresponding to 27 or 125 patches.
- A. Proposed network: Each flattened patch is mapped into a D-dimensional embedding, combined with learned positional encoding, and supplied to the encoder.The positional encoding preserves patch ordering for the transformer.
- A. Proposed network: Each encoder stage applies multi-head self-attention followed by a two-layer feed-forward network, with residual connections, ReLU activations, and layer normalization.Separate attention heads compute query, key, and value sequences before their outputs are stacked and reprojected.
- A. Proposed network: The final encoder output is projected to nclass values per patch position and reshaped into the predicted segmentation for the center patch.For binary segmentation, nclass=2.
B. Implementation and training
The model is evaluated against 3D UNet++ using Dice-based training, and pre-training on unlabeled images is proposed to improve performance when labeled data are scarce.
- B. Implementation and training: The proposed model is compared with 3D UNet++, a state-of-the-art CNN for medical image segmentation.The implementation uses TensorFlow 1.16 and runs on an NVIDIA GeForce GTX 1080 GPU.
- B. Implementation and training: Training maximizes the Dice similarity coefficient between the center-patch prediction and ground-truth segmentation using Adam.The batch size is 10 and the learning rate starts at 10^-4, halving when validation loss does not decrease after an epoch.
- B. Implementation and training: Pre-training uses denoising or inpainting on unlabeled images, then fine-tunes the entire network for segmentation with a new softmax output layer.Denoising adds Gaussian noise at SNR=10 dB, while inpainting replaces the center patch with zeros; both target the noise-free center patch.
C. Data
The study uses randomly split medical-image datasets, reserving about one-fifth of training images for validation and evaluating final models on test images.
- C. Data: About one-fifth of the training images was used for validation, with final models evaluated on test images.Images were randomly split into training, validation, and test sets.
- C. Data: The experiments used multiple medical-image datasets, as summarized in Table I.
III. RESULTS AND DISCUSSION
The convolution-free network achieved results at least comparable to, and in some cases better than, UNet++ across three medical-image segmentation datasets. Pre-training improved performance with fewer labeled images, while attention analyses and ablations highlighted multiscale and multi-head design effects.
- III. RESULTS AND DISCUSSION: The proposed method was significantly better than, or at least on par with, UNet++ across DSC, HD95, and ASSD on three datasets.Table II reports test segmentation accuracy and paired t-tests for each dataset and criterion.
- III. RESULTS AND DISCUSSION: With pre-training, the convolution-free network achieved significantly more accurate segmentations using fewer labeled training images.The reduced-label experiments used 5, 10, and 15 labeled images; cortical-plate pre-training used 500 unlabeled dHCP images, while pancreas pre-training used remaining training images.
- III. RESULTS AND DISCUSSION: Early network stages attended broadly to surrounding anatomical structures, whereas deeper stages focused more closely on the target organ or region.This pattern was observed for pancreas and cortical-plate segmentation.
- III. RESULTS AND DISCUSSION: Different multi-head attention heads learned distinct attention patterns, providing flexibility in modeling relations among image patches.The attention-map analysis links this variability to the multi-head design.
- III. RESULTS AND DISCUSSION: Larger patches or greater depth typically produced slight accuracy improvements, while fixed or absent positional encoding slightly reduced accuracy.The ablation used the pancreas dataset and compared these choices with the baseline configuration.
- III. RESULTS AND DISCUSSION: Single-head attention significantly reduced segmentation accuracy, supporting the importance of multi-head attention for learning more complex patch relations.
IV. CONCLUSIONS
The paper presents self-attention between neighboring 3D patches as a convolution-free alternative for medical-image segmentation. It reports better performance than a state-of-the-art CNN across three datasets and with only 5–15 labeled images after pre-training.
- IV. CONCLUSIONS: The proposed model uses self-attention between neighboring 3D patches instead of convolutions and outperformed a state-of-the-art CNN on three medical-image segmentation datasets.
- IV. CONCLUSIONS: With denoising and in-painting pre-training on unlabeled images, the network outperformed a CNN when only 5–15 labeled training images were available.The authors also expect the model to be effective for anomaly detection and classification in medical image analysis.