ScaleFormer: Revisiting the Transformer-based Backbones from a Scale-wise Perspective for Medical Image Segmentation

1 Introduction

In the past decades, Convolutional Neural Networks (CNNs) have greatly promoted the development of dense prediction tasks [Shelhamer et al., 2015], and a variety of CNN-based segmentation models have been developed [Ronneberger et al., 2015; Chen et al., 2017; Huang et al., 2021]. Despite its great success, CNNs failed to model explicit long-range relations beyond local regions (shown in Fig.2.(b) and (f)), since the effective receptive field of a networkÕs units is severely limited [Luo et al., 2016]. 1

* Corresponding Authors: Lanfen Lin (llf@zju.edu.cn), YenWei Chen (chen@is.ritsumei.ac.jp).

Recently, the transformer architecture has been introduced to visual tasks and shown to be an efficient way of modeling global context [Dosovitskiy et al., 2020]. This largely attributes to the multi-head self-attention (MSA) and Multilayer Perceptron (MLP) structure. Fig.1. (a) illustrates one of the transformer-based architectures for medical image segmentation, e.g. SwinUNet [Cao et al., 2021] and Missformer [Huang et al., 2021]. They typically utilized a stack of transformer blocks to model global representations. However, these pure transformer-based architectures often failed to achieve satisfactory performance, since lack of spatial inductive-bias in modelling local information. As shown in Fig.2(c) and (g), SwinUNet activated the background extensively and was difficult to distinguish objects from background. Another popular architecture, e.g. TransUNet [Chen et al., 2021] and AFTer-UNet [Yan et al., 2022], in Fig.1(b), used a CNNbased encoder to explore detailed high-resolution spatial information, and added several transformers into the deepest layer to learn global context. However, this kind of deepest

Huimin Huang and Shiao Xie are co-first authors.

hybrid transformer structure still led to blurred and inaccurate activation on the salient objects in Fig.2.(d) and (h).

The observed unsatisfactory activations of current transformer-based backbones may be imputed to two main challenges in a scale-wise perspective. (i) Intra-scale problem: As we know, both contextual local-level details and positional global-level cues are essential complements in each scale, especially for small objects, whose fine signals in the shallow layer (in Fig.2(b)-(d)) may be gradually diluted during downsampling (in Fig.2 (f)-(h)). But the integration of intra-scale local features with global representations remains a fundamental challenge. (ii) Inter-scale problem: As shown in Fig.2, the objects in medial image often have widely variable size, shape and location. The existing methods were incapable of extracting sufficient information from multiple scales generated by hierarchical encoder, which failed to notice the complete salient objects. To address these challenges, we propose a novel backbone for medial image segmentation task (in Fig.1. (c)), namely ScaleFormer, which is motivated by the following two aspects:

Firstly, considering both the spatial difference and global distribution of objects are essential characteristics, a scalewise intra-scale transformer is designed to couple the CNNbased local features with the transformer-based global representations in each scale. In this way, we are able to highlight both detailed spatial information (e.g. contextual cues) and long-range dependencies (e.g. positional cues). However, the quadratic computational cost of previous transformers becomes the main obstacle in such intra-scale structure, especially for the shallow scales with high-resolution feature maps. To achieve this, we also design a lightweight and universal Dual-Axis MSA to suppress less-relevant information and recognize salient parts with fast speed, which is realized by capturing row-wise and column-wise global dependencies.

Secondly, as many previous studies demonstrated, multiscale features explore distinctive information and are capable of resolving the complex scale variations. In this inter-scale circumstance, for a specific object (e.g. liver), instead of calculating the redundant relationships with long-distant objects (e.g. spleen and aorta) at various scales, we focus on the longrange cross-scale dependency of the object itself and the nearby objects (e.g. liver and gallbladder), which can effectively capture the scale difference (e.g. shape variation) of each object, and achieve promising performance with satisfactory boundary. Inspired by this, we design a spatial-aware inter-scale transformer considering that there is a spatial correspondence among patches in various scales through downsampling operation. In this way, we focus on the cross-scale similarities among these spatial-aware patches, which efficiently learns the mutual information in a lightweight manner.

