[Paper Review] A hybrid approach for improving U-Net variants in medical image segmentation

Aitik Gupta, Joydip Dhar|arXiv (Cornell University)|Jul 31, 2023

AI in cancer detectionComputer Science3 citations

TL;DR

This paper proposes a hybrid U-Net variant that integrates depthwise separable convolutions, residual connections, and an attention pooling mechanism to reduce model parameters by 97% while maintaining or improving segmentation accuracy on skin lesion and thyroid gland datasets. The approach achieves state-of-the-art performance with only 2.3 million parameters, significantly enhancing inference efficiency without sacrificing precision.

ABSTRACT

Medical image segmentation is vital to the area of medical imaging because it enables professionals to more accurately examine and understand the information offered by different imaging modalities. The technique of splitting a medical image into various segments or regions of interest is known as medical image segmentation. The segmented images that are produced can be used for many different things, including diagnosis, surgery planning, and therapy evaluation. In initial phase of research, major focus has been given to review existing deep-learning approaches, including researches like MultiResUNet, Attention U-Net, classical U-Net, and other variants. The attention feature vectors or maps dynamically add important weights to critical information, and most of these variants use these to increase accuracy, but the network parameter requirements are somewhat more stringent. They face certain problems such as overfitting, as their number of trainable parameters is very high, and so is their inference time. Therefore, the aim of this research is to reduce the network parameter requirements using depthwise separable convolutions, while maintaining performance over some medical image segmentation tasks such as skin lesion segmentation using attention system and residual connections.

Motivation & Objective

• To reduce the number of trainable parameters in U-Net variants for improved inference efficiency in medical image segmentation.
• To maintain or improve segmentation accuracy on challenging medical imaging tasks such as skin lesion and thyroid gland segmentation.
• To integrate depthwise separable convolutions to reduce parameter count while preserving feature representation capacity.
• To enhance feature learning through residual connections and attention pooling for better gradient flow and focus on relevant regions.
• To validate the hybrid architecture on public datasets and compare it with state-of-the-art models like MultiResUNet and Attention U-Net.

Proposed method

• Depthwise separable convolutions are applied in the encoder to reduce the number of parameters by decoupling spatial and channel-wise convolutions.
• Residual connections are added between encoder and decoder blocks to stabilize training and improve gradient flow.
• An attention pooling mechanism is introduced to dynamically emphasize important features in the encoder feature maps before decoding.
• The hybrid architecture combines U-Net’s skip connections with the efficiency of depthwise separable convolutions and the attention mechanism’s selective feature weighting.
• The model is trained end-to-end using binary cross-entropy loss and optimized via Adam with a base initial learning rate of 1e-4.
• Ablation studies are conducted to isolate the contribution of each component: depthwise convolutions, residual connections, and attention pooling.

Experimental results

Research questions

• RQ1 Can depthwise separable convolutions significantly reduce the number of parameters in U-Net variants without degrading segmentation performance?
• RQ2 How does the integration of residual connections affect training stability and segmentation accuracy in lightweight U-Net architectures?
• RQ3 To what extent does the proposed attention pooling mechanism improve feature representation and segmentation accuracy in low-parameter models?
• RQ4 How does the hybrid model compare to state-of-the-art U-Net variants like MultiResUNet and Attention U-Net in terms of accuracy and parameter efficiency?
• RQ5 Does the proposed architecture generalize well across diverse medical imaging tasks such as skin lesion and thyroid gland segmentation?

Key findings

• The proposed hybrid model achieved a Dice coefficient of 0.8872 on the HAM10000 skin lesion dataset, outperforming U-Net (0.8739) and Attention U-Net (0.8854).
• The model reduced the number of parameters to 2.3 million, a 97% reduction compared to MultiResUNet (64.8M) and a 93% reduction compared to Attention U-Net (34.88M).
• On the thyroid gland segmentation dataset, the hybrid model achieved a Dice coefficient of 0.9493 and an accuracy of 0.9766, surpassing U-Net (0.9332 and 0.9347) and matching MultiResUNet’s performance with far fewer parameters.
• The ablation study confirmed that the combination of depthwise convolutions, residual connections, and attention pooling yielded the highest IoU (0.9544) and lowest ASSD (0.0953) on the thyroid dataset.
• Visualizations showed the model outperformed U-Net in challenging cases such as blurred lesion boundaries, hair occlusions, and low-contrast regions.
• The model achieved the best balance between parameter efficiency and segmentation accuracy, demonstrating state-of-the-art performance with minimal computational cost.