MSF-TransUNet: A Multi-Scale Feature Fusion Transformer-Based U-Net for Medical Image Segmentation with Uniform Attention

Figure 1 shows the MSF-TransUNet architecture with three components: encoder, skip connection, and decoder. (1) The encoder learns global features by dividing the provided medical images into 2D patches, flattening them into sequences and using the self-attention mechanism of Transformer model. Besides, CB module enhances feature learning and model generalization. (2) Skip connections are redesigned in this new architecture by combining the low-dimensional features with the original skip connection layers, hence enhancing the decoder's capacity to include multi-scale features. Rather than directly adding features from a single encoder level to the decoder, the model makes use of multiple levels' information, which is obtained from the newly designed skip connections, through FFA-Block to enrich the learned representations. This ensures that the new skip connections' effectiveness by selectively enhancing and refining features obtained. Such an approach also enables the decoder to balance global contextual information with local fine details more efficiently. (3) The decoder consists of three core parts: Feature Concatenation, the FFA-block and up-sampling operations. The FFA-block, comprised of a Pre-Fusion block, Feature Fusion block, and Feature Selection block, which could adaptively grasp the most informative features, hence improving segmentation accuracy.

1.png

Figure 1. The structure of proposed MSF-TransUNet

2.1 Encoder incorporating CB module

MSF-TransUNet, like TransUNet [16], employs a hybrid architecture combining CNNs and Transformers in its encoder. It also incorporates the CB module [11], which enhances the model by promoting dense interactions. Previous research has shown the benefits of adding extra dense interactions in ViTs. Given the inherent challenge in learning dense attention through gradient descent, researchers manually implemented this process using a straightforward yet effective module called CB. The CB module is smoothly incorporated into the MLP layers of the MSF-TransUNet encoder, achieved with just a single line of code: X = 0.5﹡X + 0.5﹡X.mean(dim=1, keepdim=True) [11]. This CB module, by introducing uniform attention, inserts the token resulting from average pooling to each token individually, for infusing dense interactions. This integration not only simplifies the overall optimization process for MSF-TransUNet but also enhances its generalization. The CB module redirect MSF-TransUNet's focus from modeling dense attention maps to acquiring other valuable information, incurring only negligible additional operations for both inference and training.

2.2 Decoder

The architecture of the MSF-TransUNet decoder is akin to that of the TransUNet decoder. As shown in Figure 1, the MSF-TransUNet decoder utilizes three main operations: up-sampling operations, feature concatenation, and the FFA-Block. The model's final layer is a segmentation head that generates a feature map representing the prediction results.

2.2.1 FFA-block

To leverage features of varied scales extracted from the CNN part of the encoder and minimize the semantic disparity within the encoder-decoder architecture, optimizing the integration process of multi-scale features becomes imperative. In this paper, an FFA-Block, as illustrated in Figure 2, is introduced. The FFA-Block consists of a Pre-Fusion Block, a Feature Fusion Block, and a Feature Selection Block. This design aims to bridge the semantic gap effectively by utilizing multi-scale information while enhancing attention to capture salient features.

2.png

Figure 2. The schematic of FFA-block

Specifically, this paper utilizes skip connections with elaborate FFA-Block to merges multiple low-dimensional features with the ones obtained at the previous decoding stage, which could attain a comprehensive hierarchical feature map. This map is subsequently incorporated into a Pre-Fusion Block for initial feature extraction, followed by a Feature Fusion Block that assigns a unique spatial importance map for channels, which could direct the model to prioritize crucial regions within every channel [24], ensuring the interaction of information across different scales.

The Feature Selection Block generates a more efficient spatial attention map leveraging a convolutional modulation operation [25] after passing through two consecutive 3×3 convolutional layers. This block selects task-relevant features and filters out low-frequency information regions, which could augment the model's capacity to extract meaningful information.

2.2.2 Pre-fusion block

