Medical Image Segmentation using Squeeze-and-Expansion Transformers

1 Introduction

Automated Medical image segmentation, i.e., automated delineation of anatomical structures and other regions of interest (ROIs), is an important step in computer-aided diagnosis; for example it is used to quantify tissue volumes, extract key quantitative measurements, and localize pathology [Schlemper et al., 2019; Orlando et al., 2020]. Good segmentation demands the model to see the big picture and fine details at the same time, i.e., learn image features that incorporate large context while keep high spatial resolutions to output finegrained segmentation masks. However, these two demands pose a dilemma for CNNs, as CNNs often incorporate larger context at the cost of reduced feature resolution. A good meaeffective receptive field

∗Corresponding Author.

(effective RF) [Luo et al., 2016], i.e., the input areas which have non-negligible impacts to the model output.

Since the advent of U-Net [Ronneberger et al., 2015], it has shown excellent performance across medical image segmentation tasks. A U-Net consists of an encoder and a decoder, in which the encoder progressively downsamples the features and generates coarse contextual features that focus on contextual patterns, and the decoder progressively upsamples the contextual features and fuses them with fine-grained local visual features. The integration of multiple scale features enlarges the RF of U-Net, accounting for its good performance. However, as the convolutional layers deepen, the impact from far-away pixels decay quickly. As a results, the effective RF of a U-Net is much smaller than its theoretical RF. As shown in Fig.2, the effective RFs of a standard U-Net and DeepLabV3+ are merely around 90 pixels. This implies that they make decisions mainly based on individual small patches, and have difficulty to model larger context. However, in many tasks, the heights/widths of the ROIs are greater than 200 pixels, far beyond their effective RFs. Without "seeing the bigger picture", U-Net and other models may be misled by local visual cues and make segmentation errors.

Many improvements of U-Net have been proposed. A few typical variants include: U-Net++ [Zhou et al., 2018] and UNet 3+ [Huang et al., 2020], in which more complicated skip connections are added to better utilize multi-scale contextual features; attention U-Net [Schlemper et al., 2019], which employs attention gates to focus on foreground regions; 3D UNet [Çiçek et al., 2016] and V-Net [Milletari et al., 2016], which extend U-Net to 3D images, such as MRI volumes; Eff-UNet [Baheti et al., 2020], which replaces the encoder of U-Net with a pretrained EfficientNet [Tan and Le, 2019].

Transformers [Vaswani et al., 2017] are increasingly popular in computer vision tasks. A transformer calculates the pairwise interactions ("self-attention") between all input units, combines their features and generates contextualized features. The contextualization brought by a transformer is analogous to the upsampling path in a U-Net, except that

Input Encoder

Transformer

it has unlimited effective receptive field, good at capturing long-range correlations. Thus, it is natural to adopt transformers for image segmentation. In this work, we present Segtran, an alternative segmentation architecture based on transformers. A straightforward incorporation of transformers into segmentation only yields moderate performance. As transformers were originally designed for Natural Language Processing (NLP) tasks, there are several aspects that could be improved to better suit image applications. To this end, we propose a novel transformer design Squeeze-and-Expansion Transformer, in which a squeezed attention block helps regularize the huge attention matrix, and an expansion block learns diversified representations. In addition, we propose a learnable sinusoidal positional encoding that imposes a continuity inductive bias for the transformer. Experiments demonstrate that they lead to improved segmentation performance.

We evaluated Segtran on two 2D medical image segmentation tasks: optic disc/cup segmentation in fundus images of the REFUGE'20 challenge, and polyp segmentation in colonoscopy images. Additionally, we also evaluated it on a 3D image segmentation task: brain tumor segmentation in MRI scans of the BraTS'19 challenge. Segtran has consistently shown better performance than U-Net and its variants (UNet++, UNet3+, PraNet, and nnU-Net), as well as DeepLabV3+ [Chen et al., 2018].

2 RelatedWork

