1M parameters are enough? A lightweight CNN-based model for medical image segmentation

TL;DR

This paper introduces U-Lite, a lightweight CNN-based model for medical image segmentation that uses depthwise separable convolutions to significantly reduce parameters and computational cost while maintaining or improving performance.

Contribution

The paper proposes U-Lite, a novel lightweight U-Net variant utilizing Axial Depthwise and Dilated Depthwise Convolutions to achieve high accuracy with only 878K parameters, far fewer than traditional models.

Findings

U-Lite contains only 878K parameters, 35 times less than U-Net.

U-Lite achieves comparable or better segmentation performance.

The model significantly reduces computational complexity.

Abstract

Convolutional neural networks (CNNs) and Transformer-based models are being widely applied in medical image segmentation thanks to their ability to extract high-level features and capture important aspects of the image. However, there is often a trade-off between the need for high accuracy and the desire for low computational cost. A model with higher parameters can theoretically achieve better performance but also result in more computational complexity and higher memory usage, and thus is not practical to implement. In this paper, we look for a lightweight U-Net-based model which can remain the same or even achieve better performance, namely U-Lite. We design U-Lite based on the principle of Depthwise Separable Convolution so that the model can both leverage the strength of CNNs and reduce a remarkable number of computing parameters. Specifically, we propose Axial Depthwise Convolutions with kernels 7x7 in both the encoder and decoder to enlarge the model receptive field. To further improve the performance, we use several Axial Dilated Depthwise Convolutions with filters 3x3 for the bottleneck as one of our branches. Overall, U-Lite contains only 878K parameters, 35 times less than the traditional U-Net, and much more times less than other modern Transformer-based models. The proposed model cuts down a large amount of computational complexity while attaining an impressive performance on medical segmentation tasks compared to other state-of-the-art architectures. The code will be available at: https://github.com/duong-db/U-Lite.