MixKNet: A Modified U-shaped Network with Hybrid Channel Convolution for Medical Image Segmentation

Summary

This study presents a lightweight U-Net variant with improved segmentation accuracy using hybrid convolution and channel attention, achieving state-of-the-art results on MoNuseg and GlaS with fewer parameters.

Abstract

Artificial neural network-based computer image processing technology has rapidly developed in recent years and found widespread applications across multiple fields. In medical image processing, UNet and its variants have shown great success in lesion detection, cell segmentation, and polyp segmentation tasks. This research presents a modified U-shaped network with reduced network parameters achieved by decreasing the network depth and increasing the network channels to enhance the model's learning ability. To counteract the reduction in learning capacity caused by depth compression, a hybrid-channel convolutional module is introduced to replace the original network's convolutional module. The model introduces a channel attention mechanism between the layer and the point convolution layer to improve the practical channel feature extraction ability. The article also concludes that using mixed-depth convolution can effectively solve the size span of the segmentation target, making it difficult for one model to be fit for multiple data sets. The proposed model achieves state-of-the-art results on two popular public datasets, MoNuseg and GlaS, with a mean dice increase of 1.0% and 1.37%, respectively. The total parameters of the modified model are reduced to 1.71M, representing a 38.6x reduction compared to UCtransNet, with 65.6M parameters.

Introduction

Medical image segmentation is critical in various clinical applications, including diagnosis, treatment planning, and disease monitoring. Traditional image processing methods based on feature engineering algorithms are no longer suitable for handling the large volumes of data generated in modern healthcare environments. Fortunately, the advancement of artificial neural network technology and improvements in computer hardware capabilities have made it possible to train and deploy large-scale neural network models. As a result, computer vision processing algorithms based on machine learning models are continuously being developed and applied in various fields.

The most representative model structure in the semantic segmentation of medical images is UNet¹ with an encoding-decoding architecture. The model first uses a multi-level encoder to extract features of different scales from the input image. Then, the decoder up-samples step by step while combining the semantic features output by the encoder of the corresponding level as a skip connection. However, the performance of the UNet architecture can be further improved by applying several methods. For instance, attention mechanisms have been shown to be effective in enhancing the performance of convolutional neural networks. These mechanisms can selectively emphasize the essential features in the input data, leading to better representation learning and segmentation accuracy.

At present, many researchers focus on effectively fusing the features extracted by the encoder in the decoding stage. Some of the most common variants of UNet include Attention UNet², Recurrent UNet³, and V-Net⁴. These approaches employ attention mechanisms⁵, such as spatial attention, to highlight informative regions of the feature maps and suppress the non-informative regions. Additionally, some variants incorporate recurrent neural networks to model the sequential nature of the medical images and improve the feature representations. Pre-training models, such as VGG⁶, ResNet⁷, and DenseNet⁸, have been widely used to improve the performance of U-shaped networks by leveraging their rich knowledge learned from large-scale datasets. Additionally, transformer mechanisms, such as self-attention and multi-head attention, have been introduced to long-range model dependencies and capture global contextual information, leading to better segmentation performance.

However, while these variants have improved over the original UNet¹, they still have certain limitations in handling complex medical images with varying scales and shapes. Moreover, the increased complexity of these variants often leads to higher computational costs and longer training times, which limits their applicability in practical settings.

To enhance the UNet¹ architecture, a modified model named MixKNet (Mix Kernel Size U-shape Net) is proposed, which reduces network depth while increasing the number of channels to improve learning capacity. However, the reduction in depth can lead to a decline in overall performance. To mitigate this issue, a hybrid-channel convolutional module is introduced, replacing the original convolutional neural module. This module incorporates a channel attention mechanism between the depthwise and point-wise convolution layers, aiming to strengthen the model's ability to extract effective channel features.

To guide practical applicability, it is important to note that the proposed method was evaluated on relatively small-scale medical image datasets, with individual image resolutions ranging from 500 × 500 to 1000 × 1000 pixels. For model training, all images were resized to 224 × 224 pixels. The implementation was carried out using PyTorch on a single NVIDIA RTX 3090 GPU. These operational parameters suggest that the method is computationally feasible for moderate experimental setups and adaptable to limited dataset sizes.

