Source-linked AI summary

CMUNeXt: An Efficient Medical Image Segmentation Network based on Large Kernel and Skip Fusion

Fenghe Tang, Jianrui Ding, Lingtao Wang, Chunping Ning, S. Kevin Zhou

arXiv:2308.01239v2eess.IVcs.CV

TL;DR

CMUNeXt targets the difficulty of extracting global context under the computational constraints of lightweight medical image segmentation, while retaining convolutional inductive bias for scarce data. It uses large-kernel and inverted-bottleneck convolutional blocks with Skip-Fusion, and reports an effective trade-off between segmentation performance and computation for edge-oriented deployment.

  • Problem

    U-shaped convolutional networks have limited global-context extraction, while Transformer-based alternatives face resource constraints and lack convolutional inductive bias for scarce medical data.

  • Method

    CMUNeXt is a lightweight fully convolutional network using large-kernel depthwise convolutions, inverted bottlenecks, and Skip-Fusion for global-context extraction and feature fusion.

  • Results

    CMUNeXt achieves an optimal trade-off between segmentation performance and computation consumption across the reported experiments.

  • Takeaways & Limitations

    Three CMUNeXt variants target different edge and device environments for fast and accurate auxiliary diagnosis in real-world scenarios.

Abstract

from arXiv · show

The U-shaped architecture has emerged as a crucial paradigm in the design of medical image segmentation networks. However, due to the inherent local limitations of convolution, a fully convolutional segmentation network with U-shaped architecture struggles to effectively extract global context information, which is vital for the precise localization of lesions. While hybrid architectures combining CNNs and Transformers can address these issues, their application in real medical scenarios is limited due to the computational resource constraints imposed by the environment and edge devices. In addition, the convolutional inductive bias in lightweight networks adeptly fits the scarce medical data, which is lacking in the Transformer based network. In order to extract global context information while taking advantage of the inductive bias, we propose CMUNeXt, an efficient fully convolutional lightweight medical image segmentation network, which enables fast and accurate auxiliary diagnosis in real scene scenarios. CMUNeXt leverages large kernel and inverted bottleneck design to thoroughly mix distant spatial and location information, efficiently extracting global context information. We also introduce the Skip-Fusion block, designed to enable smooth skip-connections and ensure ample feature fusion. Experimental results on multiple medical image datasets demonstrate that CMUNeXt outperforms existing heavyweight and lightweight medical image segmentation networks in terms of segmentation performance, while offering a faster inference speed, lighter weights, and a reduced computational cost. The code is available at https://github.com/FengheTan9/CMUNeXt.

Introduction

CMUNeXt addresses the difficulty of extracting global context in lightweight medical image segmentation while retaining convolutional inductive bias and deployment efficiency. It combines large-kernel feature extraction with Skip-Fusion and reports favorable segmentation-performance and computational trade-offs.

  • Motivation: U-shaped convolutional networks struggle to extract global information because ordinary convolutions have inherent local limitations.Global context is important for precise lesion localization.
  • Motivation: Transformer-based segmentation can reduce parameters, but lightweight CNNs better fit scarce medical data through convolutional inductive bias.The supplied passage contrasts Transformer-based performance and data-fitting behavior with lightweight CNNs.
  • Method: CMUNeXt uses large-kernel depthwise convolutions and inverted-bottleneck pointwise convolutions to mix distant spatial and location information with few parameters.The CMUNeXt block is introduced in the encoder to extract global context efficiently.
  • Method: The Skip-Fusion block uses grouped convolution and inverted-bottleneck pointwise convolutions to smoothly fuse encoder and decoder semantic features.This design replaces ordinary convolutions in skip-connections for sufficient feature fusion.
  • Evaluation: Ultrasound datasets were selected because lesion positions vary and lesion morphology is multi-scale, providing a test of global-context extraction.These characteristics are presented as reasons the datasets can illustrate the method’s effectiveness.
  • Evaluation: Experiments demonstrate a better trade-off between segmentation performance and computational consumption, with released code and a U-shaped segmentation benchmark.The benchmark release is identified separately from the reported experimental trade-off.

Architecture Overview

CMUNeXt is organized as a five-layer encoder-decoder U-shaped network with skip-connections. CMUNeXt blocks extract multilevel global context in the encoder, while Skip-Fusion integrates encoder semantics with decoder features.

  • Architecture Overview: The architecture has five layers divided into an encoder stage and a decoder stage with skip-connections.The encoder and decoder form the network’s two main architectural stages.
  • Encoder: Each encoder level applies a CMUNeXt block, an ordinary convolution block, and down-sampling to extract and expand hierarchical features.The passage describes this sequence as part of the encoder architecture.
  • Decoder: Each decoder level uses Skip-Fusion to fuse encoder global semantic features with up-sampled decoder features.The skip-connections provide the interface for this feature fusion.

Encoder Stage

The encoder extracts hierarchical global context through five levels, combining CMUNeXt blocks, channel expansion, and down-sampling. Its stem preserves top-level skip-connection consistency, while pooling is considered suitable for low-resolution medical images.

  • Encoder Structure: The encoder has five levels, each containing a CMUNeXt block, an ordinary convolution block, and a down-sampling operation.A stem extracts original input features at the top level.
  • Stem: The stem uses a 3×3 convolution with stride 1 and padding 1 to maintain consistency with the top-level skip-connection.This avoids the output-resolution reduction associated with some ResNet- and ConvNeXt-style stems.
  • CMUNeXt Block: Depthwise separable convolution separates spatial extraction from channel mixing, reducing parameters and computational cost compared with ordinary convolution.The CMUNeXt block uses large-kernel depthwise convolution for global information and pointwise convolutions for mixing.
  • CMUNeXt Block: The inverted bottleneck expands the hidden dimension between pointwise convolutions to four times the input dimension for broader global spatial-information mixing.GELU activation and post-activation BatchNorm are also used in the block.
  • Channel Expansion: Because CMUNeXt blocks preserve resolution and channel size, an ordinary convolution block expands the channel size by two times.This channel expansion follows the CMUNeXt block.
  • Down-sampling: Pooling is presented as suitable for down-sampling medical images because they often have low resolution and minor local edge variations.The passage contrasts pooling with convolutional down-sampling for this setting.