The ability of ScaleFormer in capturing local-global cues is shown in Fig.2(a), (e). By extracting the intra-scale local details and global semantics, ScaleFormer can focus on the small object in both shallow and deep layers. Additionally, benefiting from the inter-scale interactions, ScaleFormer also attends full object extent with in various sizes and long-range scenarios. The major contributions of this work are four-fold: (i) We analyze the intra-scale and inter-scale problems faced by the current transformer-based backbones, and propose a novel backbone in a scale-wise perspective, termed as ScaleFormer, to improve the segmentation quality of medical image segmentation. (ii) We propose a scale-wise intra-scale transformer to associate CNN-based local features with intrascale transformer-based global cues in each scale. Accordingly, we design a lightweight and universal Dual-Axis MSA to capture row-wise and column-wise global dependencies. (iii) We propose a spatial-aware inter-scale transformer to interact among consensual regions in multiple scales, which is capable of capturing cross-scale dependency and resolving complex scale variations in a simple and effective manner. (iv) We conduct extensive experiments on three public datasets, which shows ScaleFormer surpasses SOTA methods.

2 Related Work

Vision transformers. In the most current research, ViT [Dosovitskiy et al. 2020] introduced transformer into visual tasks and achieved competitive performance. Motivated by the success of ViT, SegFormer [Xie et al., 2021], SETER [Zheng et al., 2020], DETR [Carion et al., 2020] have been proposed for semantic segmentation and object detection tasks, which achieved comparative results with CNN-based methods. To reduce the computational cost of self-attention operation in ViT, CMT [Guo et al., 2021] utilized depth-wise convolution to downsample key and value to extract important information; while Swin transformer [Liu et al. 2021] restricted the attention in local regions. Different from existing methods, we design a novel backbone that couples local CNN blocks and global Transformer. Such a structure not only inherits sufficient structural prior from CNN but also embeds global representations from transformer block in each scale.

Multi-scale architectures for medical image segmentation.

Multi-scale fusion has been proved to improve the segmentation performance in CNN-based methods. PSPNet [Zhao et al., 2017] exploited the capability of global context information by aggregating different region-based features. UNet++ [Zhou et al.,2020] and UNet 3+ [Huang et al., 2020] have been proposed to bridge the semantic gap between shallow encoder layers and deep decoder layers and made full use of multi-scale information through skip connection. To learn more extensive and rich context information, MS-DualGuided network [Sinha et al., 2020] applied attention in both spatial and channel dimension of feature maps on different scales. Considering the multi-scale cues, we further propose a simple and effective spatial-aware inter-scale transformer, which captures long-distance relationships among scales.

3 ScaleFormer

3.1 Overview

The proposed ScaleFormer learns intra-scale local-global cues as well as modeling inter-scale dependencies as shown in Fig.3. In special, we utilize ResUNet backbone to extract local features in a hierarchical manner. The encoder part has five CNN blocks comprised of basic resnet-34 blocks, where the channel number increases with network depth while the resolution of feature maps decreases. When the CNN goes deeper, different fine-grained features are extracted in various stages. Considering that transformer is powerful at capturing long-distance relationships, the scale-wise intra-scale transformer (in Sec.3.2) is coupled with the corresponding CNN block to highlight both detailed spatial information and long-range dependencies in each scale. To learn mutual regions at different scales, the spatial-aware inter-scale transformer (in Sec.3.3) further interacts among multiple scales. Finally, the enhanced representation is aggregated with the local-level CNN features in the same scale, and then sends to the decoder block for the final prediction. In the following, we will describe each step of ScaleFormer in detail.

3.2 Intra-Scale Transformer Block

As shown in Fig.3, the initial transformer only receives input from the CNN branch of same stage, whereas the input of other intra-scale transformer blocks incorporates global information from previous transformer stages to aggregate fine-grained details and coarse semantics information. Let i indexes the down-sampling layer along transformer branch, 𝑁# refers to the total number of transformers. The stack of feature maps 𝑋#%&'( ) can be represented as:

where function 𝐷𝑜𝑤𝑛(∙) indicates the patch embedding mechanism, which is realized by a convolution operation with stride of 2 and followed by a batch normalization and a ReLU activation function. [á] is the concatenation operation. 𝑇𝑟𝑎𝑛𝑠(∙) is the intra-scale transformer block, which consists of a lightweight Dual-Axis MSA module, and an enhanced MLP. We will describe these parts in the following.