The Pre-Fusion block inputs from the previous decoder section and the CNN component within the encoder via skip connections. In this paper, Pointwise Convolution (PWConv) is chosen as the channel context aggregator, effectively combining information across channels and enabling the model to extract rich representational details from the input feature representations. After being processed by PWConv, these features are passed to DSC, which first applies a depth-wise convolution (DConv) independently to each individual channel, after which comes a PWConv that transforms the output channels of the DConv to a new channel space [26]. This structure allows for better fusion of features with inconsistent semantics and the extraction of spatial features from the input feature map while preserving channel information. Given $X \in R^{C \times H \times W}$, where X is the input with C channels and a size of H×W, $F(X) \in R^{C \times H \times W}$ is computed as follows:

$F(X)=B(P W \operatorname{Conv}(X)(\delta)(B(\operatorname{Conv}(\delta(B(P W \operatorname{Conv}(x)))))))$      (1)

Here, $\delta$ denotes ReLU6 activation function, B denotes batch normalization, PWConv denotes pointwise convolution, DConv denotes depthwise convolution and $F(X)$ denotes the feature map obtained from the concatenated features via Pre-Fusion block.

2.2.3 Feature fusion block

As shown in following equations, the feature map F, obtained from Pre-Fusion block, which contains multi-scale information, is subsequently processed by the Feature Fusion Block. To obtain a channel-specific spatial importance map [24], spatial global average pooling ($\mathrm{GAP}_s$) and spatial global max pooling ($\mathrm{GMP}_s$) are applied on spatial dimension of feature map $F \in \mathbb{R}^{C \times H \times W}$ to produce $F_{{avg }}^s \in \mathbb{R}^{1 \times H \times W}$ and $F_{{max }}^s \in \mathbb{R}^{1 \times H \times W}$. These operations extract spatial-wise feature information to encode critical spatial dependencies and discriminative patterns. Simultaneously, using global average pooling ($\mathrm{GAP}_{\mathrm{c}}$) across the channel dimension, a global channel descriptor $F_{a v g}^c \in \mathbb{R}^{c \times 1 \times 1}$ is generated from the feature map F.

$F_{ {avg }}^s=G A P_s(F)$        (2)

$F_{max }^s=G M P_s(F)$          (3)

$F_{a v g}^c=G A P_c(F)=\frac{1}{H \times W} \sum_{h=1}^H \sum_{w=1}^W F(h, w)$     (4)

$G A P_c$ denotes global average pooling across the spatial dimension. $G M P_s$ denotes global max pooling along the spatial dimension, highlighting the most salient spatial feature per channel. $G A P_c$ denotes global average pooling across the channel dimension. For every position (h, w), the feature values from all channels are averaged.

By concatenating $F_{a v g}^s$ and $F_{max }^s$ along the channel dimension to form a spatial descriptor and performing a convolution with 7×7 kernel size, more contextual information is captured, resulting in a spatial importance map $M_s$. Channel attention map is then obtained using Eq. (6), following the approach of SENet [27], where the core process involves 1×1 convolutions for channel compression and expansion, which adjusts the weight coefficients of each channel, thereby enhancing feature representation capabilities.

Additionally, residual connections address the vanishing gradient issue by incorporating the previous input feature F directly into later layers. By fusing $M_c$ and $M_s$ through element-wise summation, the semantic content of the input feature F serves as a guiding signal to derive the final channel-wise spatial importance map, ensuring alignment with the original feature structure. Finally, the fused features $F^{\prime}$ are computed using Eq. (7). This fusion approach achieves a dual contribution: it dynamically generates distinct spatial importance map per channel to prioritize critical regions, while simultaneously retaining fine-grained details from shallow layers and semantic context from deeper layers. The synergy of these properties ensures precise segmentation of anatomical structures in medical imaging by balancing semantic and detailed features.

$M_s={Conv}_{7 \times 7}\left[F_{{avg }}^s, F_{max }^s\right]$        (5)

$M_c={Conv}_{1 \times 1}\left(\delta\left({Conv}_{1 \times 1}\left(F_{a v g}^c\right)\right)\right)$      (6)

$F^{\prime}=F \odot \sigma\left(D {Conv}_{7 \times 7}\left(\left[F,\left(M_s+M_c\right)\right]\right)\right)$      (7)

