PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Ho Hin Lee, Shunxing Bao, Yuankai Huo, Bennett A. Landman

arXiv:2209.15076v4cs.CVcs.LG

TL;DR

Volumetric medical ViTs achieve strong segmentation performance, but their complexity motivates efficient alternatives. The paper introduces 3D UX-Net, a lightweight pure ConvNet using large-kernel depth-wise convolutions and transformer-inspired block designs. Across three public datasets and training settings, it consistently outperforms transformer SOTAs with fewer model parameters.

  • Problem

    3D volumetric ViTs provide strong segmentation performance, but adapting them as generic backbones involves high complexity and computationally unscalable self-attention.

  • Method

    3D UX-Net adapts hierarchical transformer behavior with pure ConvNet modules, using large-kernel volumetric depth-wise convolutions, pointwise depth convolution scaling, and reduced normalization and activation layers.

  • Results

    3D UX-Net consistently outperforms current transformer SOTAs with fewer model parameters across three public datasets in supervised training and transfer learning scenarios.

  • Takeaways & Limitations

    Pure ConvNet modules can adapt hierarchical transformer capabilities for volumetric segmentation while preserving competitive performance with fewer model parameters.

Abstract

from arXiv · show

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\times7\times7$) 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.

1 INTRODUCTION

The paper asks whether convolution modules can reproduce hierarchical transformer capabilities in volumetric segmentation while reducing complexity. It proposes 3D UX-Net, a lightweight pure ConvNet evaluated across direct training and transfer learning.

  • Large receptive fields and model capacity are identified as key contributors to transformer performance, but 3D ViTs remain difficult to adapt as generic backbones.
  • Volumetric depth-wise convolutions with large kernels simulate Swin Transformer receptive fields, while pointwise depth convolution widens channel representations with fewer parameters.
  • 3D UX-Net adapts hierarchical transformer behavior using purely convolutional modules for volumetric medical image segmentation.
  • The study evaluates 3D UX-Net on three public datasets using direct training and finetuning scenarios for volumetric multi-organ and tissue segmentation.

2 RELATED WORK

Related work shows that transformer-based volumetric segmentation achieves strong performance but carries computational and training costs. Prior depth-wise convolution studies motivate revisiting efficient ConvNet designs for robust segmentation.

  • SwinUNETR uses Swin Transformer encoder blocks and achieves state-of-the-art performance on several 3D segmentation benchmarks.
  • Figure 1 compares designed convolution blocks and volumetric ConvNet backbones with Swin Transformer block architecture.
  • Transformer-based volumetric frameworks require lengthy training and high computational complexity for multi-scale feature extraction.
  • Depth-wise convolution has been used in segmentation to reduce parameters, sharpen features, and support transfer learning.

3 3D UX-NET: INTUITION

3D UX-Net redesigns hierarchical-transformer-inspired encoder blocks with large-kernel depth-wise convolutions, patch projection, channel scaling, residual connections, and normalization choices.

  • Patch-wise Features Projection: 3D UX-Net uses a large-kernel projection layer to extract patch-wise features as encoder inputs instead of flattening patches for a linear layer.
  • Volumetric Depth-wise Convolution with LKs: Large-kernel volumetric depth-wise convolution provides per-channel feature extraction with a large receptive field analogous to Swin Transformer self-attention.
  • Inverted Bottleneck with Depthwise Convolutional Scaling: The inverted bottleneck widens hidden representations to four times the input dimension and uses 1 × 1 × 1 depthwise convolution for channel-wise scaling.
  • The encoder design adapts layer-wise differences between ConvNets and Swin Transformers to refine transformer-like behavior at the macro level.
  • Applying Residual Connections: Residual connections link input and extracted features after scaling without normalization or activation around the residual summation.

4 3D UX-NET: COMPLETE NETWORK DESCRIPTION

3D UX-Net is a U-shaped volumetric ConvNet that uses large-kernel depthwise convolutions to emulate transformer behavior while retaining hierarchical multi-scale features.

  • Architecture: 3D UX-Net processes 3D patches with multiple redesigned volumetric convolution blocks and connects encoder features to a convolution-based decoder through skip connections.The decoder predicts segmentation probabilities after residual processing and feature concatenation.
  • Depth-wise convolution encoder: Large-kernel depthwise convolutions starting from 7 × 7 × 7 provide partitioned feature extraction and are intended to supply a global receptive field.The design uses depthwise convolution as a computational alternative to global self-attention for volumetric features.
  • Depth-wise convolution encoder: 3D UX-Net replaces Swin Transformer’s W-MSA and SW-MSA modules with two depthwise convolution layers.Layer normalization and depthwise convolution scaling are retained in the encoder-block formulation.
  • Hierarchical encoder: The encoder has four stages with two large-kernel convolution blocks per stage, producing hierarchical multi-scale representations.The DCS layer expands feature dimensions to four times the input channel size, while downsampling progressively reduces spatial resolution.
  • Qualitative evaluation: Qualitative evaluation across three public datasets compares tissue and multi-organ segmentation, with boxed regions enlarged to show segmentation differences.The caption reports that 3D UX-Net shows the best segmentation quality compared with the ground truth.

5 EXPERIMENTAL SETUP