Dual-Axis MSA. In the Fig.4, we illustrate the previous two MSA architecture and our proposed method. In original MSA [Dosovitskiy et al., 2020] (shown in Fig.4.(a)), the input 𝑋 ∈ ℝB×D×E is linearly transformed into query 𝑄, key 𝐾 and value 𝑉 with the same shape of 𝐻𝑊×𝐶. The computational cost of the original MSA is 𝑂(𝐻M𝑊M𝐶). To achieve lower cost, some works [Guo et al., 2021] used a spatial reduction mechanism by applying a reduction ratio 𝑅 to the key 𝐾 and value 𝑉 (shown in Fig.4.(b)) and the computational cost is decreased to 𝑂(𝐻M𝑊M𝐶 𝑅M ). However, it still requires large reduction ratio 𝑅 when applied to high-resolution feature map, which degrades the horizontal and vertical information.

Different from the previous MSAs, we design a lightweight and universal Dual-Axis MSA mechanism, which replaces the full-image relation (𝐻𝑊×𝐻𝑊) with a row-wise (𝐻×𝐻) and column-wise (𝑊×𝑊) relation, as illustrated inFig.4.(c). To adaptively learn the horizontal and vertical representation [Chi et al., 2020], we first apply four distinctive depth-wise convolutions on 𝑄 and 𝐾 to obtain row-wise enhanced 𝑄%, 𝐾% and column-wise enhanced 𝑄1, 𝐾1. We then pass these into average pooling operation in two directions, which estimates global and compact conditions (𝑄%O, 𝐾%O, 𝑄1O, 𝐾1O) over the whole feature map. The process can be formulated as:

when applying the row-wise MSA, the value 𝑉 is reshaped into 𝐻×𝑊𝐶; while using the column-wise MSA, 𝑉 is reshaped into 𝐻𝐶×𝑊. In this way, the complexity of our DualAxis MSA is 𝑂(𝐻M𝑊𝐶 +𝐻𝑊M𝐶). Compared with Axial Attention [Ho et al.,2019] that requires consecutive two blocks to capture long-term relations, our method learns global conditions using a single block. Then, we utilize an enhanced MLP [Huang et al., 2021] with recursive skip connection before the depth-wise convolution for feature delivery.

3.3 Spatial-Aware Inter-Scale Transformer Block

Multi-scale cues have been proven to be powerful in resolving complex scale variations, which are essential for medical image segmentation. Hence, we aim to design a transformerbased block to interact among different scales based on the hierarchical structure of ScaleFormer. The input sequence of enhanced intra-scale tokens is first rearranged into 2D lattice. Instead of calculating the dependencies on a long sequence that directly concatenated from multi-scale tokens (flattened from the complete 2D feature map of each scale), we take full advantage of the spatial correspondence among patches in different scales through down-sampling operation, which is called Patch Matching. As an example, Fig.5.(a) illustrates how to find the spatial-aware patches on three successive feature maps from i-th to (i+2)-th scale. Note that we first regularize feature maps into the same channel dimension 𝐶. Let denote 𝑝d) as the j-th patch in the i-th scale, and the corresponding down-sampeled patches are 𝑝d)e; and 𝑝d)eM sharing the bounding box with the same color (i.e. blue). In this way, we focus on the most related patches, and thus maintain the spatial correspondence as well as reduce the redundancy.

Then we apply Scale Fusion to concatenate spatial corresponding inter-scale patches, which can be formulated as:

where 𝑓𝑙𝑎𝑡𝑡𝑒𝑛(∙) operation rearrange the patch 𝑝d) into the shape of 𝐻𝑊 2M) ×𝐶 and [·] represents the concatenation operation. Because the sequence of tokens lacks horizontal and vertical information, in this part, we utilize the original MSA and MLP proposed by [Dosovitskiy et al., 2020] in our Transformer block, instead of using our Dual-Axis MSA (shown in eq.(2)). The spatial-aware inter-scale dependencies among patches can obtained as follows:

where 𝐿𝑁(∙) denotes the layer normalization operator. Then we utilize Scale Division to reverse the enhanced sequence back to patches according to the order of concatenation:

where 𝑆𝑝𝑙𝑖𝑡(∙) is an inverse operation of previous concatenation operation [·]. In the reverse procedure of patch matching, patch reversion assembles all patches in the same scale into a feature map. Then we concatenate the global-level transformer feature and the local-level CNN feature as skip connection, which further sends to the decoder stage.

