Source-linked AI summary
BiTr-Unet: a CNN-Transformer Combined Network for MRI Brain Tumor Segmentation
Qiran Jia, Hai Shu
TL;DR
Automated segmentation of 3D medical images is important for diagnosis and treatment planning, while annotated-data shortages persist. BiTr-Unet refines a CNN-Transformer approach for BraTS2021, and its validation results are described as indicating a stable network for extracting local and long-range MRI features.
Problem
Automated biomedical image segmentation supports disease diagnosis and treatment planning, but annotated-data shortages persist.
Source: §1 · pp. 1–2Method
BiTr-Unet refines TransBTS by combining CNN and Transformer components and using two sets of ViT layers.
Source: §1 · p. 4, Modifications from TransBTS · p. 8Results
Validation results indicate BiTr-Unet is a stable and powerful network for extracting local and long-range dependencies in 3D MRI scans.
Source: §4 · p. 10Takeaways & Limitations
The authors describe the Transformer module as a potential performance advantage over U-Net-like models without attention modules.
Source: §4 · p. 10The authors did not complete experiments on deeper encoder-decoder skip connections through ViT blocks, different embedding dimensions, or complexity-reducing Transformer algorithms.
Source: §4 · p. 11
Abstract
from arXiv · showhide
Convolutional neural networks (CNNs) have achieved remarkable success in automatically segmenting organs or lesions on 3D medical images. Recently, vision transformer networks have exhibited exceptional performance in 2D image classification tasks. Compared with CNNs, transformer networks have an appealing advantage of extracting long-range features due to their self-attention algorithm. Therefore, we propose a CNN-Transformer combined model, called BiTr-Unet, with specific modifications for brain tumor segmentation on multi-modal MRI scans. Our BiTr-Unet achieves good performance on the BraTS2021 validation dataset with median Dice score 0.9335, 0.9304 and 0.8899, and median Hausdorff distance 2.8284, 2.2361 and 1.4142 for the whole tumor, tumor core, and enhancing tumor, respectively. On the BraTS2021 testing dataset, the corresponding results are 0.9257, 0.9350 and 0.8874 for Dice score, and 3, 2.2361 and 1.4142 for Hausdorff distance. The code is publicly available at https://github.com/JustaTinyDot/BiTr-Unet.
1 Introduction
BiTr-Unet adapts a CNN-Transformer design for BraTS2021 brain-tumor segmentation, building on TransBTS while adding a second set of ViT layers. The approach responds to the limited research on Transformer-based 3D biomedical image segmentation.
1 Introduction
- Biomedical image segmentation supports disease diagnosis and treatment planning, while annotated-data shortages persist despite deep learning’s speed and accuracy advantages over traditional methods.The introduction describes deep learning as using multiple processing layers to extract data representations.Source: §1 · pp. 1–2
- BraTS2021 increased its curated multi-institutional glioma MRI scans from 660 to 2,000 and provides four 3D modalities with physician-labeled training segmentations.The modalities are native T1, post-contrast T1, T2, and T2-FLAIR.Source: §1 · p. 2
- BraTS2021 withholds reference segmentations for validation and test data and ranks whole tumor, tumor core, and enhancing tumor using four metrics.The metrics are Dice score, Hausdorff distance (95%), sensitivity, and specificity.Source: §1 · p. 2
- U-Net’s encoder-decoder with skip connections became a frequently modified BraTS baseline, and nnU-Net placed first in BraTS2020.Other leading BraTS2020 networks added modules such as variational autoencoders, attention gates, or CBAM to U-Net-like structures.Source: §1 · p. 3
- Vision Transformers use self-attention over image patches to capture long-range dependencies, addressing the locality limitation of convolution-only networks.The introduction describes this as a contrast between Transformer and convolution-based approaches.Source: §1 · p. 3
- Transformer-based 3D biomedical image segmentation research remains limited, although UNETR and TransBTS demonstrate possible approaches.TransBTS places ViT layers between the CNN encoder and decoder to compute long-range dependencies with a global receptive field.Source: §1 · p. 3
- BiTr-Unet refines TransBTS for BraTS2021 and adds a second set of ViT layers where TransBTS uses one.The authors retain its CNN-Transformer combination, feature embedding, and feature mapping while modifying the network for the challenge.Source: §1 · p. 4
2 Method
BiTr-Unet combines 3D CNN components with Transformer layers to segment multimodal MRI scans. Its design adds attention and deeper Transformer use to the TransBTS backbone.
2.1 Basic Structure of BiTr-Unet
- T1, T1c, T2, and T2-FLAIR scans are stacked as input X, which BiTr-Unet maps to a predicted segmentation image Y.Training labels are stored separately from the MRI input.Source: §2.1 · p. 4, Preprocessing · p. 5
- The 3D CNN encoder downsamples feature maps and uses 3D CBAM attention blocks in place of regular convolutional blocks for adaptive feature refinement.The encoder has four stride-2 convolutional blocks after an initial stride-1 block.Source: 3D CNN Encoder with Attention Module · p. 5
- Transformer outputs are reshaped and convolved into 3D features for upsampling, while encoder outputs provide skip connections to corresponding decoder layers.The decoder uses four stride-2 convolutional blocks and a final stride-1 block to produce the segmentation image.Source: Feature Mapping · pp. 6–7, Skip Connections · p. 7, 3D CNN Decoder · p. 7
- Small predicted segmentation volumes are removed, and differently trained BiTr-Unet models are combined by voxel-wise majority voting.When vote counts tie, the category with the highest averaged prediction probability is selected.Source: Postprocessing and Model Ensemble · p. 7
2.2 Comparison with Related Work
- Projected 3D feature maps become position-embedded tokens for MHA and FFN Transformer layers, used at both the fourth and fifth encoder levels.The network increases its depth and adds attention modules to refine the CNN encoder and decoder.Source: 3D CNN Encoder with Attention Module · p. 5, Feature Embedding of Feature… · p. 6, Transformer Layers · p. 6
- Unlike BiTr-Unet and TransBTS, UNETR has no convolutional encoder blocks; the paper notes this may limit local-context extraction and that its unchanged sequence dimension is computationally expensive for high-resolution 3D images.The paper characterizes UNETR as using a pure Transformer for encoding 3D medical images.Source: Differences from UNETR · p. 8
3 Result
BiTr-Unet was evaluated on BraTS2021 validation and testing data, with an ensemble used for validation and the 7050-epoch model alone submitted for testing. Training used extensive cropping and intensity augmentation.
3.1 Implementation Details
- Training used four RTX8000 GPUs for 7050 epochs with batch size 16, random 128×128×128 crops, intensity shifts and scaling, and seven-flip validation TTA.The learning rate began at 0.0002 and decayed by each iteration with power 0.9.Source: §3.1 · pp. 8–9
3.2 Segmentation Result
- Validation evaluated three checkpoints and their ensemble; Table 1 reports mean scores, while Fig. 4 visualizes three ensembled validation cases.The validation data were evaluated through the Synapse platform.Source: §3.2 · p. 9, Table 1 · p. 9, Figure 4 · p. 9
- The 7050-epoch BiTr-Unet was submitted without ensembling for testing to avoid Docker image size, runtime, and stability risks; Table 3 reports testing results.The reported testing results are for BraTS2021 testing data.Source: §3.2 · p. 9, Table 3 · p. 10
4 Discussion
The discussion characterizes validation performance as supporting the model’s extraction of local and long-range MRI dependencies, while noting weaker testing statistics on more dissimilar data. It also identifies untested architectural variants and computation-reduction methods for future exploration.
4 Discussion
- Testing shared validation-like medians and 75th percentiles but had inferior mean, standard deviation, and 25th percentile; the authors say greater data dissimilarity may expose difficulty with unseen patterns.The authors link this discrepancy to testing data being more dissimilar from training data than validation data.Source: §4 · p. 10
- Validation results are described as indicating that BiTr-Unet extracts local and long-range dependencies, with potentially better performance than U-Net-like models without attention modules.The authors describe the Transformer-related performance advantage as potential.Source: §4 · p. 10
- Time limits left deeper encoder-decoder variants with ViT skip connections and different embedding dimensions untested; the authors also propose exploring methods to reduce computation complexity.The proposed complexity-reduction methods are motivated by recent Transformer work for image classification or segmentation.Source: §4 · pp. 10–11