MDEU-Net: Medical Image Segmentation Network Based on Multi-Head Multi-Scale Cross-Axis

1. Introduction

Medical image segmentation is a fundamental task in medical image processing with broad applications in areas such as disease diagnosis[1,2], surgical planning, and lesion monitoring [3,4,5,6]. Lesions are precisely identified and located using this technique that enables more accurate diagnoses. The development of precise surgical plans to minimize risks and complications is also supported while measurements of the size and volume of organs, tissues, or lesions are provided. Additionally, it offers quantitative metrics for assessing the condition and monitoring treatment outcomes.

Currently, two-dimensional segmentation models based on convolutional neural networks (CNNs)[7,8] have become widely adopted in medical image segmentation[9,10] . One of the most prominent models is U-Net[11] that features a unique U-shaped encoder-decoder architecture and improves segmentation accuracy through the introduction of skip connections. This is particularly useful for medical image processing tasks. So the success of U-Net has inspired the development of various network architectures including several U-Net variants. U-Net++[12] enhances model performance by replacing traditional connectivity with nested and dense skip connections[13,14]. Additionally, ResU-Net[15] models incorporate residual connections into U-Net to mitigate the gradient vanishing problem[16] during training and demonstrate strong performance in medical image segmentation tasks.With the continuous development and refinement of network models, MCANET has to some extent compensated for the inadequacy of traditional CNN networks in capturing long-distance dependencies through an attention mechanism.

Traditional convolutional neural networks (CNNs) and Transformer models are widely used in medical image segmentation to capture both local and global features. However, these networks struggle to achieve accurate segmentation and face particular challenges when dealing with complex and variable morphologies, such as lesions or organs with intricate structures and limited performance. In medical image segmentation the accuracy of segmentation is significantly affected by factors that contribute to challenges in achieving precise results, such as morphology, scale, and boundary ambiguity. As a result, developing methods to effectively capture multi-scale information in medical images[17,18] and improving model robustness and accuracy through flexible network architectures have become key research focuses in the field.

An improved method for medical image segmentation integrating cross-axis attention [19,20,21] and multi-scale feature fusion to enhance segmentation accuracy is presented in this paper. First, the introduction of the cross-axis attention mechanism enables the model to effectively capture global information from multiple directions within the image, thereby boosting the boundary recognition capabilities of complex lesion regions. Simultaneously, multi-scale feature extraction refines the feature representation of lesion regions at different scales and enhances the model’s ability to process lesions of varying sizes. Next, efficient feature fusion[22,23,24] is implemented through a gating mechanism that selectively emphasizes key features and mitigates the gradient vanishing problem via residual connections. This enhances the model’s ability to capture complex structures and details. Finally, the multi-head attention mechanism learns attention patterns across multiple subspaces in parallel to enhance the model’s robustness and segmentation accuracy for medical images with complex morphologies and ambiguous boundaries.

2. Network Architecture

The MDEU-Net model is proposed in this paper to address the challenges of multi-scale information extraction and detail capture in complex images and it incorporates the MDMSC Multihead Attention[25] module and the EF Efficient Feature Fusion module.

2.1. General Organization

The architecture of the MDEU-Net model enhances the U-Net encoder-decoder structure as shown in Figure 1. The encoder is composed of four convolutional modules that each consist of two convolutional layers[26] followed by a ReLU activation function[27] and a max pooling[28] operation. Furthermore, a multi-head multi-scale cross-axis attention mechanism is incorporated to enhance feature extraction. The decoder also consists of four convolutional modules that each contain an upsampling layer, an efficient feature fusion module, and two convolutional layers. The key innovations of the model include the incorporation of a multi-head multi-scale cross-axis attention mechanism at each downsampling stage which enhances multi-scale feature modeling. The model also adopts an efficient feature fusion module that replaces the traditional U-Net skip connections and improves feature fusion and information transfer efficiency.

2.2. Multi-Head Multi-Scale Cross-Axis Attention Module

As shown in Figure 2,the proposed architecture consists of eight parallel multi-scale channel (MSC) modules that extract multi-scale features from the image through downsampling and distribute the channels evenly to eight single-head attention modules. Each single-head attention module follows the MSC principle by extracting multi-scale features and then sequentially stacking the output channels of all modules. A residual connection (RC) is introduced that facilitates effective information transfer to mitigate the vanishing gradient problem and reduce information loss.