Decoder Stage with skip-connections

The decoder contains five levels that pair Skip-Fusion with upsampling. Skip-Fusion uses grouped and pointwise convolutions, followed by GELU and BatchNorm, to combine encoder and decoder features efficiently.

  • Decoder Structure: The decoder consists of five bottom-to-top levels, each containing a Skip-Fusion block and an upsampling block.This establishes the repeated structure of the decoder stage.
  • Skip-Fusion: Skip-Fusion uses grouped convolution to perform feature-by-feature extraction on skipped encoder features and up-sampled decoder features.The grouped design addresses the burden of ordinary convolutional feature fusion.
  • Skip-Fusion: Each convolution in Skip-Fusion is followed by GELU activation and BatchNorm.These operations are part of the block definition described in the decoder section.
  • Feature Fusion: The fusion output is formed from encoder and decoder features, denoted fE and fD, respectively.The supplied definition identifies ffusion as the output fusion feature map.
  • Upsampling: The upsampling block uses bilinear interpolation by a factor of two, followed by 3×3 convolution, BatchNorm, and ReLU.This increases feature-map resolution while preserving important features.

Datasets

The study evaluates CMUNeXt on four ultrasound segmentation datasets, using varied clinical sources and repeated 70/30 train-validation splits to assess generalization.

  • BUS contains 562 breast ultrasound images from five devices, including 306 benign and 256 malignant cases with ground-truth labels.
  • BUSI contains 780 images from 600 female patients, covering normal, benign, and malignant cases; only benign and malignant cases are used.
  • TNSCUI includes 3,644 ultrasound cases across different ages and genders, each paired with ground-truth labels.
  • The private TUS dataset contains 1,942 thyroid ultrasound images from 192 cases, with labels produced by three experienced radiologists.
  • Each dataset is randomly split three times into 70% training and 30% validation sets to assess generalization across datasets and random splits.

Implementation Details

CMUNeXt is evaluated through variants that change channel counts, block lengths, and kernel sizes, using BCE-plus-Dice training under standardized augmentation and optimization settings.

  • CMUNeXt-S and CMUNeXt-L vary compute ratios, while CMUNeXt-S also uses a 9-sized bottom kernel to enlarge its receptive field.
  • The training loss combines binary cross entropy and Dice loss as L = 0.5 × BCE (ŷ, y) + Dice (ŷ, y).
  • Images are resized to 256×256 and augmented with random rotations and flips before training.
  • Training uses SGD with 1e-4 weight decay, 0.9 momentum, an initial learning rate of 0.01, batch size 8, and 300 epochs.
  • All experiments run on a single NVIDIA GeForce RTX4090 GPU using a poly learning-rate schedule.

Evaluation Metrics and Comparison Methods

The evaluation compares segmentation quality with model size, computational complexity, and inference speed across medical ultrasound datasets and competing networks.

  • The study evaluates IoU, F1-score, parameter count, GFLOPs, and FPS to compare segmentation performance and computational efficiency.
  • Table 2 reports comparison results between CMUNeXt and recent, widely used medical image segmentation networks on ultrasound datasets.
  • Figure 2 is included as a comparison chart within the evaluation of competing segmentation models.

Discussion

CMUNeXt combines strong segmentation performance with lower computational cost, faster inference, and lighter weights. Ablation, visualization, and training analyses attribute these results to large-kernel global context extraction and Skip-Fusion feature transfer.

  • CMUNeXt achieves state-of-the-art segmentation performance while maintaining reasonable computational complexity.
  • 6.03x fewer parameters, 2.24x fewer GFLOPs, and at least 2.49x higher FPS distinguish CMUNeXt-L from CMU-Net and TransUNet.
  • CMUNeXt-S improves IoU by 1.1-3.8% and F1 by 0.7-3.2% over UNeXt, using 0.41M parameters at 657 FPS.
  • Averaged across four datasets, CMUNeXt provides the best segmentation performance with faster inference, lighter weight, and lower GFLOPs than other models.
  • Ablation Study: Introducing the CMUNeXt block greatly improves segmentation while reducing parameters and GFLOPs and increasing FPS in the reduced U-Net.
  • Ablation Study: Adding Skip-Fusion further improves segmentation performance, demonstrating enhanced knowledge transfer capability.
  • Visualization of Results: CMUNeXt produces more accurate lesion localization and shape for multiple microcalcifications and low-contrast cases by learning global positional context.
  • Training Stability: CMUNeXt trains more stably and converges faster than U-Net and UNeXt, especially during early training.

Conclusion

CMUNeXt is a lightweight medical image segmentation network designed to balance segmentation performance and computation consumption. Its variants target different edge-device requirements and support fast, accurate auxiliary diagnosis in real-world scenarios.

  • CMUNeXt achieves an optimal trade-off between segmentation performance and computation consumption through refined network design.
  • Three CMUNeXt variants are tailored to deployment environments and performance requirements across different edge and device platforms.
  • The variants enable fast and accurate auxiliary diagnosis in various real-world scenarios.
  • CMUNeXt provides a promising solution with potential to improve diagnostic accuracy and efficiency in medical image segmentation.
Loading 2308.01239v2…