Source-linked AI summary
Medical Transformer: Gated Axial-Attention for Medical Image Segmentation
Jeya Maria Jose Valanarasu, Poojan Oza, Ilker Hacihaliloglu, Vishal M. Patel
TL;DR
Medical segmentation requires long-range image dependencies, but ConvNets emphasize local patterns and transformers are difficult to train on scarce medical datasets. The paper introduces gated axial attention and LoGo training in MedT, which combines global and local processing. Across three datasets, MedT achieves better performance than convolutional and related transformer-based architectures.
Problem
ConvNets lack long-range dependency modeling, while transformer-based models are difficult to train effectively on relatively small medical imaging datasets.
Method
MedT combines gated position-sensitive axial attention, which controls positional-encoding influence, with LoGo training on whole images and patches.
Results
MedT improves performance over convolutional and related transformer-based architectures across three medical image segmentation datasets.
Takeaways & Limitations
The proposed architecture provides a transformer-based medical segmentation approach that does not require pre-training on large-scale datasets.
Abstract
from arXiv · showhide
Over the past decade, Deep Convolutional Neural Networks have been widely adopted for medical image segmentation and shown to achieve adequate performance. However, due to the inherent inductive biases present in the convolutional architectures, they lack understanding of long-range dependencies in the image. Recently proposed Transformer-based architectures that leverage self-attention mechanism encode long-range dependencies and learn representations that are highly expressive. This motivates us to explore Transformer-based solutions and study the feasibility of using Transformer-based network architectures for medical image segmentation tasks. Majority of existing Transformer-based network architectures proposed for vision applications require large-scale datasets to train properly. However, compared to the datasets for vision applications, for medical imaging the number of data samples is relatively low, making it difficult to efficiently train transformers for medical applications. To this end, we propose a Gated Axial-Attention model which extends the existing architectures by introducing an additional control mechanism in the self-attention module. Furthermore, to train the model effectively on medical images, we propose a Local-Global training strategy (LoGo) which further improves the performance. Specifically, we operate on the whole image and patches to learn global and local features, respectively. The proposed Medical Transformer (MedT) is evaluated on three different medical image segmentation datasets and it is shown that it achieves better performance than the convolutional and other related transformer-based architectures. Code: https://github.com/jeya-maria-jose/Medical-Transformer
1 Introduction
Medical image segmentation needs accurate, robust methods because segmentation supports diagnosis, surgical planning, and treatment. ConvNets capture local patterns but struggle with long-range dependencies, while transformers face training challenges on scarce medical datasets; MedT addresses these issues with gated attention and LoGo training.
- Medical image segmentation supports computer-aided diagnosis, image-guided surgery, surgical planning, and treatment strategies.
- ConvNets attend mainly to local pixel subsets, limiting their ability to model long-range dependencies and global image context.
- Modeling long-range dependencies can reduce false positives by distinguishing scattered background pixels from segmentation-mask pixels.
- Transformer self-attention encodes long-range dependencies, but transformer-based models generally work best with large-scale datasets, unlike scarce medical imaging datasets.
- The paper proposes gated position-sensitive axial attention, LoGo training, and MedT for medical image segmentation on smaller datasets.
2 Medical Transformer (MedT)
MedT combines self-attention, axial factorization, learnable gating, and Local-Global training. The design uses global image context and local patches while controlling positional-encoding influence in small-data settings.
- 2.1 Self-Attention Overview: Self-attention projects an input feature map into queries, keys, and values, then pools values using global affinities to capture non-local information.
- Axial-Attention: Axial attention decomposes two-dimensional self-attention into height-wise and width-wise modules, improving computational efficiency while modeling long-range interactions.
- 2.2 Gated Axial-Attention: Gated axial attention adds learnable gates that control how strongly learned relative positional encodings influence non-local context.
- 2.3 Local-Global Training: LoGo combines a global branch operating at original resolution with a local branch processing 16 image patches, then adds their feature maps for segmentation.
- 2.3 Local-Global Training: The global branch models long-range dependencies, whereas the local branch focuses on finer features through patch-based processing.
3 Experiments and Results
MedT is evaluated against convolutional and transformer-based baselines on three medical image segmentation datasets using F1 and IoU scores. It outperforms the compared methods, with qualitative results highlighting improved handling of long-range dependencies.
- The evaluation uses Brain anatomy ultrasound, Gland segmentation, and MoNuSeg datasets, comparing convolutional and transformer-based baselines.
- F1 and IoU scores provide the quantitative measures for comparing the proposed methods with convolutional and transformer-based baselines.
- MedT performs better than gated axial attention, LoGo, and previous methods across the evaluated datasets.
- MedT improves over fully attention baselines by 0.92%, 4.76%, and 2.72% on Brain US, GlaS, and MoNuSeg, respectively, in F1 score.
- MedT improves over the best convolutional baseline by 1.32%, 2.19%, and 0.06% on Brain US, GlaS, and MoNuSeg, respectively, in F1 score.
- Qualitative examples show MedT capturing small masks and avoiding false predictions near segmentation masks through pixel-wise dependency modeling.
4 Conclusion
The paper proposes gated axial attention and LoGo training for MedT, a transformer-based medical image segmentation model. MedT combines global and local feature learning and achieves good performance without pre-training on large-scale datasets.
- MedT uses gated axial attention as its main encoder building block for transformer-based medical image segmentation.
- LoGo trains MedT at full image resolution and on patches to learn global context and finer local features, respectively.
- Unlike other transformer-based models, MedT does not require pre-training on large-scale datasets.
- Experiments on three datasets report good MedT performance over ConvNets and related transformer-based architectures.
Supplementary Material for Medical Transformer: Gated Axial-Attention for Medical
The supplementary material expands the paper’s documentation of datasets, architecture, training strategy, ablations, parameter counts, and additional results.
- The supplementary material provides more detailed descriptions of the datasets used in the paper.
- It gives more intricate details about the proposed architecture and training strategy.
- It includes an ablation study, parameter-count analysis, and additional experimental results.
1 Dataset details
The study evaluates medical image segmentation using Brain US, GLAS, and MoNuSeg datasets with different image sources, sample sizes, annotations, and resized resolutions.
- Brain US contains 1,629 annotated images from 20 premature neonates, split into 1,300 training and 329 testing images at 128 × 128 resolution.
- GLAS contains 165 H&E-stained microscopic images, split into 85 training and 80 testing images, resized to 128 × 128.
- MoNuSeg contains H&E-stained tissue images across multiple organs and patients, with 30 training and 14 testing images resized to 512 × 512.
2 MedT details
MedT combines gated axial attention with global and local processing branches. The global branch processes whole-image features, while the local branch operates on image patches.
- MedT uses gated axial attention as its basic building block and applies the LoGo strategy during training.The architecture contains global and local branches fed by feature maps from an initial convolutional block.
- The global and local branches use transformer layers in their encoders and convolutional blocks in their decoders.
- The global branch contains two encoder and two decoder blocks, whereas the local branch contains five of each.
3 Training details
The experiments use fixed optimization and hardware settings, with gate training delayed during the initial epochs.
- MedT is trained for 400 epochs with batch size 4, Adam optimization, and a 0.001 learning rate.
- Gate training remains inactive for the first 10 epochs of gated axial attention training.
- Experiments run on an Nvidia Quadro 8000 GPU.
4 Analysis
The analysis evaluates ablations, parameter counts, and whether MedT’s performance advantage is attributable to its parameterization. MedT reduces parameters through shallow global processing and patch-based local processing.
- 4.1 Ablation Study: The ablation study starts with U-Net, adds residual connections for Res-UNet, and replaces encoder convolutions with axial attention layers.The resulting Axial Attention UNet includes an additional front-end convolutional block for feature extraction.
- 4.2 Number of Parameters: Table 2 compares MedT’s parameter count with existing methods, including modified U-Net and Res-UNet configurations with adjusted filter counts.
- 4.2 Number of Parameters: MedT reduces parameters by using two global encoder-decoder layers and restricting the local branch to image patches.The gated axial attention block adds only four learnable parameters to the layer.
- 4.2 Number of Parameters: The parameter-matching analysis reports that baselines with more parameters do not exceed MedT in performance.This comparison is intended to assess whether MedT’s improvement is due to a slight parameter-count change.
5 Results
Qualitative results compare MedT with U-Net and Res-UNet on ultrasound scans. The highlighted regions identify errors made by convolutional baselines that MedT avoids.
- In highlighted regions of ultrasound scans, U-Net and Res-UNet misclassify areas as brain ventricles, whereas MedT does not make the same mistake.
6 Concurrent works
Concurrent work includes TransUNet, which combines a transformer patch encoder with a convolutional decoder, and TransFuse, which fuses parallel CNN and transformer branches. In contrast, this work investigates a self-attention-only transformer encoder without pre-training.
- TransUNet uses a transformer-based encoder over image-patch sequences and a convolutional decoder with skip connections for medical image segmentation.
- TransFuse combines parallel CNN and transformer branches through a BiFusion module for polyp segmentation.
- Unlike these approaches, the paper explores a self-attention-only transformer encoder for medical image segmentation without pre-trained weights.