The multiscale cross-axis attention structure illustrated in Figure 2b features a network design that consists of two parallel branches, each of which is responsible for extracting the horizontal and vertical features of the image, respectively. Each branch employs a distinct 1D convolutional kernel size to capture multi-scale contextual information along a single spatial dimension.These features are integrated and enhanced in the orthogonal spatial dimension through a cross-axis attention mechanism that facilitates the capture of both horizontal and vertical information. The core idea of this approach is to leverage the multi-layer perceptron (MLP) capabilities of convolution to construct fine-grained hierarchical feature representations and to enable comprehensive and detailed extraction of image features. Taking the topmost branch in Figure 2b as an example,the input is $F^{*}\left({\displaystyle \frac{H}{8}},W,C\right)$ and the output[20]:

$$F_x=Conv1\times 1\left({\sum }_{i=0}^{2}Conv1D_i^x\left(Norm\left(F^{*}\right)\right)\right)$$

(1)

Where it denotes the one-dimensional convolution along the axis that we set to 1 × 7, 1 × 11 and 1 × 21 according to SegNeXt[29].$Norm\left(·\right)$ denoting layer normalization. For the bottom branch, the $F_y$ output can be expressed as[20]:

$$F_y=Conv1\times 1\left({\sum }_{i=0}^{2}Conv1D_i^y\left(Norm\left(F^{*}\right)\right)\right)$$

(2)

We propose to compute the cross-attention between $F_x$ and $F_y$ that aims to better utilize the multi-scale convolutional features in both spatial directions. $F_x$ will be used as the key matrix[30] and the value matrix[31], $F_y$ will be used as the query matrix[32]. The computation is as follows[20]:

$$F_t=MHC{A}_y\left(F_y,F_x,F_x\right)$$

(3)

Where $MHC{A}_y\left(·,·,·\right)$ denotes the cross attention along the x-axis. For the bottom branch feature extraction is performed along the y-axis in the same way as denoted below[20]:

$$F_b=MHC{A}_y\left(F_x,F_y,F_y\right)$$

(4)

$MHC{A}_x\left(·,·,·\right)$ denotes cross attention along the y-axis. For the obtained $F_t$ and $F_b$, the output of the proposed multiscale transaxial attention can be expressed as[20]:

$$F_n=Conv1\times 1\left(F_t\right)+Conv1\times 1\left(F_b\right)$$

(5)

The single-head attention works as above and the overall MDMSC output is expressed as:

$$Fout=F+F_m$$

(6)

Where F is the downsampled output. $F_m$ represents the feature output after fusion of 8 single-head attention mechanisms.

2.3. Efficient Feature Fusion Module

As illustrated in Figure 3, the EF module consists of two sub-modules: EA and EC. The EA module fuses low-level and high-level semantic features through skip connections and refines them using residual connections. The fused features are then passed to the EC module for further feature extraction.

The structure of the EA module is depicted in Figure 3c. Initially, group convolution is employed for intra-group feature fusion that contrasts with standard convolution and reduces computational complexity significantly. Following, the ReLU activation function is applied after convolving the input features. The low-level semantic features are then concatenated through residual connections that are transmitted via skip connections. Finally, the residual connection helps mitigate the interference of high-level semantic features on low-level features when the correlation between two input features is weak and this helps prevent degradation in overall model performance. The computational expression is as follows[33]:

$$EA\left(g,x\right)=x\times \left(1+Sigmoid\left(Conv1\times 1\left(Relu\left(W_g+W_x\right)\right)\right)\right)$$

(7)

$$W_g=Relu\left(BN\left(GroupConv32\left(g\right)\right)\right)$$

(8)

$$W_x=Relu\left(BN\left(GroupConv32\left(x\right)\right)\right)$$

(9)

Where sigmoid and ReLU is the activation function, BN is the Batch Normalization operation. GroupConv32 is the 32-component group convolution, and Conv1×1 is the regular convolution with a convolution kernel size of 1×1. In this model, g is the semantic feature obtained by up-sampling and x is the low-level semantic feature passed by jump connection.

The structure of the EC module is depicted in Figure 3b and it incorporates a residual network based on the ECA from S-Unet[34]. The EC module eliminates the need for additional weight parameters through the global average pooling (GAP) layer which consists of a GAP layer followed by a 1×1 convolutional layer. This significantly reduces the model’s parameter count and mitigates the risk of overfitting while improving computational efficiency. However, the GAP operation compresses the spatial dimensions of the feature map into a single global average which results in the loss of spatial structure information and potentially causes the loss of local spatial details. The limitation is effectively addressed by introducing the residual network that enhances the model’s ability to recover spatial information.