$\delta$ is the ReLU activation function, while ${Conv}_{k \times k}$ represents a convolution with a k×k kernel, [·] represents channel-wise concatenation. To optimize computational efficiency and minimize model complexity, the architecture employs a bottleneck design (Eq. (6)), a ${Conv}_{1 \times 1}$ first compresses the channel dimension and a subsequent ${Conv}_{1 \times 1}$ restores the channel dimension from C/r to C. In this paper, r is empirically set to 8. ${DConv}_{7 \times 7}$ denotes a 7×7 depth-wise convolution to capture spatial dependencies. $\odot$ denotes the Hadamard (element-wise) product, and $\sigma$ represents the Sigmoid activation function.

2.2.4 Feature selection block

After features pass through Pre-Fusion block and Feature Fusion Block, they are forwarded to two standard convolution layers with a 3×3 kernel. Subsequently, a convolutional modulation (ConvMod) is employed to enhance spatial encoding efficiency within the Transformer U-Net framework. ConvMod is achieved by utilizing large kernels (≥7×7) within depth-wise convolutional layers [25, 28, 29], demonstrated in Eqs. (10)-(12). It facilitates task-relevant feature selection and attenuates redundant low-frequency information from the feature map $F^{\prime}$.

Self-attention [9] takes a $X \in \mathbb{R}^{N \times C}$, where N=H×W denotes the flattened spatial dimensions (height H, width W), and C represents the channel depth. The feature map X is first mapped through three learnable linear transformations to derive the key (K), query (Q), and value (V) matrices, each with dimensions $\mathbb{R}^{N \times C}$. The self-attention output is calculated as a weighted sum of the value matrix V, with the weights are determined by a normalized similarity score matrix A. This score matrix A is derived from the pairwise interactions between Q and K, capturing long-range dependencies across the input sequence.

${Attention}(Q, K, V)=A \cdot V$        (8)

$A=\frac{\operatorname{Softmax}\left(\mathrm{QK}^{\mathrm{T}}\right)}{\sqrt{d}}$      (9)

where, d is the dimension of Q and K vectors, serving as a scaling factor. The Softmax function normalizes the similarity scores into a probability distribution across the input sequence. The attention mechanism computes the output as context-aware aggregation of the V matrix, where the aggregation weights are determined by A. These weights encode pairwise affinities between Q and K, enabling the model to prioritize semantically relevant regions of the input.

Within ConvMod module, instead of computing the self-attention similarity matrix A as in traditional ViTs, self-attention is approximated by replacing matrix multiplications with the Hadamard product combined with depth-wise convolution operations to compute the output. To be specific, given $X \in \mathbb{R}^{H \times W \times C}$, the ConvMod module applies a simple depth-wise convolution with a K×K kernel to capture spatial patterns. An element-wise Hadamard product operation is subsequently performed to compute the output Z. The process is formally expressed as follows:

$Z=A \odot V$           (10)

$A=D {Conv}_{k \times k}\left(W_1 X\right)$       (11)

$V=W_2 X$               (12)

In those equations, $\odot$ represents the Hadamard product. $W_1$ and $W_2$ denote the weight matrices of two linear layers, while $DConv v_{k \times k}$ refers to a depthwise convolution with a K×K kernel. This formulation allows each spatial location (h, w) to capture correlations with all pixels in its surrounding K×K receptive field, ensuring effective local and global feature interactions.

Specifically, the Feature Selection block in MSF-TransUNet employs ConvMod with a 7×7 kernel size following the deepest skip connection, while the Feature Selection Blocks at other stages utilize a 11×11 kernel size. ConvMod streamlines feature selection by replacing the computationally expensive self-attention mechanism with a combination of depth-wise convolutions and the Hadamard product, significantly reducing computational overhead while preserving spatial dependencies.

As shown in Table 1, the computational complexity of ConvMod is O(N·C), whereas the complexity of self-attention is O(N²∙C). This quadratic complexity in self-attention becomes prohibitive when processing high-resolution medical images. In contrast, ConvMod provides a more scalable and efficient alternative, ensuring computational feasibility without sacrificing performance. By leveraging large kernels in a computationally efficient manner, ConvMod improves the extraction of features at various scales, rendering it especially effective for medical image segmentation tasks.

Table 1. Complexity of self-attention and ConvMod

