[Paper Review] 3D UX-Net: A Large Kernel Volumetric ConvNet Modernizing Hierarchical Transformer for Medical Image Segmentation

Ho Hin Lee, Shunxing Bao+2 more|arXiv (Cornell University)|Sep 29, 2022

Medical Imaging and Analysis99 citations

TL;DR

This paper presents 3D UX-Net, a pure ConvNet with large kernel depthwise convolutions that emulates hierarchical transformer behavior for volumetric medical image segmentation, achieving state-of-the-art results on public datasets with fewer parameters.

ABSTRACT

The recent 3D medical ViTs (e.g., SwinUNETR) achieve the state-of-the-art performances on several 3D volumetric data benchmarks, including 3D medical image segmentation. Hierarchical transformers (e.g., Swin Transformers) reintroduced several ConvNet priors and further enhanced the practical viability of adapting volumetric segmentation in 3D medical datasets. The effectiveness of hybrid approaches is largely credited to the large receptive field for non-local self-attention and the large number of model parameters. In this work, we propose a lightweight volumetric ConvNet, termed 3D UX-Net, which adapts the hierarchical transformer using ConvNet modules for robust volumetric segmentation. Specifically, we revisit volumetric depth-wise convolutions with large kernel size (e.g. starting from $7 imes7 imes7$) to enable the larger global receptive fields, inspired by Swin Transformer. We further substitute the multi-layer perceptron (MLP) in Swin Transformer blocks with pointwise depth convolutions and enhance model performances with fewer normalization and activation layers, thus reducing the number of model parameters. 3D UX-Net competes favorably with current SOTA transformers (e.g. SwinUNETR) using three challenging public datasets on volumetric brain and abdominal imaging: 1) MICCAI Challenge 2021 FLARE, 2) MICCAI Challenge 2021 FeTA, and 3) MICCAI Challenge 2022 AMOS. 3D UX-Net consistently outperforms SwinUNETR with improvement from 0.929 to 0.938 Dice (FLARE2021) and 0.867 to 0.874 Dice (Feta2021). We further evaluate the transfer learning capability of 3D UX-Net with AMOS2022 and demonstrates another improvement of $2.27\%$ Dice (from 0.880 to 0.900). The source code with our proposed model are available at https://github.com/MASILab/3DUX-Net.

Motivation & Objective

• Motivate the need for efficient 3D segmentation backbones that balance performance and model size.
• Propose a lightweight volumetric ConvNet that mimics hierarchical transformer behavior using large kernel depthwise convolutions.
• Reduce parameters and normalization requirements while preserving or improving segmentation accuracy.
• Showcase strong empirical results on public brain and abdominal datasets in supervised and transfer-learning settings.

Proposed method

• Introduce volumetric depthwise convolutions with large kernel sizes (LK) to simulate large receptive fields.
• Replace MLP in transformer blocks with pointwise depth convolutions to widen features with fewer parameters (DCS).
• Use an inverted bottleneck design with depthwise convolutions to expand and compress channel features across layers.
• Replace batch normalization with layer normalization in encoder blocks and use GELU activations.
• Construct a 4-stage encoder with two LK blocks per stage, followed by a ConvNet-based U-shaped decoder with skip connections.

Experimental results

Research questions

• RQ1 Can a pure ConvNet with large-kernel depthwise convolutions match or exceed transformer-based 3D segmentation performance?
• RQ2 Does the proposed 3D UX-Net achieve comparable or better accuracy with fewer parameters and reduced normalization layers?
• RQ3 How does 3D UX-Net perform in supervised training and transfer learning on public volumetric datasets?
• RQ4 What is the impact of kernel size and depthwise scaling (DCS) on segmentation performance across datasets?

Key findings

• 3D UX-Net outperforms SwinUNETR on FeTA2021 (0.874 vs 0.867 Dice) and FLARE2021 (0.934 vs 0.929 Dice).
• In transfer learning to AMOS2022, 3D UX-Net achieves Dice 0.900, a 2.27% improvement over the best transformer baseline.
• Ablation studies show kernel size and depthwise scaling influence performance, with LK sizes around 7x7x7 to 13x13x13 providing optimal gains per dataset.
• 3D UX-Net converges faster in FeTA2021 and demonstrates robust transfer learning behavior on AMOS2022.
• The model achieves competitive Dice scores with fewer parameters (53.0M) compared to several transformer baselines.