3. Experiments and Discussion of Results

3.1. Experimental Preparation

3.1.1. Experimental Data

This study utilizes the PanNuKe dataset for cell nucleus segmentation that consists of 300 accurately labeled microscopic images. For model training and evaluation, the dataset is split into 260 for training and 40 for testing.The image annotations in the PanNuKe dataset were manually performed by skilled technicians who covered the precise contours and locations of the cell nuclei to ensure that the model can learn accurate segmentation information. For the retinal blood vessel segmentation task, we utilized the Retinal Blood Vessels in the Eye (RITE) dataset from the DRIVE database that consists of 40 high-quality retinal blood vessel images. The dataset was split into 30 training and 10 testing groups for model training and evaluation. Each image is meticulously labeled with the exact contours and structural details of the blood vessels that aid the model in learning the intricate features of retinal blood vessels. Data augmentation techniques were applied during training to enhance model performance and fully utilize the dataset. These techniques expand the dataset in a way that increases sample diversity, reduces the risk of overfitting and improves the model’s generalization ability and robustness. The images in the training set were subjected to rotation, translation, scaling, and color transformations. Additionally, data augmentation helps mitigate noise, uneven illumination, and biases in biomedical images, which in turn improves segmentation accuracy.

3.1.2. Experimental Methods

All experiments were conducted using the PyTorch 2.0.0 framework on a computing platform with Python 3.8 (Ubuntu 20.04 operating system) and an Nvidia RTX 4090 GPU (24GB graphics memory). The MDEU-Net model was trained with the Adam optimizer, a learning rate of 0.0001, momentum of 0.9, and weight decay of 1e-8. The model was trained for 200 epochs on the DSB2018 dataset with a batch size of 4, and for 40 epochs on the RITE dataset with a batch size of 2. After training, the loss function curves showed a significant decrease in the nucleus segmentation task between epochs 0 and 100, which stabilized at approximately 0.04. Similarly, the ocular vascular segmentation task showed a notable reduction between epochs 0 and 30 that ultimately converged to around 0.1. The effectiveness of the MDEU-Net model was demonstrated during the training process when its performance gradually improved over time.

3.2. Experimental Evaluation Indicators and Results

3.2.1. Evaluation Indicators

Semantic segmentation is a pixel-level task in image segmentation. Commonly used evaluation metrics include recall, mean intersection over union (mIoU), and accuracy (Acc). All of these metrics are computed from the confusion matrix. The definitions of true positive (TP), false positive (FP), false negative (FN), and true negative (TN) are provided in Table 1.

TP: indicates that a sample is predicted to be a positive example and the true label is a positive example.

FN: indicates that a sample is predicted to be a counterexample and the true label is a positive example.

FP: indicates that a sample is predicted to be a positive example and the true label is a negative example.

TN: indicates that a sample is predicted to be a counterexample and the true label is a counterexample.

The binary semantic segmentation evaluation metric can be expressed as:

$$Acc={\displaystyle \frac{TN+TP}{TP+TN+FP+FN}}$$

(10)

$$mIoU={\displaystyle \frac{\frac{TP}{TP+FP+FN}+\frac{TN}{TN+FN+FP}}{2}}$$

(11)

$$Recall={\displaystyle \frac{TP}{TP+FN}}$$

(12)

3.2.2. Experimental Results and Analysis

Table 2 presents the performance comparison of seven deep learning models for cell nucleus and retinal blood vessel segmentation tasks. As shown in Table 2, the MDEU-Net model significantly outperforms other medical image segmentation models in both cell nucleus and retinal blood vessel segmentation. Compared to the original U-Net model in the task of cell nucleus segmentation, MDEU-Net improves the mIoU, accuracy (Acc), and recall by %11.52, %3.51, and %6.19,respectively. Improvements of %8.33, %0.6, and %5.59 in mIoU, accuracy, and recall are observed for retinal blood vessel segmentation, respectively. MDEU-Net achieves a %0.47 improvement in mIoU for the cell nucleus segmentation task that demonstrates its effectiveness compared to the more advanced MCANET model. For the retinal vessel segmentation task it demonstrates improvements of %1.02 in mIoU, %0.13 in accuracy (Acc), and %2.40 in recall, respectively. These results demonstrate that MDEU-Net achieves superior segmentation performance which is particularly evident in recall rate. It also maintains a relatively small model parameter size and shows notable improvements in both accuracy and intersection over union (IoU).