2.3 Dataset and evaluation

This paper utilized the Synapse [30] and ACDC [31] datasets. The Synapse dataset comprises 30 CT-enhanced abdominal scans, each containing a varying number of slices. This dataset includes eight abdominal organs with corresponding labels. The Synapse dataset could be accessed through https://www.synapse.org/Synapse:syn3193805/wiki/217789.

The ACDC dataset consists of cardiac MRI scans from 100 patients, with labels for the left ventricle (LV), right ventricle (RV), and myocardium (MYO). The data splitting for both datasets follow the approach outlined in study [32]. The ACDC dataset could be accessed through https://www.creatis.insa-lyon.fr/Challenge/acdc/.

For the evaluation of our approach’s segmentation results, this paper employed DSC [33] and HD [34]. The corresponding equations are presented in Eq. (13) and Eq. (14):

$\operatorname{DSC}(M, N)=2 \frac{|M \cap N|}{|M|+|N|}$          (13)

$H D(M, N)=\max (h(M, N), h(N, M))$          (14)

Here, M and N represent the point sets corresponding to the Ground Truth and the predicted values, respectively. The term h(M,N) refers to the supremum of the minimum distances from all points in set M to set N, as shown in Eq. (15). This metric quantifies the largest deviation between the two sets.

$h(M, N)=\max _{m \in M} \min _{n \in N}\|m-n\|$     (15)

2.4 Loss function

To optimize the effectiveness of the segmentation approach, a composite loss function which incorporates both Cross-Entropy Loss and Dice Loss is implemented. The formulation is as follows:

$\begin{array}{r}{ Loss }=0.4 \times{ Cross }-{ Entropy\,\,Loss }+0.6 \times { Dice\,\,Loss }\end{array}$    (16)

$Cross-Entropy\,\,Loss =-\sum_{i=1}^n y_i\,\,log\,\,p_i$            (17)

$Dice\,\,Loss =1-\frac{2|A \cap B|}{|A|+|B|}$         (18)

Here, n represents the total amount of pixels from an input image, $y_i$ denotes the ground-truth for the i-th pixel and $p_i$ is the predicted probability for the same pixel. The sets A and B are the point sets for the Ground Truth and the predicted segmentation, correspondingly.

The designed composite loss function integrates the merits of both loss functions. Cross-Entropy Loss has a good performance in classifying each pixel independently, and therefore the predicted probabilities closely match the actual class labels. Dice Loss, on the other hand, focuses on where the ground truth and the expected segmentations overlap, which is particularly beneficial for addressing class imbalance and improving boundary precision of segmentation. By employing these two loss functions, the model strikes a compromise between accurate pixel classification and well-defined boundaries and therefore attains better overall segmentation performance.

2.5 Implementation details

The MSF-TransUNet model was implemented using Python 3.9, PyTorch 2.0.0, and CUDA 11.8 with an NVIDIA GeForce 4090 GPU. The images were pre-resized to 224×224 pixel. The model parameters were initialized by the pre-trained "R50-ViT" model. For the models trained with the SGD optimizer, a weight decay of 0.0001 was applied, and training was conducted for 150 epochs. The momentum was set to 0.9, and an initial learning rate of 0.01. The learning rate had a dynamic decay schedule, which was set as:

$l r=base_{l r} \times\left(1.0-\frac{\text { iter_num }}{\text { max_rations }}\right)^{0.9}$     (19)

where, iter_num denotes the current iteration, max_iterations represents the total number of iterations and $base _{l r}$ is the initial learning rate. This decay schedule enables a gradual reduction of the learning rate, which helps stabilize the optimization and promotes convergence.

MSF-TransUNet uses techniques for data enhancement commonly employed in TransUNet, such as random rotation, horizontal and vertical flipping, and spatial rescaling to the desired resolution, to increase the model's generalizability. The augmentations allow the model to support diversity in medical images and enhance the model's cross-dataset robustness.

4.1 Discussions on synapse dataset