Our work is largely inspired by DETR [Carion et al., 2020]. DETR uses transformer layers to generate contextualized features that represent objects, and learns a set of object queries to extract the positions and classes of objects in an image. Although DETR is also explored to do panoptic segmentation [Kirillov et al., 2019], it adopts a two-stage approach which is not applicable to medical image segmentation. A followup work of DETR, Cell-DETR [Prangemeier et al., 2020] also employs transformer for biomedical image segmentation, but its architecture is just a simplified DETR, lacking novel components like our Squeeze-and-Expansion transformer. Most recently, SETR [Zheng et al., 2021] and TransU-Net [Chen et al., 2021] were released concurrently or after our paper submission. Both of them employ a Vision Transformer (ViT) [Dosovitskiy et al., 2021] as the encoder to extract image features, which already contain global contextual information. A few convolutional layers are used as the decoder to generate the segmentation mask. In contrast, in Segtran, the transformer layers build global context on top of the local image features extracted from a CNN backbone, and a Feature Pyramid Network (FPN) generates the segmentation mask.

[Murase et al., 2020] extends CNNs with positional encoding channels, and evaluates them on segmentation tasks. Mixed results were observed. In contrast, we verified through an ablation study that positional encodings indeed help Segtran to do segmentation to a certain degree.

Receptive fields of U-Nets may be enlarged by adding more downsampling layers. However, this increases the number of parameters and adds the risk of overfitting. Another way of increasing receptive fields is using larger stride sizes of the convolutions in the downsampling path. Doing so, however, sacrifices spatial precision of feature maps, which is often disadvantageous for segmentation [Liu and Guo, 2020].

3 Squeeze-and-Expansion Transformer

The core concept in a transformer is Self Attention, which can be understood as computing an affinity matrix between different units, and using it to aggregate features:

where K, Q, V are key, query, and value projections, respectively. f is softmax after dot product. Att_weight(X,X) is the pairwise attention matrix between input units, whose i, jth element defines how much the features of unit j contributes to the fused (contextualized) features of unit i. FFN is a feedforward network to further transform the fused features.

The basic transformer above is extended to a multi-head attention (MHA) [Vaswani et al., 2017; Voita et al., 2019], aiming to capture different types of associations between input units. Each of the Nh heads computes individual attention wights and output features (C/Nh-dimensional), and their output features are concatenated along the channel dimension into C-dimensional features. Different heads operate in exclusively different feature subspaces.

We argue that transformers can be improved in four aspects make them better suited for images:

1. In Eq.(2), the intermediate features Attention(X) are obtained by linearly combining the projected input features, where the attention matrix specifies the combination coefficients. As the attention matrix is huge: N×N, with typically N > 1000, it is inherently vulnerable to noises and overfitting. Reducing the attention matrix to lower rank matrices may help.
2. In traditional transformers, the output features are monomorphic: it has only one set of feature transformations (the multi-head transformer also has one set of transformations after concatenation), which may not have enough capacity to fully model data variations. Just like a mixture of Gaussians almost always better depicts a data distribution than a single Gaussian, data variations can be better captured using a mixture of k transformers.
3. In traditional transformers, the key and query projections are independently learned, enabling them to capture asymmetric relationships between tokens in natural language. However, the relationships between image units are often symmetric, such as whether two pixels belong to the same segmentation class.
4. Pixels in an image have strong locality and semantic continuity. The two mainstream positional encoding

schemes [Carion et al., 2020; Dosovitskiy et al., 2021] do not fully impose such an inductive bias. This bias could be imposed by an improved positional encoding.

The Squeeze-and-Expansion Transformer aims to improve in all the four aspects. The Squeezed Attention Block computes attention between the input and M inducing points [Lee et al., 2019], and compresses the attention matrices to N×M. The Expanded Attention Block is a mixture-of-experts model with Nm modes ("experts"). In both blocks, the query projections and key projections are tied to make the attention symmetric, for better modeling of the symmetric relationships between image units. In addition, a Learnable Sinusoidal Positional Encoding helps the model capture spatial relationships.

3.1 Squeezed Attention Block

[Lee et al., 2019] proposes Induced Set Attention Block (ISAB) by bringing inducing points into the transformer. It was originally designed to learn good features of a set of unordered objects. Here we employ this design to "squeeze" the bloated attention matrix, so as to reduce noises and overfitting in image tasks. We rename ISAB as Squeezed Attention Block (SAB) to highlight its new role in this context 1 .

In SAB, inducing points are a set of M learned embeddings c1, · · · , cM in an external discrete codebook. Usually M ≪N, the number of input units. The inducing points are first transformed into new embeddings C′ = c′1 , · · · , c′M after attending with the input. The combination of these embeddings form the output features Xout = x′1 , · · · , x′N (Fig.3):

where Single-Head(·, ·) is a single-head transformer, and EAB(·, ·) is an Expanded Attention Block presented in the next subsection. In each of the two steps, the attention matrix is of N ×M, much more compact than vanilla transformers.