In conclusion, this article presents a novel modification to the U-shaped network structure and introduces a hybrid channel convolution module along with a channel attention mechanism, both of which improve segmentation performance on medical image datasets. The main contributions of this work are as follows: (1) Proposing a modified U-shaped network structure with a hybrid deep-wised convolution module that replaces the original convolutional neural module. (2) Introducing a channel attention mechanism between the deep-wised and the point-wised convolution layer to improve the model's effective channel feature extraction ability. (3) Achieving state-of-the-art results simultaneously on two popular public datasets in the MoNuseg⁹ nuclear segmentation dataset with significantly fewer model parameters (compared to UCtransNet¹⁰ with 64M parameters) and improved performance on the GlaS¹¹ gland segmentation dataset.

The rest of this article is organized as follows. Step 2 provides a comprehensive review of previous studies on UNet¹ and its variations in medical image segmentation. The strengths and limitations of existing approaches are examined, along with related work on attention mechanisms, mixed-depthwise convolutional networks, and their applications in segmentation models. Step 3 details the architecture of the proposed MixKNet model, including modifications to the UNet structure, the integration of a channel attention mechanism, and the use of mixed-depth convolution. Step 4 reports the results of extensive experiments, accompanied by a discussion and an ablation study to evaluate the contributions of each component. Finally, step 5 concludes the paper and outlines potential directions for future research.

This section begins with a review of prior studies on UNet and its variants in medical image segmentation, outlining the strengths and limitations of existing methods. In addition, related research on attention mechanisms and their applications in segmentation models is discussed. Recent developments involving depthwise convolution techniques for enhancing segmentation performance are also examined. A comparative analysis of the proposed MixKNet (c) and other models is presented in Figure 1, where dashed lines indicate skip connections.

UNet and its variations
The UNet architecture was first proposed by Ronneberger et al. in 2015¹ and has since been widely used in medical image segmentation. The model consists of an encoding path and a decoding path. The encoder extracts high-level features from the input image while the decoder up-samples and combines the features to generate a segmentation map. Since then, various modifications have been made to the UNet architecture to improve its performance in medical image segmentation. One of the most common modifications is the introduction of skip connections, which have been shown to improve the accuracy of the segmentation. Zhou et al.¹² proposed a UNet variant that uses dense skip connections, allowing the model to preserve spatial information better.

Another popular modification to the UNet architecture is the addition of attention mechanisms. These mechanisms can help the model selectively emphasize the most important features, leading to better representation learning and segmentation accuracy. Some examples of variants with attention mechanisms include Attention UNet², ResUNet with attention gates¹³, and Dense-UNet with attention gates¹⁴. In addition to the above modifications, many other variants of the UNet architecture have been proposed for medical image segmentation. UCTransNet¹⁰ is a variant of the U-Net architecture that incorporates skip connections and a Transformer module. The skip connections in UCTransNet¹⁰ are rethought from a channel-wise perspective, which allows for better feature reuse and reduces the number of parameters. The Transformer module is added to capture long-range dependencies and improve translation quality. UCTransNet¹⁰ has been shown to achieve state-of-the-art results on several machine translation benchmarks. Zihan et al. proposed a deep learning model called LViT¹⁵, which they applied to medical image analysis tasks. To extend the semi-supervised version of LViT¹⁵ and compensate for the quality deficiency in image data, they proposed the Exponential Pseudo Label Iteration mechanism (EPI). Additionally, to preserve the local features of images, they proposed the Pixel-Level Attention Module (PLAM), which selectively focuses on important image features.

Attention mechanism
Attention mechanisms have been widely studied in machine learning, particularly in natural language processing and computer vision. In recent years, attention mechanisms have also been applied to medical image analysis tasks, including image segmentation. Bahdanau et al.¹⁶ introduced an attention mechanism in neural machine translation to improve the quality of translation. The mechanism allows the model to selectively focus on relevant parts of the input sequence, improving the translation quality. Since then, various attention mechanisms have been proposed for image analysis tasks. One common type of attention mechanism is the spatial attention mechanism, which involves applying a weight to each pixel in the feature map based on its importance. For example, Guo et al.¹⁷proposed a spatial attention mechanism for the task of retinal vessel segmentation. The mechanism uses a gating mechanism to adjust the weighting of the spatial features, improving the model's ability to distinguish between the vessel and non-vessel pixels. Another type of attention mechanism is channel attention, which focuses on adjusting the weights of feature maps across channels. Hu et al.¹⁸ proposed a squeeze-and-excitation network (SENet) that applies channel-wise attention to the feature maps, improving the performance of the image classification task. More recently, attention mechanisms have been combined with convolutional neural networks (CNNs) for image segmentation tasks. For example, Chen et al.¹⁹ proposed an attention-guided network for brain tumor segmentation that combines spatial and channel-wise attention mechanisms.

