NestedFormer: Nested Modality-Aware Transformer for Brain Tumor Segmentation

Recent works show that transformer is more conducive to modeling global information than CNNs. To better extract local context information for each modality, we extend the Poolformer [yu2021metaformer] as the modality-specific encoder. As discussed in [yu2021metaformer], replacing the computation-intensive attention module in Transformer with average pooling can achieve superior performance than recent transformer and MLP-like models. Therefore, to enhance global information, we design Global Poolformer Block (GPB), which leverages global pooling instead of average pooling in Poolformer, followed by a fully connected layer. As shown in Fig. 2, given the input feature embedding $X$, a GPB block consists of a learnable global pooling (GP) and a MLP sub-block. The output $Z$ is computed as,

where $LN(\ast )$ denotes the layer normalization and $W_g$ is the learnable parameter in the FC layer. Our Global Poolformer encoder contains five groups of one feature embedding (FE) layer and two GPB blocks. Each FE layer is a 3D-convolution, while the first FE layer has a convolution patch size of $1\times 1\times 1$ and the rest layers have a patch size of $2\times 2\times 2$ and a stride of 2. The encoders gradually encode each modality image into high-level feature $F_{L,i}\in R^{d\times w\times h\times C},i\in [1,M]$, where $(d,w,h)=(\frac{D}{16},\frac{W}{16},\frac{H}{16})$ are $1/16$ of input spatial resolutions $H$, $W$ and depth dimension $D$; $M$ is the number of modal images, the channel dimension $C$ and the layer number $L$ are set as $C=128$, $L=5$.

Given high-level features $F_{L,i},i\in [1,M]$, NMaFA leverages a spatial-attention based transformer $T_{tsa}$ and a cross-modality attention based transformer $T_{cma}$ in a nested manner; see Fig. 3. First, transformer $T_{tsa}$ utilizes the self-attention scheme to compute the long-range correlation between different patches in the space within each modality. Specifically, $F_{L,i}$ is concatenated in the channel dimension to obtain high-level embedding $F_s\in R^{d\times w\times h\times MC}$. In this work, each location of $F_s$ is considered as one “patch”. Then a patch embedding layer maps $F_s$ to a token sequence ${\hat{F}}_s\in R^{dwh\times C}$. $T_{tsa}$ takes ${\hat{F}}_s$ and the position encoding [vaswani2017attention] as the input, and outputs spatially-enhanced feature ${\tilde{F}}_s\in R^{dwh\times C}$.

Second, transformer $T_{cma}$ utilizes the cross-attention scheme to further compute the global relation among different modalities to achieve inter-modality fusion. To this end, $F_{L,i}$ is concatenated in the spatial dimension to obtain the flatten sequence ${\hat{F}}_c\in R^{MP\times C}$. Here, $P(P=32)$ denotes the number of dominant tokens learnt via the Token Learner strategy [ryoo2021tokenlearner], which helps to reduce the computational scope especially when the number of tokens increases greatly along with more modalities. After that, both ${\tilde{F}}_s$ and ${\hat{F}}_c$ are fed into $T_{cma}$ to get the modality-enhanced feature embedding $\tilde{F}$.

Also note that our two modules are different from traditional channel-spatial attention networks, which reweigh feature maps channel-wise and spatial-wise. Our NMaFA relies on transformer mechanism and the two transformers are fused in a nested form, rather than serial [khanh2020enhancing] or parallel [mou2019cs] fusion.

In feature decoding, we first fold the tokens $\tilde{F}$ back to a high-level 4D feature map $R_L\in R^{d\times w\times h\times C}$. $R_L$ is progressively processed in a regular bottom-up style with a 3D convolution and $2\times$ upsampling operation to recover a full resolution feature map $R_1\in R^{D\times W\times H\times N_c}$ for segmentation, where $N_c$ is the number of segments. Note that the encoder features are multi-modal. Hence, we design a modality-sensitive gating strategy in skip connection, to filter the encoder features $\{F_{l,i},l\in [1,4],i\in [1,M\left]\right\}$ according to the modality importance. To be specific, for the $l$-th layer, an modality importance map $I_l\in R^{\frac{D}{2^{l-1}}\times \frac{W}{2^{l-1}}\times \frac{H}{2^{l-1}}\times M}$ is learnt from $\tilde{F}$ that is the output of NMaFA, as follows,

where $FC(\ast )$ is a $1\times 1\times 1$ full connection layer, $U_l^{2\times }$ denotes upsampling $l$ times, and $\sigma (\ast )$ is the sigmoid function. Denote $\odot$ as element-wise multiplication. Then the filtered encoder feature $F_l$ is formulated as

Our NestedFormer was implemented in PyTorch1.7.0 on a NVIDIA GTX 3090 GPU. The parameters were initialized via Xavier [glorot2010understanding]. The loss function was a combination of soft dice loss and cross-entropy loss and we adopted the AdamW optimizer [loshchilov2017decoupled] with a weight decay of ${10}^{-5}$. The learning rate was empirically set as ${10}^{-4}$.We adopted two $T_{tsa}$ sequentially and just one $T_{cma}$. In $MH{A}_w$, the window-size was set as (2, 2, 2) for BraTS2020 and (2, 4, 4) for MeniSeg.

For evaluation, we use a public brain tumor segmentation dataset BraTS2020 [menze2014multimodal] and a private 3D meningioma segmentation dataset (MeniSeg) collected from Brain Medical Center of Tianjin University, Tianjin Huanhu hospital. Dice score and 95% Hausdorff Distance (HD95) are adopted for quantitative comparison.

We compare our network against seven SOTA segmentation methods, including three CNN-based methods (3D-UNet [cciccek20163d], SegResNet [myronenko20183d], MAML [zhang2021modality], nnUNet [isensee2021nnu]), and three transformer-based methods (SwinUNet(2D) [cao2021swin], TransBTS [wang2021transbts], and UNETR [hatamizadeh2022unetr]). For a fair comparison, we utilize the public implementations of compared methods to re-train their networks for generating their best segmentation results. Considering the computation power, all the methods are trained for at most 300 epochs on BraTS2020 and 200 epochs on MeniSeg.

We conduct ablation studies on the MeniSeg dataset to evaluate the contributions of main modules in our method; see Table 3. We not only compared the effects of three different encoder backbones based on CNN, PB, and GP, but also verified the effect of our proposed fusion modules. Among them, $baseline1$ uses multiple U-Net encoders to extract features of different modal images, and performs feature fusion by concatenation. $baseline2$-$baseline4$ uses multiple GPB encoders to extract features, and conducts skip connection via simple convolution, w/o $T_{tsa}$ and $T_{cma}$ (see Fig. 3), respectively. $baseline5$ replaces GPB block with the original PoolFormer block (PB) in the encoder, using the proposed NMaFA module (including $T_{tsa}$ and $T_{cma}$) as well as MSG. It can be observed clearly that compared with $baseline2$, using the NMaFA module enhances the extraction of long-distance dependency information and effectively improves the segmentation results, while GPB outperforms PB by considering global information. Moreover, the MSG module is added to increase the feature reuse capability of skip connections, which further improves the segmentation effect, achieving the best average segmentation Dice (0.765) on the MeniSeg dataset.