Discussion. In this section, we analyze the successive operations of the intra-scale and inter-scale transformers. The local features firstly passed through intra-scale transformer to extract the contextual information, which learns the entire pixelwise dependency in each scale. Then, the enhanced feature maps are fed into inter-scale transformer, which explores the corresponding patch-wise dependencies among multiple scales. In this way, even though two pixels (𝑥)& , 𝑥ds ) in different scales ( 𝑖 and 𝑗 ) and irrelevant patches (𝑎 and 𝑏), in Fig5.(c), their relationship can be also captured via first finding the intra-scale dependencies of 𝑥)& and 𝑥)s , then calculating the inter-scale dependencies of 𝑥)s and 𝑥ds . These two consecutive steps realize the goal of capturing pixel-wise dependencies in full scales with less computational cost.

4 Experiments

4.1 Dataset and Evaluation

We perform experiments on three public medical image datasets:

Synapse. The multi-organ segmentation challenge [Landman et al., 2015] consists 30 abdominal CT scans. Following [Chen et al., 2021], we split 18 cases for training and remaining 12 cases for testing. We reported the Dice Coefficient (DSC) and Hausdorff Distance (HD) on 8 different organs.

MoNuSeg. The Multi-Organ Nucleus Segmentation Challenge [Kumar et al., 2017] contains 30 images for training, and 14 images for testing. Following [Wang et al., 2021], we utilized DSC and Intersection over Union (IoU) metrics.

ACDC. The automated cardiac diagnosis challenge [Bernard et al., 2018] contains 100 MRI scans with three organs, left ventricle (LV), right ventricle (RV) and myocardium (MYO). Following [Chen et al., 2021], we reported the DSC with a random split of 70 training cases, 10 validation cases and 20

testing cases. We split the dataset randomly because the previous division was unavailable.

4.2 Implementation Details

We utilized data augmentations to avoid overfitting, including random rotation and flipping. Note that our ScaleFormer did not require pretrain on any large datasets. Here, we itemized the batch size (bs), learning rate (lr), maximum training epochs (ep), optimizer (opt) for three datasets:

• ∞ Synapse: bs=8; lr=3e-3; ep=600; opt=SGD;
• ∞ ACDC: bs=8; lr=3e-3; ep=200; opt=SGD;
• ∞ MoNuSeg: bs=4; lr=1e-3; ep=200; opt=Adam;

All models were trained with momentum 0.9 and weight decay 1e-4. For fair comparison, we used the same settings and combined cross entropy loss and dice loss for all experiments.

4.3 Comparison with the State-of-the-Arts

To demonstrate the effectiveness of ScaleFormer, we first compare it with 11 state-of- the-art methods on the Synapse dataset, including V-Net [Milletari et al., 2016], DARR [Fu et al., 2020], U-Net, Att-UNet [Schlemper et al., 2019], ViT, TransUnet, SwinUNet, AFTer-UNet, MISSFormer. Both TransUnet and SwinUNet are pretrained on the ImageNet. Here, we abbreviate ResNet-50 as ÒR50Ó, which is utilized as encoder and combined with U-Net, Att-UNet and ViT.

Experimental results are reported in Table 1 and the best results are highlighted in bold. As shown, our ScaleFormer without inter-scale former (w/o inter) outperforms all the other methods in both regional measures DSC (82.08%) and boundary-aware measure HD (18.03mm). In particular, our ScaleFormer (w/o inter) surpasses other transformer-based methods, including SwinUNet and MISSFormer with pure transformers, as well as TransUNet and AFTer-UNet that stack several transformers in the deep layers of encoder. It indicates that the structure combined with CNN and intrascale transformer has the ability of capturing local-global clues. Additionally, an increment of (DSC: +0.78% and HD: -1.22mm) is achieved when utilizing inter-scale transformer, which achieves the new SOTA on synapse dataset.

To demonstrate the generality of ScaleFormer, we conduct experiments on another two public datasets. Additionally, we compare it with UNet++, MRUNet [Ibtehaz et al.,2020], MedT [Valanarasu et al.,2021] and UCTransNet [Wang et al., 2021], which are four SOTA methods on MoNuSeg dataset. Table 2 presents DSC and IoU metrics on MoNuSeg, and Table 3 reports DSC on our divided ACDC dataset (*) with retrained TransUnet and SwinUnet using their original settings. Our ScaleFormer achieves the best result for capturing more useful representations, which demonstrates the superior generalization and robustness of ScaleFormer.