Recent advances in semi-supervised and multimodal learning for medical image analysis and remote sensing have demonstrated promising performance improvements. Yang et al.²⁰ proposed UMSCS, a novel unpaired multimodal segmentation framework that leverages cross-modality generative learning and semi-supervised strategies. Zhang et al. introduced several cutting-edge techniques: an evidential-enhanced tri-branch consistency model for semi-supervised segmentation²¹, a self-aware and cross-sample prototypical learning framework for medical image segmentation²², and a time-frequency-aware hierarchical feature optimization approach for synthetic aperture radar (SAR) jamming recognition in the aerospace domain²³. These works collectively highlight the importance of hybrid supervision and domain-aware feature modeling in both medical and engineering imaging tasks.

Depth-wised convolution
Depth-wise convolution is designed to capture channel-wise correlations in the input feature maps and has been shown to improve the efficiency and accuracy of convolutional neural networks.

One of the earliest and most influential depth-wise convolution models is MobileNet²⁴, proposed by Howard et al. in 2017. MobileNet uses depth-wise separable convolution, which applies a depth-wise convolution followed by a point-wise convolution, to reduce the number of parameters and computation required for convolutional layers. This allows MobileNet to achieve state-of-the-art accuracy on image classification tasks while using significantly fewer parameters than traditional convolutional neural networks. Since the introduction of MobileNet, a number of other depth-wise convolution architectures have been proposed, including ShuffleNet²⁵, Xception²⁶, and ResNeXt²⁷. These models vary in their specific implementation of deep-wise convolution, but all share the goal of reducing the computational complexity of convolutional layers while maintaining or improving accuracy. Deep-wise convolution has also been applied to the task of semantic segmentation, with models such as PSPNet²⁸ using depth-wise separable convolution to reduce the computation and memory requirements of large-scale convolutional layers. MixNet²⁹ further extends the idea of deep-wise convolution by introducing a mixed depth-wise convolution, which utilizes multiple kernel sizes to capture features at different scales. It has been shown to outperform other state-of-the-art models while maintaining comparable parameters and FLOPs. These models have demonstrated significant improvements in accuracy and efficiency over traditional convolutional neural networks on various semantic segmentation tasks.

Protocol

This section describes the proposed MixKNet architecture for medical image segmentation. The MixKNet model extends the UNet architecture, incorporating attention mechanisms and mixed depthwise convolution layers. The software used in this study is listed in the Table of Materials.

1. Overall architecture

1. The MixKNet architecture, shown in Figure 2, follows the standard UNet structure, consisting of an encoder and a decoder. Set the input shape to 224 × 224 × 3. In the encoder, use two DownBlock modules-each includes a 1×1 depth-wise convolution, a 2×2 max pooling layer, and a multi-kernel convolution (MDConv) with kernel sizes such as [1, 3, 5] or [3-13].
2. In the decoder, use two UpBlock modules-each performs bilinear upsampling, skip connection concatenation, a 1×1 depth-wise convolution, and a multi-kernel convolution.
3. Use ReLU activations after each convolutional block. Apply a final 1×1 convolution to produce the output, followed by a Sigmoid activation for binary segmentation, or leave unactivated for external Softmax in multi-class tasks.

2. Flattened U-shape

NOTE: Reduce the depth of the neural network architecture to decrease computational complexity and model size, while being aware that this may reduce the capacity to learn complex data representations.