This paper provides a comprehensive comparison performance of MSF-TransUNet, TransUNet, DA-TransUNet, Swin-UNet, and U-Net on the Synapse dataset with qualitative comparison illustrated in Figure 3 and quantitative comparison showed in Table 2. The results, illustrated in Figure 3, lead to several important conclusions. (1) Compared to other Transformer-based and hybrid models, models based entirely on CNNs tend to produce more mis-segmentation issues. Specifically, as shown in the third row of Figure 3, U-Net misclassifies the spleen, while in the first row, it produces false positives for the pancreas. This indicates that other Transformer-based and hybrid models have stronger global context encoding and semantic differentiation capabilities. (2) The results in the third row of Figure 3 indicate that MSF-TransUNet preserves the overall shape of the organs more precisely compared to other approaches and effectively reducing mis-segmentation. In contrast, U-Net, TransUNet, and DA-TransUNet exhibit varying levels of misclassification in the spleen segmentation, while Swin-UNet shows coarse segmentation boundaries and false-positive regions. These results reflect that MSF-TransUNet achieves superior accuracy and robustness in segmentation.

Experiments show that MSF-TransUNet performs better in segmentation tasks, preserving shape information and effectively representing global context as well as fine-grained information. This is evidenced by the observation that MSF-TransUNet performs better in most organs with the highest DSC values for the aorta (88.11%), gallbladder (71.33%), liver (94.83%), and pancreas (61.97%), and its performance on the kidneys and spleen remain stable. To ensure the statistical significance of the enhanced performance, the baseline model and MSF-TransUNet were compared through a paired t-test on the Synapse datasets. The results confirmed a statistically significant improvement in Dice score across the Synapse dataset (t = 3.031, p = 0.0191). This overall improvement reflects the success of MSF-TransUNet in effectively integrating global and local feature, making it especially robust for challenging segmentation tasks of medical imaging.

4.2 Discussions on ACDC dataset

MSF-TransUNet was applied to the segmentation task on ACDC dataset with performance measured using average DSC in percentage for RV, MYO, and LV. The experimental results, as shown in Table 3, indicate that MSF-TransUNet performs the best with an average DSC of 91.52%, surpassing other models like DA-TransUNet with 91.09% and TransUNet with 90.08%. It achieves the highest DSC for RV at 90.25% and LV at 96.59%, and demonstrates strong performance for MYO with 87.72%, which is better than the performance of most existing segmentation methods.

These results prove the effectiveness and strength of MSF-TransUNet in accurately segmenting different cardiac structures. In the case of the RV, the DSC of 90.25%, outperforming other models like VA-TransUNet (88.88%) and TransUNet (88.27%), indicates that MSF-TransUNet excels in identifying the structural edges and details of the RV, much better than the previous models. Similarly, the highest DSC for the LV also demonstrates excellent accuracy in delineating the LV, which is crucial for accurate cardiac assessments and treatments. The high performance in MYO segmentation also reflects the model’s ability to deal with the complexities of myocardial tissue, which is typically difficult to segment due to its complicated and highly variable nature.

4.3 Further work

The outstanding performance of MSF-TransUNet on the Synapse and ACDC datasets can be attributed to its hybrid Transformer and CNN structure. The structure effectively integrates global context with local feature. The incorporation of FFA-block and CB module allows the model to learn dense spatial interactions and extract key information more effectively, leading to more accurate and robust segmentation results.

Despite these advantages are observed, MSF-TransUNet faces some challenges. Firstly, the addition of FFA-blocks comes with a computational cost. Nevertheless, this increase in complexity is mitigated by the use of convolutional modulation (ConvMod) instead of self-attention, which effectively limits the overall computational burden. Further optimizations are still needed for real-time applications, particularly in low-resource environments. Future research will explore techniques such as model pruning and lightweight attention mechanisms to further reduce computational costs while preserving segmentation accuracy. Additionally, this present work addresses two-dimensional medical image segmentation, but most medical imaging modalities like CT and MRI provide three-dimensional volumetric data. In order to fill this gap, future work will focus on extending MSF-TransUNet to volumetric 3D medical image segmentation, assessing its effectiveness in volumetric scenarios. Moreover, future efforts will aim to optimize the model for real-time applications as well as further enhance its capabilities to three-dimensional medical image segmentation to leverage its full potential in clinical practice.