Fig.6 shows the qualitative comparison of classic UNet, TransUNet, SwinUNet and our ScaleFormer on Synapse dataset. It can be observed that ScaleFormer not only accurately localizes organs but also produces coherent boundaries, even in small object circumstances.

4.4 Ablation Study

Ablation of intra-scale transformers on different stages. We investigate how intra-scale transformers affect performance among different scales on synapse dataset, which is varied from Stage6 (deepest layer with size of 7×7) to Stage2 (shallowest layer with size of 112×112). As shown in Fig.7, the accuracy is improved with stacking the transformer blocks, which achieves the highest results (DSC: 82.08%, HD:18.03mm) when extracting intra-scale local-global cues on four consecutive stages (Stage6543). However, applying the transformer on the Stage2 may affect the genuine characteristics and increase the computational cost (required two GPUs). Therefore, in our experiments, we introduce the intrascale transformers on Stage3-Stage6, which provides the best results within a considerable computational complexity.

Ablation of the input for intra-scale transformer block. Based on the architecture with four-stage intra-scale transformers, we present the ablation of its input, i.e. the samestage CNN and the previous-stage transformer (Pre-Trans). As shown in Fig.8 that only using pre-trans in a staking manner achieves inferior results than using CNN feature (DSC: - 0.80%, HD: +3.05mm), which means that staking transform ers may dilute detailed signals by extracting global information. By combining two inputs, the best results (DSC: 82.08%, HD: 18.03mm) are achieved for aggregating finegrained details and coarse-grained semantics.

Effectiveness of our Dual-Axis MSA. To verify the capability of Dual-Axis MSA in capturing global cues, we compare it with CMT [Guo et al., 2021], PVT [Wang et al. 2021], and Axial attention [Ho et al., 2019] in Fig.9(b). We find that our method achieves a better accuracy (higher DSC and lower HD) than these methods with less computation cost and fewer parameters. It indicates the superiority of Dual-Axis MSA via splitting the full-image dependency in the row-wise and the column-wise manner.

Ablation of inter-scale transformers on various stages. Additionally, we conduct several experiments to assess the influence of inter-scale transformers on various stages. As shown in Fig.9 (a), by learning interactive information on deeper stages, an increment from 81.37% to 82.86% is achieved, where 0.67% and 0.82% improvements are produced by adding Stage5 and Stage6. It indicates that the learned mutual relations can enhance segmentation quality.

4.5 Interpretation of ScaleFormer

We visualize the feature map from the final decoder in Fig.10. As shown, ResUNet (in Fig.10.(f)) has the capability of focusing on the local areas but it fails to activate larger regions; while TransUNet (in Fig.10.(e)) and SwinUNet (in Fig.10.(d)) can globally model the long-term relationships but decline the important detailed local features. Impressively, our ScaleFormer (w/o inter) (in Fig.10.(b)) not only inherits sufficient structural prior from CNN but also embeds the global representations from transformer block, which helps to find the tiny objects. Considering the multi-scale cues, ScaleFormer (in Fig.10.(c)) has more complete attention areas on objects with widely variable size, shape and location.

5 Conclusion

In this paper, we proposed a novel transformer-based backbone, namely ScaleFormer, from a scale-wise perspective to improve the medical image segmentation quality even for the small objects. In each scale, a scale-wise intra-scale transformer with a lightweight Dual-Axis MSA is designed to combine the CNN-based local features and transformerbased global features; while among multiple scales, a simple and effective spatial-aware inter-scale transformer was proposed to interact across scales. Experiments showed that ScaleFormer outperformed both CNN-based and transformer-based segmentation networks on three datasets.

Acknowledgements

This work was supported in part by Zhejiang Provincial Natural Science Foundation of China under the Grant No. LZ22F020012, Major Scientific Research Project of Zhejiang Lab under the Grant No. 2020ND8AD01, and in part by the Grant-in Aid for Scientific Research from the Japanese Ministry for Education, Science, Culture and Sports (MEXT) under the Grant No. 20KK0234, No. 21H03470 and No. 20K21821.