1. Reduce a four-level UNet architecture to a two-level version, as shown in Figure 3, to decrease the number of model parameters while retaining the core encoding and decoding components.
   1. Construct the flattened architecture by removing levels 3 and 4 from both the encoder and decoder paths, and preserve feature connections at level 2.
   2. Adjust the downsampling and upsampling operations accordingly to match the reduced depth-use only one downsampling step before the bottleneck and one upsampling step after.
      NOTE: This reduction in depth may decrease the model's learning ability, as it may not capture complex relationships between input and output. Increasing the number of base convolution channels in each model layer can help overcome this limitation by enhancing the model's capacity to learn discriminative features. With more convolution channels, the model can capture more complex patterns and relationships in the input data, compensating for the reduced depth and improving performance.
2. Consider that increasing the number of convolution channels can also raise the model's computational complexity and memory usage, presenting a trade-off with training and inference speed. Find an appropriate balance between model capacity and computational efficiency.

3. Mix kernel size for the encoder

1. To enhance encoder performance, employ a mixed kernel size approach in the DC-CA (Depthwise Convolution and Channel Attention) module, as illustrated in Figure 4. Instead of relying on a single kernel size, apply multiple depthwise convolutions in parallel with varying kernel sizes (e.g., 1, 3, 5, 7, 9, 11, 13) to extract features at multiple receptive fields.
2. Apply batch normalization and ReLU activation to each branch individually before concatenation. Concatenate the outputs from each kernel branch along the channel dimension.
3. Use a point-wise 1×1 convolution after concatenation to fuse features and project them to a fixed number of output channels, maintaining consistency with the expected input size of the next layer.
   NOTE: The mixed depthwise convolution layer consists of multiple depthwise convolutions with varying kernel sizes, followed by a point-wise convolution. This design enables the capture of features at different scales, which is critical in medical image segmentation tasks. The first module employs three kernel sizes-1, 3, and 5-to increase the receptive field, while the second module utilizes larger kernel sizes and more groups, specifically 3, 5, 7, 9, 11, and 13.

4. Channel attention

1. Integrate a channel attention mechanism into the DC-CA module to enhance feature representation.
   1. Use a 1×1 convolution with "same" padding to generate channel-wise attention maps. Apply global average pooling across spatial dimensions to produce a compact descriptor for each channel.
   2. Use a lightweight convolutional network consisting of two 1×1 convolutions with a ReLU activation in between. Avoid weight sharing across channels-learn unique attention weights per channel.
   3. Apply a sigmoid activation to the final output to obtain normalized attention scores.Reweight the input feature maps via channel-wise multiplication.
      NOTE: Let the input tensor be x. Channel attention is implemented through the following expression:
      (attention)_tensor = σ(conv(ReLU(conv(x)))) (1)
      where σ represents the sigmoid activation function, conv is the convolution operation, and ReLU is the rectified linear unit activation function.
2. Obtain the attention map by performing an adaptive average pooling operation on the attention tensor, followed by expanding it to match the size of the output tensor after convolution with a different kernel size y:
   attention = expand(avg_pool([attention]_tensor),size(y)) (2)
3. Compute the output tensor z by performing element-wise multiplication of the input tensor y with the attention map:
   z = z * attention (3)
   where * represents element-wise multiplication

5. Datasets

NOTE: Two publicly available medical image datasets, as presented in Table 1 and Table 2, are utilized in this study: GlaS (Sirinukunwattana et al.¹¹) and MoNuSeg (Kumar et al.⁹).

1. Use the Gland Segmentation dataset (Sirinukunwattana et al.¹¹), acquired using a digital pathology scanner and manually annotated to segment individual glandular structures within tissue samples.
   NOTE: The dataset consists of 165 images, each with a resolution of 500 × 500 pixels, along with corresponding ground truth segmentation masks. The MoNuSeg dataset (Kumar et al.⁹) comprises 51 images of microscopic nuclei, each with a resolution of 1000 × 1000 pixels.

6. Implementation details

1. Implement the model using PyTorch on a single NVIDIA RTX 3090 GPU. Use Ubuntu 22.04 as the operating system.
2. Resize all input images to a fixed size of 224 × 224 pixels. Apply data augmentation techniques to enhance training diversity.
   1. Apply augmentations randomly in the following order: horizontal flip → vertical flip → rotation → gamma correction → logarithmic transformation → random scaling.
   2. Apply random horizontal flips with a probability of 0.5, vertical flips with a probability of 0.5, and rotations within ±15 degrees with a probability of 0.5.
   3. Apply gamma correction with a probability of 0.3, logarithmic transformation with a probability of 0.3, and random scaling with a scaling factor between 0.9 and 1.1, applied with a probability of 0.3.
   4. Normalize images using mean values of [0.485, 0.456, 0.406] and standard deviations of [0.229, 0.224, 0.225].
