PaperlayerView exact revision on arXiv

Source-linked AI summary

UNeXt: MLP-based Rapid Medical Image Segmentation Network

Jeya Maria Jose Valanarasu, Vishal M. Patel

arXiv:2203.04967v1eess.IVcs.CV

TL;DR

Point-of-care imaging motivates segmentation networks that are less parameter-heavy, computationally complex, and slow than existing architectures. UNeXt addresses this gap with a convolutional encoder-decoder and latent tokenized MLP blocks with shifted features. Across multiple datasets, it achieves faster inference, lower complexity, fewer parameters, and state-of-the-art segmentation performance.

  • Problem

    Point-of-care applications require segmentation networks that address the large parameter counts, computational overhead, and slow inference of existing architectures.

  • Method

    UNeXt combines an initial convolutional stage with latent tokenized MLP blocks that shift and project convolutional features into tokens within an encoder-decoder architecture.

  • Results

    Across multiple medical-image segmentation datasets, UNeXt achieves faster inference, reduced complexity, fewer parameters, and state-of-the-art performance.

  • Takeaways & Limitations

    UNeXt provides a lightweight convolutional-MLP architecture for medical image segmentation focused on point-of-care applications.

Abstract

from arXiv · show

UNet and its latest extensions like TransUNet have been the leading medical image segmentation methods in recent years. However, these networks cannot be effectively adopted for rapid image segmentation in point-of-care applications as they are parameter-heavy, computationally complex and slow to use. To this end, we propose UNeXt which is a Convolutional multilayer perceptron (MLP) based network for image segmentation. We design UNeXt in an effective way with an early convolutional stage and a MLP stage in the latent stage. We propose a tokenized MLP block where we efficiently tokenize and project the convolutional features and use MLPs to model the representation. To further boost the performance, we propose shifting the channels of the inputs while feeding in to MLPs so as to focus on learning local dependencies. Using tokenized MLPs in latent space reduces the number of parameters and computational complexity while being able to result in a better representation to help segmentation. The network also consists of skip connections between various levels of encoder and decoder. We test UNeXt on multiple medical image segmentation datasets and show that we reduce the number of parameters by 72x, decrease the computational complexity by 68x, and improve the inference speed by 10x while also obtaining better segmentation performance over the state-of-the-art medical image segmentation architectures. Code is available at https://github.com/jeya-maria-jose/UNeXt-pytorch

1 Introduction

Point-of-care imaging increases the need for medical-image segmentation networks that are lightweight, computationally efficient, and fast. UNeXt addresses this need with a convolutional-MLP design that targets efficient segmentation while maintaining performance.

  • Motivation: UNet and transformer-based extensions have advanced medical image segmentation, but many methods give limited attention to computational complexity and inference time.The introduction identifies this gap alongside the large parameter counts that hinder point-of-care use.
  • Motivation: Point-of-care imaging enables bedside acquisition and faster analysis, increasing the importance of efficient segmentation networks.Examples include POCUS devices, phone-based skin-condition imaging, and bedside MRI.
  • Proposed approach: UNeXt combines convolutional networks and MLPs in a five-layer encoder-decoder architecture with skip connections.Its early and final blocks use convolutional processing, while the bottleneck uses tokenized MLP blocks.
  • Proposed approach: The tokenized MLP block projects convolutional features into tokens, while shifting operations extract local information along different axes.The design targets representation learning with reduced computation and parameter count.
  • Contributions: UNeXt improves medical-image segmentation performance while using fewer parameters, higher inference speed, and lower computational complexity than prior architectures.These outcomes are stated as the paper’s central contributions across evaluated medical segmentation tasks.

2 UNeXt