SAB is conceptually similar to the codebook used for discrete representation learning in [Esser et al., 2020], but the discretized features are further processed by a transformer. SAB can trace its lineage back to low-rank matrix factorization, i.e., approximating a data matrix Xn×n ≈ Pn×d · Qd×n, which is a traditional regularization technique against data noises and overfitting. Confirmed by an ablation study, SAB helps fight against noises and overfitting as well.

3.2 Expanded Attention Block

The Expanded Attention Block (EAB) consists of Nm modes, each an individual single-head transformer. They output Nm sets of contextualized features, which are then aggregated into one set using dynamic mode attention:

where the mode attention G ∈ RNu×Nm is obtained by doing a linear transformation of each mode features, and taking softmax over all the modes. Eq.(9) takes a weighted sum over the modes to get the final output features Xout. This dynamic attention is inspired by the Split Attention of the ResNest model [Zhang et al., 2020b].

EAB is a type of Mixture-of-Experts [Shazeer et al., 2017], an effective way to increase model capacity. Although there is resemblance between multi-head attention (MHA) and EAB, they are essentially different, as shown in Fig.4. In MHA, each head resides in an exclusive feature subspace and provides unique features. In contrast, different modes in EAB share the same feature space, and the representation power largely remains after removing any single mode. The modes join forces to offer more capacity to model diverse data, as shown in an ablation study. In addition, EAB is also different from the Mixture of Softmaxes (MoS) transformer [Zhang et al., 2020a]. Although MoS transformer also uses k sets of queries and keys, it shares one set of value transformation.

1 We clarify that our contribution is a novel transformer architecture that combines SAB with an Expanded Attention Block.

3.3 Learnable Sinusoidal Positional Encoding

A crucial inductive bias for images is the pixel locality and semantic continuity, which is naturally encoded by convolutional kernels. As the input to transformers is flattened into 1-D sequences, positional encoding (PE) is the only source to inject information about spatial relationships. On the one hand, this makes transformers flexible to model arbitrary shapes of input. On the other hand, the continuity bias of images is non-trivial to fully incorporate. This is a limitation of the two mainstream PE schemes: Fixed Sinusoidal Encoding and Discretely Learned Encoding [Carion et al., 2020; Dosovitskiy et al., 2021]. The former is spatially continuous but lacks adaptability, as the code is predefined. The latter learns a discrete code for each coordinate without enforcing spatial continuity.

We propose Learnable Sinusoidal Positional Encoding, aiming to bring in the continuity bias with adaptability. Given a pixel coordinate (x, y), our positional encoding vector pos(x, y) is a combination of sine and cosine functions of linear transformations of (x, y):

where i indexes the elements in pos, {ai, bi, ci} are learnable weights of a linear layer, and C is the dimensionality of image features. To make the PE behave consistently across different image sizes, we normalize (x, y) into [0, 1]2 . When the input image is 3D, Eq.(10) is trivially extended by using 3D coordinates (x, y, z).

The encoding in Eq.(10) changes smoothly with pixel coordinates, and thus nearby units receive similar positional encodings, pushing the attention weights between them towards larger values, which is the spirit of the continuity bias. The learnable weights and sinusoidal activation functions make the code both adaptable and nonlinear to model complex spatial relationships [Tancik et al., 2020].

4 Segtran Architecture

As a context-dependent pixel-wise classification task, segmentation faces a conflict between larger context (lower resolution) and localization accuracy (higher resolution). Segtran partly circumvents this conflict by doing pairwise feature contextualization, without sacrificing spatial resolutions. There are five main components in Segtran (Fig.1): 1) a CNN backbone to extract image features, 2) input/output feature pyramids to do upsampling, 3) learnable sinusoidal positional encoding, 4) Squeeze-and-Expansion transformer layers to contextualize features, and 5) a segmentation head.

4.1 CNN Backbone

We employ a pretrained CNN backbone to extract features maps with rich semantics. Suppose the input image is X0 ∈ RH0×W0×D0 , where for a 2D image, D0 = 1 or 3 is the number of color channels. For a 3D image, D0 ≫ 3 is the number of slices in the depth dimension. For 2D and 3D images, the extracted features are CNN(X0) ∈ RC×H×W , and CNN(X0) ∈ RC×H×W×D , respectively.