3. Train the model using the Adam optimizer with a learning rate of 0.001 and a batch size of 4. Employ the Cosine Annealing with Warm Restarts learning rate scheduling technique to introduce variability into the learning rate schedule and prevent convergence to local optima.
   NOTE: Use PyTorch's built-in torch.optim.lr_scheduler.CosineAnnealingWarmRestarts scheduler. The optimizer and scheduler were implemented as follows:
   optimizer = torch.optim.Adam(model.parameters(), lr=0.001, weight_decay=1e-5)
   scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(
   optimizer, T_0=10, T_mult=2, eta_min=1e-6)
   Here, T_0=10 denotes the number of epochs before the first restart, T_mult=2 doubles the restart period after each cycle, and eta_min=1e-6 sets the minimum learning rate. Update the learning rate at the end of each epoch by calling scheduler.step().
4. Use binary cross-entropy as the loss function. Apply early stopping to prevent overfitting. Train the model for a maximum of 5000 epochs.
5. Evaluate model performance using the mean Intersection over Union (mIoU) and the Dice coefficient.
   Dice =(2|A∩B|)/(|A|+|B|) (4)
   IoU=(|A∩B|)/(|A∪B|) (5)
   NOTE: The mIoU is the ratio of the intersection between the predicted and ground truth segmentation masks to their union. At the same time, the Dice coefficient is the ratio of twice the intersection to the sum of the predicted and ground truth segmentation masks.
6. Conduct a rigorous evaluation on small datasets using three rounds of 5-fold cross-validation, resulting in a total of 15 evaluations. Randomly shuffle the cross-validation splits in each round to ensure variability across folds. Stratify the folds based on class distributions to maintain label balance within each fold.
7. Compute the mean performance across folds to mitigate the risk of convergence to local optima. Perform a statistical test to demonstrate that the proposed approach achieves statistically significant improvements over comparable methods.
8. Show representative validation predictions during training in Figure 5 to illustrate the progressive improvement in segmentation quality. After training completes, load the best-performing model checkpoint based on validation performance (e.g., highest mIoU and Dice score).
   1. Run the model on the validation set to generate predicted segmentation masks using the saved parameters.
   2. Visualize the results using Matplotlib by displaying the input image, ground truth mask, and predicted mask in a side-by-side format.
   3. Highlight the superiority of the proposed method in segmentation performance by marking representative regions with red boxes in the visualization.

Results

Comparison with state-of-the-art methods
The MixKNet architecture has been evaluated on two medical image segmentation datasets, GlaS and MoNuSeg, and compared with state-of-the-art methods in the field. The evaluation has been conducted by reporting the dice and IoU scores of the proposed method and several comparable models, as shown in Table 3. The results indicate that the MixKNet architecture outperforms the other models, achieving the highest me...

Discussion

This study presents MixKNet, a novel modification of the UNet¹ architecture for medical image segmentation, integrating mixed-depth convolution and a channel attention mechanism to enhance segmentation performance while significantly reducing computational complexity. The hybrid channel convolution combines multiple convolutional kernels operating on separate channel groups. This design allows the network to capture diverse spatial and contextual features at different receptive field scales within...

Materials

List of materials used in this article

Name	Company	Catalog Number	Comments
Linux OS (Ubuntu 22.04)	Various (Open Source Community)	N/A(Version 22.04 LTS)	Open-source operating system for development https://releases.ubuntu.com/22.04/
NVIDIA RTX 3090 GPU	NVIDIA	3090	High-performance GPU for deep learning https://www.nvidia.com/en-us/geforce/graphics-cards/30-series/rtx-3090/
Python	Python Software Foundation	N/A (Version ≥ 3.8)	Programming language for AI/ML development https://www.python.org/
PyTorch	Meta AI	N/A (Version 1.13)	Open-source machine learning framework https://pytorch.org/
Visual Studio Code (VS Code)	Microsoft	N/A	Lightweight and powerful code editor https://code.visualstudio.com/