The experiments evaluate 3D UX-Net on three public volumetric segmentation datasets under direct supervised training, transfer learning, and architectural ablation settings.

  • Datasets: Experiments use FeTA2021, FLARE2021, and AMOS2022 public multi-modality datasets for volumetric segmentation.The datasets cover infant brain MRI and abdominal multi-organ segmentation tasks.
  • Evaluation protocol: FeTA2021 and FLARE2021 are evaluated with direct supervised training and five-fold cross-validation, while AMOS2022 is used for transfer learning.AMOS finetuning uses pretrained weights from the best FLARE2021 fold model.
  • Ablation studies: Ablation studies examine different architecture choices, kernel sizes, and replacing linear layers with depthwise convolution for feature extraction.The studies use FeTA2021 and FLARE2021 to assess component effects.

6 RESULTS

3D UX-Net achieves strong segmentation performance across direct-training and transfer-learning evaluations, with fast convergence and favorable qualitative results.

  • Direct training: 0.870 to 0.874 Dice on FeTA2021 and 0.929 to 0.934 Dice on FLARE2021 mark the reported improvements over compared methods.The paper describes 3D UX-Net as achieving the best performance across the evaluated segmentation tasks.
  • Convergence and quality: 3D UX-Net converges fastest on FeTA2021, while convergence becomes comparable to SwinUNETR as FLARE2021 training sample size increases.The reported qualitative results also describe preserved organ and tissue morphology relative to ground-truth labels.
  • Transfer learning: 0.900 mean Dice, a 2.27% enhancement, is reported for 3D UX-Net finetuning on AMOS2022.The paper also reports consistent improvement for most organ segmentations.
  • Kernel-size ablation: 13 × 13 × 13 kernels perform best on FLARE2021, whereas 7 × 7 × 7 kernels are optimal on FeTA2021.The paper attributes the FLARE2021 preference to a larger receptive field for neighboring abdominal organs and the FeTA2021 preference to localized infant brains.

7 DISCUSSION

The design emulates Swin Transformer behavior with pure ConvNet modules and uses inverted bottlenecks and pointwise depthwise convolutions to reduce parameters.

  • Design rationale: 3D UX-Net combines sliding-window-like depthwise convolution with an inverted bottleneck that widens computed feature channels.These components are identified as key perspectives underlying the improved design.
  • Feature scaling: Pointwise depthwise convolution replaces linear scaling for wider-channel features and slightly improves performance with fewer model parameters.The reported changes are FeTA2021: 0.872 to 0.874 and FLARE2021: 0.933 to 0.934.
  • Future work: The limited number of scaling layers in each encoder block may affect performance slightly, motivating future investigation of 3D linear-scaling scalability.The paper presents this as a future-work direction rather than a resolved design choice.

8 CONCLUSION

3D UX-Net adapts hierarchical-transformer capabilities using pure ConvNet modules for volumetric medical image segmentation, outperforming transformer SOTAs with fewer parameters across supervised and transfer-learning settings.

  • 3D UX-Net is a volumetric network that adapts hierarchical-transformer capabilities using pure ConvNet modules.Its encoder uses depthwise convolutions and projections to simulate hierarchical-transformer behavior.
  • The model improves segmentation performance across different training settings, including supervised training and transfer learning.
  • 3D UX-Net outperforms current transformer SOTAs with fewer model parameters on three challenging public datasets.

A.1 DATA PREPROCESSING & MODEL TRAINING

The training pipeline applies hierarchical intensity preprocessing, foreground sub-volume cropping, augmentation, and AdamW optimization across the evaluated datasets.

  • Data preprocessing: Intensity clipping and percentile-based min-max normalization preprocess each volume before training.For FLARE2021 and AMOS2022, clipping uses limits of -175 and 250; normalized intensities range from 0 to 1.
  • Data preprocessing: Training randomly crops 96 × 96 × 96 foreground sub-volumes and applies rotations, intensity shifting, and scaling augmentation.The stated scaling factor is 0.1.
  • Model training: The study uses AdamW optimization for 40000 training steps with a learning rate of 0.0001 on FeTA2021 and FLARE2021.Training uses NVIDIA-Quadro RTX 5000 for these datasets and RTX A6000 for AMOS2022.
  • Model training: Table 4 summarizes hyperparameters for direct training and finetuning scenarios across three public datasets.

A.2 PUBLIC DATASETS DETAILS

The supplied passages cover dataset evaluation, nn-UNet comparisons, efficiency optimization, and future architectural directions rather than a complete dataset-details section.

  • Comparisons and ablations: The study compares 3D UX-Net with CNN-based SOTA networks and nn-UNet on FeTA2021 and FLARE2021.The passages report ablation studies and mean organ Dice improvements when replacing plain 3D U-Net with 3D UX-Net: 0.874 to 0.881 on FeTA2021 and 0.934 to 0.944 on FLARE2021.
  • Efficiency optimization: 62.2M to 53.0M parameters compares 3D UX-Net with SwinUNETR, while architectural optimization reduces parameters from 53.0M to 32.1M.The optimized design also reduces FLOPs from 639.4G to 536.1G while preserving performance, as reported in Table 7.
  • Limitations and future work: The decoder adds a 3D ResNet block after transpose convolution, increasing efficiency costs and motivating further decoder-design investigation.The authors identify depthwise convolution as a potential solution for reducing this burden.
  • Limitations and future work: Re-parameterization, parallel branches, and stacks of small kernels are proposed as future directions for improving large-kernel training and inference efficiency.
Loading 2209.15076v4…