On 2D images, typically ResNet-101 or EfficientNet-D4 is used as the backbone. For increased spatial resolution, we change the stride of the first convolution from 2 to 1. Then H,W = H0/16,W0/16. On 3D images, 3D backbones like I3D [Carreira and Zisserman, 2017] could be adopted.

4.2 Transformer Layers

Before being fed into the transformer, the visual features and positional encodings of each unit are added up before being fed to the transformer: Xspatial = Xvisual + pos(coordinates(X)). Xspatial is flattened across spatial dimensions to a 1-D sequence X0 ∈ RNu×C , where Nu is the total number of image units, i.e., points in the feature maps.

The transformer consists of a few stacked transformer layers. Each layer takes input features X, computes the pairwise interactions between input units, and outputs contextualized features Xout of the same number of units. The transformer layers used are our novel design Squeeze-andExpansion Transformer (Section 3).

4.3 Feature Pyramids and Segmentation Head

Although the spatial resolution of features is not reduced after passing through the transformer layers, for richer semantics, the input features to transformers are usually high-level features from the backbone. They are of a low spatial resolution, however. Hence, we increase their spatial resolution with an input Feature Pyramid Network (FPN) [Liu et al., 2018] and an output FPN, which upsample the feature maps at the transformer input end and output end, respectively.

Without loss of generality, let us assume the EfficientNet is the backbone. The stages 3, 4, 6, and 9 of the network are commonly used to extract multi-scale feature maps. Let us denote the corresponding feature maps as f1, f2, f3, f4, respectively. Their shapes are fi ∈ RCi×Hi×Wi , with Hi = H0 2i ,Wi = W0 2i .

As described above, f(X0) = f4 is 1/16 of the original image, which is too coarse for accurate segmentation. Hence, we upsample it with an input FPN, and obtain upsampled feature maps f34:

where conv34 is a 1 × 1 convolution that aligns the channels of f3 to f4, and upsample×2(·) is bilinear interpolation.

f34 is 1/8 of the original image, and is used as the input features to the transformer layers. As the transformer layers keep the spatial resolutions unchanged from input to output feature maps, the output feature maps g34 is also 1/8 of the input image. Still, this spatial resolution is too low for segmentation. Therefore, we adopt an output FPN to upsample the feature maps by a factor of 4 (i.e., 1/2 of the original images). The output FPN consists of two upsampling steps:

where conv12 and conv24 are 1 × 1 convolutional layers that align the channels of f1 to f2, and f2 to f4, respectively.

This FPN scheme is the bottom-up FPN proposed in [Liu et al., 2018]. Empirically, it performs better than the original top-down FPN [Lin et al., 2017], as richer semantics in top layers are better preserved.

The segmentation head is simply a 1 × 1 convolutional layer, outputting confidence scores of each class in the mask.

5 Experiments

Three tasks were evaluated in our experiments:

REFUGE20: Optic Disc/Cup Segmentation in Fundus Images. This task does segmentation of the optic disc and cup in fundus images, which are 2D images of the rear of eyes (Fig. 5). It is a subtask of the REFUGE Challenge 2 [Orlando et al., 2020], MICCAI 2020. 1200 images were provided for training, and 400 for validation. We also used two extra datasets, Drishti-GS dataset [Sivaswamy et al., 2015] and RIM-ONE v3 [Fumero et al., 2011] when training all models. The Disc/Cup dice scores of validation images were obtained from the official evaluation server.

Polyp: Polyp Segmentation in Colonoscopy Images. Polyps are fleshy growths in the colon lining that may become cancerous. This task does polyp segmentation in colonoscopy images (Fig. 5). Two image datasets [Fan et al., 2020] were used: CVC612 (CVC in short; 612 images) and Kvasir (1000 images). Each was randomly split into 80% training and 20% validation, and the training images were merged into one set.

BraTS19: Tumor Segmentation in MRI Images. This task focuses on the segmentation of gliomas, a common brain tumor in MRI scans. It was a subtask of the BraTS'19 challenge 3 [Menze et al., 2015; Bakas et al., 2017], MICCAI 2019. It involves four classes: the whole tumor (WT), the tumor core (TC), the enhancing tumor (ET) and background. Among them, the tumor core consists of the necrotic regions and non-enhancing tumors (red), as well as the enhancing tumor (yellow). 335 scans were provided for training, and 125

2 https://refuge.grand-challenge.org/Home2020/

3 https://www.med.upenn.edu/cbica/brats-2019/

