A Data-scalable Transformer for Medical Image Segmentation: Architecture, Model Efficiency, and Benchmark

TL;DR

MedFormer presents a data-scalable hybrid Transformer for 3D medical image segmentation that learns from limited data without pre-training. It introduces convolutional inductive bias, an efficient Bidirectional Multi-Head Attention (B-MHA) with linear complexity, and a global multi-scale semantic map fusion to capture global context with minimal overhead. Across a large cardiac MRI collection and seven public datasets, MedFormer outperforms CNNs and Transformer baselines, showing strong data scalability, robustness to domain shifts, and superior generalization. The work provides a public codebase to enable fair, cross-task benchmarking and sets a foundation for robust, clinically applicable medical segmentation models.

Abstract

Transformers have demonstrated remarkable performance in natural language processing and computer vision. However, existing vision Transformers struggle to learn from limited medical data and are unable to generalize on diverse medical image tasks. To tackle these challenges, we present MedFormer, a data-scalable Transformer designed for generalizable 3D medical image segmentation. Our approach incorporates three key elements: a desirable inductive bias, hierarchical modeling with linear-complexity attention, and multi-scale feature fusion that integrates spatial and semantic information globally. MedFormer can learn across tiny- to large-scale data without pre-training. Comprehensive experiments demonstrate MedFormer's potential as a versatile segmentation backbone, outperforming CNNs and vision Transformers on seven public datasets covering multiple modalities (e.g., CT and MRI) and various medical targets (e.g., healthy organs, diseased tissues, and tumors). We provide public access to our models and evaluation pipeline, offering solid baselines and unbiased comparisons to advance a wide range of downstream clinical applications.

Figures (6)

• Figure 1: MedFormer exhibits outstanding generalizability across seven diverse public datasets, encompassing various modalities, target structures, and target sizes. While other Transformer-based models display good performance on certain specific datasets, they fail to achieve the same level of generalizability of nnUNet or MedFormer.
• Figure 2: (A): The architecture of MedFormer. (B): The proposed efficient bidirectional multi-head attention (B-MHA) that reduces the complexity of conventional self-attention to be linear. (C): The proposed multi-scale semantic map fusion module. (D): Example images of the collected large cardiac MRI dataset. (E): Example images of the testing data of M&Ms dataset that comes from four vendors, where domain shift in terms of visual appearance exists among vendors. The testing data samples are prepared for the evaluation of both model performance and robustness. (F): Example images of seven widely-used datasets that cover diverse tasks, including multiple modalities (e.g., CT and MRI) and various medical targets (e.g., healthy organs, diseased tissues, and tumors).
• Figure 3: Illustration of the initial semantic map generation.
• Figure 4: (A): The curve of Dice score coefficient (DSC) v.s. training data percentage in the collect cardiac dataset. (B) and (C): DSC v.s. number of parameters v.s. Flops. y-axis: DSC, x-axis: number of parameters/M, bubble size/number under model: the Flops/G. (B) shows 2D models comparison under 5% percent data, measured with $256\times 256$ input size. (C) shows 3D models comparison on BCV dataset, measured with $64\times128\times128$ input size. (D): The robustness analysis of models across data scales on four testing vendor domains. A detailed setting description can be found in the method section. (E): Segmentation boundary of models trained with 5% data on four domains. White: ground truth, Red: LV. Green: MYO. Blue: RV. The numbers on the lower right of the images are the Dice scores.
• Figure 5: Segmentation visualization. (A): ACDC dataset. (B): BCV dataset, (C): LiTS dataset. Yellow arrows indicate notable and challenging areas of segmentation errors from different models.