UNeXt combines convolutional encoder-decoder blocks with tokenized MLP blocks in latent stages, retaining skip connections and resolution changes across the network. Feature shifting, token projection, depth-wise convolution, and low-dimensional embedding are used to model representations efficiently while introducing locality and positional information.

  • Network Design: UNeXt uses a two-stage encoder-decoder: three convolutional encoder blocks, two tokenized MLP encoder blocks, two tokenized MLP decoder blocks, and three convolutional decoder blocks.Each encoder block halves feature resolution, each decoder block doubles it, and skip connections link corresponding levels.
  • Network Design: The architecture uses fewer channels than UNet variants, reducing parameters and computation.The standard channel configuration is C1 = 32, C2 = 64, C3 = 128, C4 = 160, and C5 = 256.
  • Shifted MLP: Shifting features across height and width before tokenization creates window locality along the corresponding axes.The operation sequentially shifts feature partitions to form random windows and introduce local information.
  • Tokenized MLP Stage: The tokenized MLP block shifts features, projects them into E-dimensional tokens, applies a shifted MLP, and then processes them with depth-wise convolution and GELU.Depth-wise convolution encodes positional information with fewer parameters than standard convolution, while GELU is used as a smoother activation than ReLU.
  • Tokenized MLP Stage: Tokenized MLP computations operate across an embedding dimension that is smaller than the feature-map dimensionality, limiting computational and parameter overhead.The experiments set the hidden dimension H to 768 unless otherwise stated.

3 Experiments and Results

Experiments benchmark UNeXt on ISIC and BUSI-related settings against convolutional and transformer baselines, emphasizing segmentation quality alongside efficiency for point-of-care use. UNeXt reports stronger segmentation performance with substantially lower computation and parameter count, while qualitative results remain competitive.

  • Datasets: UNeXt is evaluated on the ISIC 2018 and BUSI datasets, selected to approximate point-of-care imaging conditions.ISIC contains camera-acquired dermatologic images, while BUSI contains ultrasound images.
  • Performance Comparison: UNeXt obtains better segmentation performance than all baselines, with TransUNet the closest competitor and improvements statistically significant at p < 10^-5.The comparison includes convolutional and transformer-based frameworks.
  • Efficiency Comparison: 0.57 GFLOPs is UNeXt’s computational cost, compared with 38.52 for TransUNet and 55.84 for UNet.The paper attributes UNeXt’s lower computation relative to TransUNet to the absence of attention blocks.
  • Efficiency Comparison: 1.58 M parameters is UNeXt’s model size, compared with 105.32 M parameters for TransUNet.The paper describes UNeXt as the most lightweight network among the baselines.
  • Efficiency Comparison: Figure 4 compares F1 score against GFLOPs, inference time, and parameter count, showing UNeXt as the most efficient network.Higher F1 is better, while lower values are preferred for each x-axis efficiency measure.
  • Qualitative Results: Qualitative comparisons on ISIC and BUSI show that UNeXt produces competitive segmentation predictions relative to the other methods.The figure includes predictions from UNet, UNet++, MedT, TransUNet, and UNeXt alongside ground truth.

4 Discussion

The discussion examines how UNeXt’s modules and channel configurations trade segmentation performance against model size and computational overhead, and contrasts its design with MLP-Mixer.

  • Ablation Study: Tokenized MLPs substantially improve performance while adding only minimal complexity and parameters in the ablation study.Adding DWConv positional embedding yields further improvement, while feature shifting improves performance without adding parameters or FLOPs.
  • Ablation Study: Shifting features across both axes produces the best ablation performance without additional parameters or computational cost.The shift operation contributes no additions or multiplications and therefore adds no FLOPs.
  • Analysis on number of channels: Increasing UNeXt’s channels improves performance but increases computational overhead, whereas decreasing them produces a lightweight model with a modest performance reduction.The channel count affects parameters, complexity, and performance; UNeXt-S is specifically described as very lightweight.
  • Difference from MLP-Mixer: Unlike MLP-Mixer’s all-MLP image-recognition design, UNeXt combines convolutional features with shifted tokenized MLPs for image segmentation.An MLP-Mixer encoder with a convolutional decoder was still heavy at around 11 M parameters and produced non-optimal segmentation performance.

5 Conclusion

The conclusion presents UNeXt as a convolutional and MLP-based medical image segmentation architecture for point-of-care applications. Across multiple datasets, it achieves faster inference, lower complexity, fewer parameters, and state-of-the-art performance.

  • 5 Conclusion: UNeXt combines an initial convolutional stage with latent-space MLPs and uses shifted tokenized MLP blocks to model representations efficiently.The architecture is designed for medical image segmentation focused on point-of-care applications.
  • 5 Conclusion: Across multiple datasets, UNeXt achieves faster inference, reduced computational complexity, fewer parameters, and state-of-the-art segmentation performance.The conclusion reports these outcomes collectively rather than giving separate dataset-specific values.
Loading 2203.04967v1…