for validation. The dice scores of ET, WT and TC on the validation scans were obtained from the official evaluation server.

5.1 Ablation Studies

Two ablation studies were performed on REFUGE20 to compare: 1) the Squeeze-and-Expansion Transformer versus Multi-Head Transformer; and 2) the Learnable Sinusoidal Positional Encoding versus two schemes as well as not using PE.

All the settings were variants of the standard one, which used three layers of Squeeze-and-Expansion transformer with four modes (Nm = 4), along with learnable sinusoidal positional encoding. Both ResNet-101 and EfficientNet-B4 were evaluated to reduce random effects from choices of the backbone. We only reported the cup dice scores, as the disc segmentation task was relatively easy, with dice scores only varying ±0.005 across most settings.

Type of Transformer Layers. Table 1 shows that Squeezeand-Expansion transformer outperformed the traditional multi-head transformers. Moreover, Both the squeeze attention block and the expansion attention block contributed to improved performance.

Positional Encoding. Table 2 compares learnable sinusoidal positional encoding with the two mainstream PE schemes and no PE. Surprisingly, without PE, performance of Segtran only dropped 1~2%. A possible explanation is that the transformer may manage to extract positional information from the CNN backbone features [Islam et al., 2020].

Number of Transformer Layers. Table 3 shows that as the number of transformer layers increased from 1 to 3, the performance improved gradually. However, one more layer caused performance drop, indicating possible overfitting.

5.2 Comparison with Baselines

Ten methods were evaluated on the 2D segmentation tasks:

• U-Net [Ronneberger et al., 2015]: The implementation in a popular library Segmentation Models.PyTorch (SMP) was used 4 . The pretrained ResNet-101 was chosen as the encoder. In addition, U-Net implemented in U-Net++ (below) was evaluated as training from scratch.
• U-Net++ [Zhou et al., 2018]: A popular PyTorch implementation 5 . It does not provide options to use pretrained encoders, and thus was only trained from scratch.
• U-Net3+ [Huang et al., 2020]: The official PyTorch implementation 6 . It does not provide options to use pretrained encoders.
• PraNet [Fan et al., 2020]: The official PyTorch implementation 7 . The pretrained Res2Net-50 [Gao et al., 2020] was recommended to be used as the encoder.
• DeepLabV3+ [Chen et al., 2018]: A popular PyTorch implementation 8 , with a pretrained ResNet-101 as the encoder.
• Attention based U-Nets [Oktay et al., 2018]: Attention U-Net (AttU-Net) and AttR2U-Net (a combination of AttU-Net and Recurrent Residual U-Net) were evaluated 9 . They learn to focus on important areas by computing element-wise attention weights (as opposed to the pairwise attention of transformers).
• nnU-Net [Isensee et al., 2021]: nnU-Net generates a custom U-Net configuration for each dataset based on its statistics. It is primarily designed for 3D tasks, but can also handle 2D images after converting them to pseudo3D. The original pipeline is time-consuming, and we extracted the generated U-Net configuration and instantiated it in our pipeline to do training and test.
• Deformable U-Net [Jin et al., 2019]: Deformable UNet (DUNet) uses deformable convolution in place of ordinary convolution. The official implementation 10 of DUNet was evaluated.
• SETR [Zheng et al., 2021]: SETR uses ViT as the encoder, and a few convolutional layers as the decoder. The SETR-PUP model in the official implementation 11 was evaluated, by fine-tuning the pretrained ViT

4 https://github.com/qubvel/segmentation_models.pytorch/

5 https://github.com/4uiiurz1/pytorch-nested-unet

6 https://github.com/ZJUGiveLab/UNet-Version

7 https://github.com/DengPingFan/PraNet

8 hhttps://github.com/VainF/DeepLabV3Plus-Pytorch

9 https://github.com/LeeJunHyun/Image_Segmentation

10 https://github.com/RanSuLab/DUNet-retinal-vessel-detection

11 https://github.com/fudan-zvg/SETR/

weights.

• TransU-Net [Chen et al., 2021]: TransU-Net uses a hybrid of ResNet and ViT as the encoder, and a U-Net style decoder. The official implementation 12 was evaluated, by fine-tuning their pretrained weights.
• Segtran: Trained with either a pretrained ResNet-101 or EfficientNet-B4 as the backbone.

Three methods were evaluated on the 3D segmentation task:

• Extension of nnU-Net [Wang et al., 2019]: An extension of the nnU-Net 13 with two sampling strategies.
• Bag of tricks (2nd place solution of the BraTS'19 challenge) [Zhao et al., 2019]: The winning entry used an ensemble of five models. For fairness, we quoted the best single-model results ("BL+warmup").
• Segtran-3D: I3D [Carreira and Zisserman, 2017] was used as the backbone.

5.3 Training Protocols

All models were trained on a 24GB Titan RTX GPU with the AdamW optimizer. The learning rate for the three transformer-based models were 0.0002, and 0.001 for the other models. On REFUGE20, all models were trained with a batch size of 4 for 10,000 iterations (27 epochs); on Polyp, the total iterations were 14,000 (31 epochs). On BraTS19, Segtran was trained with a batch size of 4 for 8000 iterations.

The training loss was the average of the pixel-wise crossentropy loss and the dice loss. Segtran used 3 transformer layers on 2D images, and 1 layer on 3D images to save RAM. The number of modes in each transformer layer was 4.

5.4 Results

Tables 4 and 5 present the evaluation results on the 2D and 3D tasks, respectively. Overall, the three transformer based methods, i.e., SETR, TransU-Net and Segtran achieved best performance across all tasks. With ResNet-101 as the backbone, Segtran performed slightly better than SETR and TransU-Net. With EfficientNet-B4, Segtran exhibited greater advantages.

12 https://github.com/Beckschen/TransUNet

13 https://github.com/woodywff/brats_2019

It is worth noting that, Segtran (eff-B4) was among the top 5 teams in the semifinal and final leaderboards of the REFUGE20 challenge. Among either REFUGE20 or BraTS19 challenge participants, although there were several methods that performed slightly better than Segtran, they usually employed ad-hoc tricks and designs [Orlando et al., 2020; Wang et al., 2019; Zhao et al., 2019]. In contrast, Segtran achieved competitive performance with the same architecture and minimal hyperparameter tuning, free of domainspecific strategies.

5.5 Cross-Domain Generalization

To explore how well different methods generalize to new domains, we trained three representative methods, U-Net, DeepLabV3+ and Segtran on the 1200 training images of REFUGE20. All the methods used a pretrained ResNet-101 as the encoder/backbone. The trained models were evaluated on both the REFUGE20 training images and the RIM-One dataset [Fumero et al., 2011]. As RIM-One images have drastically different characteristics from REFUGE20, all models suffered severe performance drop, as shown in Table 6. Nevertheless, Segtran had the least performance degradation, showing the best cross-domain generalization. Fig.6 shows a RIM-One image and the corresponding soft segmentation masks (before thresholding) produced by different methods. The mask produced by Segtran contains the fewest artifacts.

5.6 Computational Efficiency

Table 7 presents the number of parameters and FLOPs of a few representative methods. In general, transformer-based methods consume more computation and GPU RAM than conventional methods.

Our profiling showed that the number of parameters/FLOPs of Segtran are dominated by the output FPN, which vary drastically across different backbones. As the bottom-up FPNs we adopt are somewhat similar to EfficientDet [Tan et al., 2020], the model size/FLOPs are optimal when using EfficientNets. With ResNets as the backbone, Segtran has a significantly higher model size/FLOPs, and hence this choice of backbone is not recommended for efficiency-sensitive scenarios.

5.7 Impact of Pretraining

Models for medical image tasks usually benefit from initialization with weights pretrained on natural images (e.g. ImageNet [Deng et al., 2009]), as medical image datasets are typically small. To quantitatively study the impact of pretraining, Table 8 compares the performance of using pretrained weights vs. training from scratch of a fewmethods. Pretraining brought ~2.5% increase of average dice scores to the two transformer-based models, and 1% to U-Net (ResNet-101).

6 Conclusions

In this work, we present Segtran, a transformer-based medical image segmentation framework. It leverages unlimited receptive fields of transformers to contextualize features. Moreover, the transformer is an improved Squeeze-and-Expansion transformer that better fits image tasks. Segtran sees both the global picture and fine details, lending itself good segmentation performance. On two 2D and one 3D medical image segmentation tasks, Segtran consistently outperformed existing methods, and generalizes well to new domains.

Acknowledgements

We are grateful for the help and support of Wei Jing. This research is supported by A*STAR under its Career Development Award (Grant No. C210112016), and its Human-Robot Collaborative Al for Advanced Manufacturing and Engineering (AME) programme (Grant No. A18A2b0046).