As shown in Figure 5, the MDEU-Net model demonstrates exceptional segmentation performance in both retinal blood vessel and cell nucleus segmentation tasks, while maintaining low computational complexity. This indicates that the model not only enhances segmentation accuracy but also reduces the demand for hardware resources.

As shown in Figure 6, it compares the segmentation results of the MDEU-Net model with those of other networks for cell nucleus and retinal blood vessel segmentation. MDEU-Net efficiently captures global features and accurately recognizes object boundaries and contours. Furthermore, the model excels in local feature extraction that is particularly effective in handling complex backgrounds and fine structures.

As shown in Figure 7, a comparison of the magnified prediction maps for different segmentation models in specific regions containing labeled maps that involve MDEU-Net, S-Unet, and Res-Unet is presented. The MDEU-Net model demonstrates superior contour accuracy in the nucleus segmentation magnification map, which shows a better fit with the labeled map than the other models. In retinal vascular segmentation magnification maps where MDEU-Net outperforms other models in terms of vascular continuity and detail restoration, it places particular emphasis on the precise reconstruction of capillaries and vessel termini that closely align with the labeled maps.

The multi-head attention mechanism is a critical component in modern neural networks. Unlike single-head attention, which shares query, key, and value matrices, multi-head attention assigns independent matrices to each head and enables the model to focus on different parts of the input simultaneously. The optimal number of heads is not universally fixed and varies depending on the specific task and model architecture. To evaluate the impact of the multi-head attention mechanism on segmentation performance, Table 3 presents segmentation metrics for different head configurations across the nucleus and retinal vessel datasets. It is important to note that the number of output channels in the first downsampling layer of the MDEU-NET model is 64 and that means the number of heads must be a divisor of 64 to ensure proper model operation. As shown in the table, the configurations MDEU2-NET, MDEU4-NET, MDEU8-NET, MDEU16-NET and MDEU32-NET correspond to 2, 4, 8, 16 and 32 heads respectively. Experimental results indicate that the model with 8 heads achieves superior segmentation performance in both the nucleus and retinal vessel tasks compared to other configurations.

The MDEU-NET model innovatively introduces the MDMSC multi-head attention module and the EF efficient feature fusion module on top of the U-Net architecture. To evaluate the contribution of these two modules to the model’s segmentation performance, we designed an ablation experiment. Specifically, we removed each module individually and conducted segmentation experiments on the PanNuKe dataset and the RITE dataset. Each experiment employed the Adam optimizer with a learning rate of 0.0001, momentum of 0.9, and weight decay set to 1e-8. The results are presented in Table 4 and Table 5, demonstrating the impact of each module on the overall segmentation performance.

Table 4 and Table 5 show the impact of the ablation experiment on the segmentation performance of the RITE dataset and the PanNuKe dataset, respectively. As can be seen from the segmentation performance results presented in the tables, the removal of either the MDMSC multi-head attention module or the EF efficient feature fusion module has an effect on the segmentation performance. To further verify the importance of these two modules, we conducted a visual analysis of the segmentation performance of the model after removing each module. These visual results reinforce the experimental evidence and also provide a more intuitive perspective for understanding the specific contribution of each component in the segmentation task. The segmentation results are shown in Figure 8.

4. Conclusion

This paper presents MDEU-NET, a novel medical image segmentation network based on the U-Net architecture, which incorporates several optimizations and enhancements. First, a multi-head attention mechanism is employed, enabling the model to capture information from multiple perspectives simultaneously. This significantly improves the model’s representational capacity and performance compared to single-head attention. Second, a multi-scale cross-axis attention mechanism is designed to address the variability of organs and lesion sites in medical images, enabling each attention head to focus on features across different scales and axes, thereby enhancing the segmentation process. Finally, the Efficient Feature Fusion (EF) module is introduced as a superior alternative to traditional feature fusion methods. This module integrates the attention mechanism with a gating module, enabling detailed information to be captured at multiple levels, thereby enhancing the emphasis on critical features and improving both feature extraction accuracy and segmentation performance. Experimental results demonstrate that the MDEU-NET model excels in segmentation tasks, with a significant improvement in recall rate and notable gains in accuracy and intersection ratio.