3D UX-Net: A Large Kernel Volumetric ConvNet Modernizing Hierarchical Transformer for Medical Image Segmentation
Authors: Ho Hin Lee
Shunxing Bao
Yuankai Huo
Bennett A. Landman
Venue: ICLR 2023 poster (2023)
Keywords: Depth-wise Convolution
Large Kernel Convolution
Convolutional Neural Network
Hierarchical Transformer
Volumetric Segmentation
Medical Image Segmentation

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. We hypothesize that volumetric ConvNets can simulate the large receptive field behavior of these learning approaches with fewer model parameters using depth-wise convolution. 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 (LK) 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.

Published as a conference paper at ICLR 2023
3D UX-NEt: A LARGE KERNEL VOLUMETRIC CON- VNET MODERNIZING HIERARCHICAL TRANSFORMER FOR MEDICAL IMAGE SEGMENTATION
Ho Hin Lee Vanderbilt University
Shunxing Bao Vanderbilt University
Yuankai Huo
Bennett A. Landman
Vanderbilt University
Vanderbilt University
ABSTRACT
The recent 3D medical VIT' (e.g.. SwinUNETR) achieve the state-of-the-art per- formances several 3D volumetric data benchmarks, including 3D medical image segmentation. Hierarchical transformers (c.g., Swin Transformers) rein- troduced several ConvNet priors and further enhanced the practical viability of adapting volumetric segmentation in 3D medical datasets. The effectiveness hybrid approaches is largely credited to the large receptive field for non-local self- attention and the large number of model parameters. We hypothesize that vol- umetric ConvNets can simulate the large receptive field behavior of these learn- ing approaches with fewer model parameters using depth-wise convolution. In this work, we propose lightweight volumetric ConvNet, termed 3D UX-Net. which adapts the hierarchical transformer using ConvNet modules for robust vol- umetric segmentation. Specifically, we revisit volumetric depth-wise convolu- tions with large kernel (LK) size (e.g. starting from 7 x 7 7) to enable the larger global receptive fields, inspired by Swin Transformer. We further substi- tute 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 (c.g. SwinUNETR) us- ing three challenging public datasets volumetric brain and abdominal imag- ing: 1) MICCAI Challenge 2021 FLARE, 2) MICCAI Challenge 2021 FeTA, and 3) MICCAI Challenge 2022 AMOS. 3D UX-Net consistently outperforms Swin- UNETR 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 A.MOS2022 and demonstrates another improvement of 2.27% Dice (from 0.880 O. 0.900). The source code with our proposed model are available httpp:// thub. com Coom/MASTIO MMtp/MMSSTb//////////a LADUUXXX 3DUX- Net.
INTRODUCTION
Significant progress has been made recently with the introduction of vision transformers (ViTs) Dosovitskiy et al. (2020) into 3D medical downstream tasks, especially for volumetric segmenta- tion benchmarks Wang et al. (202T): Hatamizadeh et al.] (202bb); Zhou et al. (2021): Xie et al. CChnn al. (2021). The characteristics VIT's are the lack of image-specific inductive bias and the scaling behaviour, which are enhanced by large model capacities and dataset sizes. Both characteristics contribute to the significant improvement compared to ConvNets on medical image segmentation Tang et al. (2022); [ao et al. (2021); He et al. (2022): Atito et al. (202I). However, it challenging to adapt 3D ViT models as generic network backbones due to the high complexity of computing global self-attention with respect to the input size, especially in high resolution images with dense features across scales. Therefore, hierarchical transformers are proposed to bridge these gaps with their intrinsic hybrid structure Zhang al. (2022): Liu al. (2021). Introducing the "slid-
"Correspondence ho. vanderbilltcduu

---

Published as a conference paper at ICLR 2023
ing window" strategy into ViTs termed Swin Transformer behave similarily with ConvNets Liu et al. (202T) SwinUNETR adapts Swin transformer blocks as the generic vision encoder backbone and achieves current state-of-the-eeeeee the performance several 3D segmentation benchmarks HHatamizadchi et al. (022a); [Tang et al. 2022). . Such performance gain is largely owing to the large receptive field from 3D shift window multi-head self-attention (MSA). However, the computation of shift window MSA is computational unscalable to achieve via traditional 3D volumetric ConvNet architectures. As the advancement of ViTs starts to bring back the concepts of convolution, the key components for such large performance differences are attributed to the scaling behavior and global self-attention with large receptive fields. As such, we further ask: Can we leverage convolution modules to enable the capabilities of hierarchical transformers?
The recent advance in I.K L-beeeed depthwise convolution design (e.g.. Liu al. Liu et al. 2022) provides a computationally scalable mechanism for large receptive field in 21 ConvNet. Inspired by such design, this study revisits the 3D volumetric ConvNet design to investigate the feasibility of (1) achieving the SOTA performance via pure ConvNet architecture, (2) yielding much less network complexity compared with 3D ViTs, and (3) providing a new direction of de- signing 3D ConvNet volumetric high resolution tasks. Unlike SwinUNETR, we propose lightweight volumetric ConvNet 3D UX-Net to adapt the intrinsic properties of Swin Transformer with ConvNet modules and enhance the volumetric segmentation performance with smaller model capacities. Specifically, we introduce volumetric depth-wise convolutions with L. sizes to simulate the operation of large receptive fields for generating self-attention in Swin transformer. Further- more, instead of linear scaling the self-attention feature across channels, we further introduce the pointwise depth convolution scaling distribute each channel-wise feature independently into wider hidden dimension (e.g., 4 x input channel), thus minimizing the redundancy of learned context across channels and preserving model performances without increasing model capacity. We evalu- ate 3D UX-Net supervised volumetric segmentation tasks with three public volumetric datasets: 1) MICCAI Challenge 2021 FeTA (infant brain imaging), 2) MICCAT Challenge 2021 FLARE (ab- dominal imaging), and 3) MICCAI Challenge 2022 AMOS (abdominal imaging). Surprisingly, 3D UX-Net, a network constructed purely from ConvNet modules, demonstrates a consistent improve- ment across all datasets comparing with current transformer SOTA. We summarize our contributions as below:
We propose the 3D UX-Net to adapt transformer behavior purely with ConvNet modules in volumetric setting. To our best knowledge, this is the first large kernel block design of leveraging 3D depthwise convolutions to compete favorably with transformer SOTAs in volumetric segmentation tasks.
We leverage depth-wise convolution with LK size the generic feature extraction back- bone, and introduce pointwise depth convolution to scale the extracted representations ef- feetively with less parameters.
We use three challenging public datasets to evaluate 3D UX-Net 1) direct training and 2) finetuning scenarios with volumetric multi-organ/tissues segmentation. 3D UX-Net achieves consistently improvement in both scenarios across all ConvNets and transform- ers SOTA with fewer model parameters.
RELATED WORK
2.1 TRANSFORMER BASED SEGMENTATION
Significant efforts have been put into integrating ViTs for dense predictions in medical imaging domain Hatamizadeh et al. (2022)); Chen et al. (202T): Zhou et al. (202I); [aag et al. (202T). With the advancement of Swin Transformer, SwinUNETR equips the encoder with the Swin Trans- former blocks to compute self-attention for enhancing brain tumor segmentation accuracy in 3D MRI Images Hatamizadeh et al. (2022a). Tang et al. extends the SwinUNETR by adding a self- supervised learning pre-training strategy for fine-tuning segmentation tasks. Another Unet-like ar- chitecture Swin-Unet further adapts Swin Transformer on both the encoder and decoder network via conneccinnttsnss to learn local and global semantic features for tion [Co et al. (202T). Similarly, SwinBTS has the similar intrinsic structure with Unett with an enhanced transformer module for detailed feature extraction Jiing et al. (2022). However, the
CT segmenta-
2

---

Published as a conference paper at ICLR 2023
Block Intrinsic
Successive Block Comparison
10 Sotal Doead
3r Owand
Laper Norm.
flrr.
Layer firrs
Eagar
30 Can Ixlal Cureed
Me Net poiet
BAK, 7ono
SANWAA Novt
+
2
VSL unanr
Lian Clare
Lapar Nore.
Laper Morm.
Layer Komm.
Leger Maurn
Chamml4
Usaus Durad's
MLP
DCs Lekel
MP
DCS txtel
Cured
GLU
4Bu
SEL
anu
grz
swth tundomuraaae alck
61 US sett Bick
Figure I: Overview of our proposed designed convolution blocks to simulate the behaviour of swin transformers. We leverage depthwise convolution and pointwise scaling adapt large receptive field and enrich the features through widening independent channels. We further compare different backbones of volumetric ConvNets and Swin Transformer block architecture. The yellow dotted line demonstrates the differences in spatial position of widening feature channels in the network bottleneck.
transformer-based volumetric segmentation frameworks still require lengthy training time and are accompanied by high computational complexity associated with extracting features at multi-scale levels Xie et al. (202I): [Shamshad et al. (2022). Therefore, such limitations motivate to rethink ConvNets can emulate transformer behavior to demonstrate efficient feature extraction.
2.2 DEPTHWISE CONVOLUTION BASED SEGMENTATION
Apart from transformer-be based framework, previous works began revisit the concept of depth- wise convolution and adapt its characteristics for robust segmentation. It has been proved to be a powerful variation of standard convolution that helps reduce the number of parameters and transfer learning Guo et al. (2019). Zunair et al. introduced depthwise convolution to sharpen the fea- tures prior fuse the decode features in UNet-like architecture [Zunair Hamzal (202T). 3D U--Net leveraged depthwise convolutions each channel Huang et al. (2019). Both studies demonstrate the feasibility of using depthwise convolution in enhancing volumetric tasks. However, only small kernel size is used to perform depthwise convolution. Several prior works have investigated the effectiveness of LK convolution in medical image segmentation. For instance, Huo et al. leveraged I.K (7x7) convolutional layers as the skip connections to address the anatomical variations for splenomegaly spleen segmentation Huo et al. (2018); Li al. proposed to adapt LK and dilated depthwise convolutions in decoder for volumetric segmentation Li et al. (2022). However, sigmificant increase of FLOPs is demon- strated with LKs and dramatically reduces both training and inference efficiency. To enhance the model efficiency with LKs, Liu et al. proposed ConvNeXt as 2D generic backbone that simulate ViTs advantages with LK depthwise convolution for downstream tasks with natural image Liu al. (2022), while Conv is proposed to extend for 2D medical image segmentation and compared only with 2D CNN-based networks (e.g.. ResUNet et al. 2021). UNet++ [Zhou al.) (2019) Han et al.] (2022). However, limited studies have been proposed to efficiently leverage depthwise convolution with LKs in a volumetric setting and compete favorably with volumetric transformer ap- proaches. With the large receptive field brought by LK depthwise convolution, we hypothesize that LK depthwise convolution can potentially emulate transformers' behavior and efficiently benefits for volumetric segmentation.
domain adaptors to extract domain-specifee features
3D UX-NET: INTUITION
Inspired by Liu et al. 2022 introduce 3D UX-Net, simple volumetric ConvNet that adapts the
capability of hierarchical transformers and preserves the advantages of using ConvNet modules such inductive biases. The basic idea of designing the encoder block in 3D UX-Net can be divided into
3

---

Published as a conference paper at ICLR 2023
Irput
Mulk Sagractztion
Layer Nerms
Lape barns
12"2 3
flor..
oestatal
033id
Bock
Stape4
Figure 2: Overview of the proposed 3D UX-Net with our designed convolutional block as the en- coder backbone. I.K convolution is used project features into patch-wise embeddings. A down- sampling block is used in each stage to mix and enrich context across all channels, while our de- signed blocks extract meaningful features in depth-wise setting.
1) bise and 2) layer-wise perspectives. First, we discuss the block-wise perspective in three views:
Patch-wise Features Projection: Comparing the similarities between ConvNets and ViTs, there is a common block that both networks use to aggressively downscale feature represen- tations into particular patch sizes. Here, instead of flattening image patches as sequential input with linear layer Dosovitskiy et al. (2020), we adopt L. projection layer to extract patch-wise features as the encoder's inputs.
Volumetric Depth-wise Convolution with LKs: One of the intrinsic specialties of the swin transformer is the sliding window strategy for computing non-local MSA. Overall, there are two hierarchical ways to compute MSA: window-based MSA (W-MSA) and 2) shifted window MSA (SW-MSA). Both ways generate global receptive field across layers and further refine the feature correspondence between non-overlapping windows. Inspired by the idea of depth-wise convolution, we have found similarities between the weighted sum approach in self-attention and the convolution per-channel basis. We argue that using depth-wise convolution with a LK size can provide a large receptive field in extracting fea- tures similar to the MSA blocks. Therefore, we propose compressing the window shifting characteristics of the Swin Transformer with volumetric depth-wise convolution using LK size (e.g., starting from x 7 x 7). Each kernel channel is convolved with the corre- sponding input channel, so that the output feature has the same channel dimension as the input.
Inverted Bottleneck with Depthwise Convolutional Scaling: Another intrinsic structure in Swin Transformer is that they are designed with the hidden dimension of the MLP block to be four times wider than the input dimension, as shown Figure 1. Such a design is in- terestingly correlated to the expansion ratio in the ResNet block et al. (2016). Therefore, we leverage the similar design in ResNet block and move up the depth- convolution to compute features. Furthermore,
introduce depthwise convolutional scaling (DCS) with x kernel to linearly scale each channel feature independently. We enrich the fea- ture representations by expanding and compressing each channel independently. thus min- imizing the redundancy cross-channel context. We enhance the cross-channel feature correspondences with the downsampling block in each stage. By using DCS, we further reduce the model complexity by 5% and demonstrates a comparable results with the block architecture using MLP.
The macro-design in convolution blocks demonstrates the possibility of adapting the large receptive field and leveraging similar operation of extracting features compared with the Swin Transformer. We want to further investigate the variation between ConvNets and the Swin Transformer in layer-
4

---

Published as a conference paper at ICLR 2023
wise settings and refine the model architecture better simulate ViTs in macro-level. Here, further define and adapt layer-wise differences into another three perspectives:
Applying Residual Connections: From Figure the golden standard 3D U-Net block demonstrates the naive approach of using small kernels to extract local representations with increased channels Çiçek et al. (2016). while the SegResNet block applies the residual sim- ilar to the transformer block My Myronenko (2018). Here, we also apply residual connections between the input and the extracted features after the last scaling layer. However, we do not apply any normalization and activation layers before and after the summation of residual to be equivalent with the swin transformer structure.
Adapting Layer Normalization (LN): In ConvNets, batch normalization (BN) is com- mon strategy that normalizes convolved representations to enhance convergence and reduce overfitting. However, previous works have demonstrated that BN can lead detrimental effect in model generalizability Wu & Johnson (2021). Although several approaches have (Saiimans
been proposed have an alternative normalization techniques Salimans & Kingma (2016); Ulyanov et al. (2016); Wu & He (2018), BN still remains as the optimal choice in volu- metric vision tasks. Motivated by vision transformers and Liu al. (2022). we directly substitute BN with LN in the encoder block and demonstrate similar operations in ViT's Ba
et
(016).
Many previous works have used the rectified lin- ear unit (ReL.U) activation layers Nair & Hinton (010), providing non-linearity in both ConvNets and ViTs. However, previously proposed transformer models demonstrate the Gaussian error linear unit (GELU) to be a smoother variant, which tackle the limitation of sudden zero in the negative input range in ReLU Hendrycks & Gimpel (2016). Therefore, we further substitute the ReLU with the GELU activation function.
Using GELU as the Activation Layer:
4 3D UX-NET: COMPLETE NETWORK DESCRIPTION
3D UX-Net comprises multiple re-designed volumetric convolution blocks that directly utilize 3D patches. Skip connections are further leveraged to connect the multi-resolution features to convolution-biie based decoder network. Figure illustrates the complete architecture of 3D UX-Net. We further describe the details of the encoder and decoder in this section.
4.1 DEPTH--ISS CONVOLUTION ENCODER
Given a set of 3D image volumes V/2 = X,, Yi.......... random sub-volumes P; E RRx are extracted be the inputs for the encoder network. Instead of flattening the patches and projecting with linear layer Hatamizadeh et al. (202bb), we leverage a LK convolutional layer to compute partitioned feature map with size z To adapt the characteristics of computing local self-attention, we use the depthwise convolution with kernel size starting from 7 7 x (DWC) with padding of 3, to act as a "shifted window" and evenly divide the feature map. As global self-attention is generally not computationally affordable with large number of patches extracted in the Swin Transformer [Lu et al.] (202T). we hypothesize that performing depthwise convolution with LK size can effectively extract features with global receptive field. Therefore, we define the output of encoder blocks layers and I ++ as follows:
x that are projected into C 48-dimensional space.
z (2 ') + 2
2' DCS(LN(3))) + 2!
(I)
*++1 DWC(LN((2)) +
,+1 = DC + g!+1
where & and ê1+1 are the outputs from the DWC layer in different depth levels; LN and DCS denote the layer normalization and the depthwise convolution scaling, respectively (see. Figure 1). Compared to the Swin Transformer, we substitute the regular and window partitioning multi- head self-attention modules, W-MSA and SW-MSA respectively, with two DWC layers.
Motivated by SwinUNETR Tang et al. (2022); Hatamizadeh et al. (2022a). the complete architecture of the encoder consists of 4 stages comprising of 2 I.K convolution blocks at each stage (i.e. I-=8
5

---

Published as a conference paper at ICLR 2023
Table I: Comparison of transformer and ConvNet SOTA approaches the Feta 2021 and FLARE 2021 testing dataset. (*: < 0.01, with Wilcoxon signed-rank test to all SOTA approaches)
PolA 2001
FLARR 2001
Methods
*aamss
FL.0P%
EXP
GM
W.3
Vest.
Cerch
DXM
815
Moxn
|Spleen
Kidaey
Liver
Pancreos
Moan
UlNaa Caakk
135.00 0G
0.925
0
0.910
0.907
0.892
SesReshoi
15.63
0.927
0.862
0
0.50%
0.520
0.925
Trans Wan atall
2021
31.54
110.4
0.585
0.778
0,552
0.561
0.513
0.576
0.877
0.868
0.964
0.555
0,974
0.711
0,902
UNET Pastinaan So Ai
2000h
92.5M
$2.60
0.561
0.262
0.927
0.862
0.508
0.558
083
0.850
0.927
0.947
0.886
0.2.0
0.885
unFonner Z-ri 2012
149.M
240.2G
0.890
0.770
0.330
0387
0.913
0.576
0328
0.863
0.933
0.900
0.975
0.217
0.905
SwisUNETs
3022
338.4G
0.929
0.869
0.914
0.367
0.065
0.980
0.788
0.929
UX-Net
0.534*
Table 2: Comparison of Finetuning performance with transformer SOTA approaches on the AMOS 2021 testing dataset.(* : p 0.01, with Wilcoxon signed-rank test to all SOTA approaches)
Methods
Spicen
R. Kid
1. Kid
Gall.
Psn.
Liver
Stom.
Aorta
IVC
Fanc:
RAG
1.4G
1300.
Blad.
Pros.
Avg
Dn-UNNe
0.959
0.889
0.950
0.948
0.739
0.869
0.839
TmnoBTS
0.885
0.931
0.916
0.817
0.744
0.969
0.837
0.914
0.855
0.724
0.630
0.566
0.704
0.741
0.650
0.792
UNETR
0.926
0.936
0.918
0.785
0,702
0.969
0.788
0.893
0.828
0.732
0.717
0.554
0.658
0.683
0.722
0.762
inFremer
0.935
03904
0.887
0.836
0.712
0.964
0.798
0.901
0.821
0.734
0.665
0.587
0.641
0.744
0.714
0. 3.790
SWODUNETR
0.959
0.960
0.894
0.979
0.745
0.817
0.875
0.841
0.880
UX-NNr
0.970
01967
0,923
0.984
0.739
0.906
0.876
a) FeTA 2021
b) FLARE 2021
AMOS 2022
0.8
wnl Munn
0.7
0.6
0.5
0.5
0.4
0.2
10k
13k
25k
30k
35K
330 Not JCNN)
SWWNUNETR (SwiT. E. CNND
UNETR (MTT ONN
:inFormer (VT)
TransBTS JCNN E.+ DJ1
Figure 3: Validation Curve with Dice Score for FeTA2021 (a), FLARE2021 (b) and AMOS2022 (c). 3D UX-Net demonstrates the fastest convergence rate with limited samples training (FeTA2021) and transfer learning (AMOS2022) scenario respectively, while the convergence rate is comparable to SwinUNETR with the increase of sample size training (FLARE2021).
total layers). Inside the block, the DCS layer is followed by the DWC layer in each block. The DCS layer helps scale the dimension of the feature map (4 times of the input channel size) without increasing model parameters and minimizes the redundancy of the learned volumetric context across channels. To exchange the information across channels, instead of using MLP, leverage stan- dard convolution block with kernel size 2 x 2 x with stride to downscale the feature resolution by factor of 2. The same procedure continues in stage 2, stage and stage 4 with the resolutions D respectively. Such hierarchical representations in x TO x To multi-scale setting extracted in each stage and are further leveraged for learning dense volumetric segmentation.
D x W of T x T x 4. g
W.
s x g and Tद
Hf W
4.2 DECODER
The multi-scale output from each stage in the encoder connected a ConvNet-based decoder via skip connections and form a "U-shaped" like network for downstream segmentation task. Specif- ically, we extract the output feature mapping of each stage i(i E 0, 1, 2, 3,4) in the encoder and further leverage a residual block comprising two post-normalized 3 x 3 x 3 convolutional layers with instance normalization to stabilize the extracted features. The processed features from each stage are then upsampled with transpose convolutional layer and concatentated with the features from the preceding stage. For downstream volumetric segmentation, we also concatenate the resid- ual features from the input patches with the upsampled features and input the features into residual

---

Published as a conference paper at ICLR 2023
UNETR
TransBTS
FeTA
FLARE 2021
AMOS 2022
Figure 4: Qualitative representations of tissues and multi-organ segmentation across three public datasets. Boxed further zoomed in and visualize the significant differences in segmentation quality. 3D UX-Net shows the best segmentation quality compared to the ground-truth.
block with X 1 convolutional layer with a softmax activation to predict the segmentation probabilities.
EXPERIMENTAL SETUP
Datasets We conduct experiments on three public multi-modality datasets for volumetric segmen- tation, which comprising with 1) MICCAI 2021 FeTA Challenge dataset (FeTT2221) Payette et al. 202T 2) MICCAI 2021 FLARE Challenge dataset (FLARE2021) Ma al. (202T). and 3) MIC- CAI 2022 AMOS Challenge dataset (AMOS2022) Ji et al. (022). For the FETA2021 dataset, employ 80 T2-weighted infant brain MRIs from the University Children's Hospital with 1.5 and 3T clinical whole-body scanners for brain tissue segmentation, with seven specific tissues well- annotated. For FLARÉ2021 and AMOS2022, we employ 511 multi-contrast abdominal CT from FLARE2001 with four anatomies manually annotated and 200 multi-contrast abdominal CT from AMOS 2022 with sixteen anatomies manually annotated for abdominal multi-organ segmentation. More details of the three public datasets can be found in appendix A.2.
Implementation Details We perform evaluations on two scenarios: 1) direct supervised training and 2) transfer learning with pretrained weights. FeTA2021 and FLARE2021 datasets are leverage to evaulate in direct training scenario, while AMOS dataset used in transfer learning scenario. We perform five-fold cross-validations both FeTA2021 and FLARE2021 datasets. More detailed information of data splits are provided in Appendix A.2. For the transfer learning scenario, we leverage the pretrained weights from the best fold model trained with FLARE2021, and finetune the model weights on AMOS2022 to evaluate the fine-tuning capability of 3D UX-Net. The complete preprocessing and training details are available at the appendix A. Overall, we evaluate 3D UX- Net performance by comparing with current volumetric transformer and ConvNet SOTA approaches for volumetric segmentation in fully- supervised setting. We use the Dice similarity coefficient as an evaluation metric to compare the overlapping regions between predictions and ground- truth labels. Furthermore, we performed ablation studies to investigate the effect on different kernel size and the variability of substituting linear layers with depthwise convolution for feature extraction.

---

Published as a conference paper at ICLR 2023
Table 3: Ablation studies of different architecture FeTA2021 and FLARE2021
FeTA2021
FLARE2021
Methods
#Params (M)
Mean Dice
SwinUNETR
62.2
0.867
0,929
Use Standard Conv.
186.9
0.875
0,937
Use Depth Conv.
53.0
0.874
0.934
Kemel=33 x x 3
52.5
0.867
0.928
Kernel=5 x X 5
52.7
0.869
0,931
Kernel=7 X X 7
53.0
0.874
0,934
Kemne=99 x x 9
53.6
0.870
0,934
Kerne==11 11
54.4
0.871
0,936
Kernel=13 x 13 x 13
55.7
0.871
0.938
No MLP
51.1
0.869
0,915
Use MLP
56.3
0.872
0,933
Use DCS 1x1x1
53.0
0.874
0.934
6 RESULTS
6.1 EVALUATION ON FETA & FLARE
Table shows the result comparison of current transformers and ConvNets SOTA medical im- age segmentation in volumetric setting. With our designed convolutional blocks as the encoder backbone, 3D UX-Net demonstrates the best performance across all segmentation task with signif- icant improvement in Dice score (FeTA2021: 0.870 to 0.874, FLARE2021: 0.929 to 0.934). From Figure 2, we observe that 3D UX-Net demonstrates the quickest convergence rate in training with FeTA2021 datasets. Interestingly, when the training sample size increases, the efficiency of training convergence starts become compatible between SwinUNETR and 3D UX-Net. Apart from the quantitative representations, Figure 3 further provides additional confidence of demonstrating the quality improvement segmentation with 3D UX-Net. The morphology of organs and tissues are well preserved compared to the ground-truth label.
6.2 TRANSFER LEARNING WITH AMOS
Apart from training from scratch scenario, we further investigate the transfer learning capability of 3D UX-Net comparing to the transformers SOTA with AMOS 2022 dataset. We observe that the finetuning performance
3D UX-Net significantly outperforms other transformer network with mean Dice of 0.900 (2.27% enhancement) and most of the organs segmentation demonstrate consistent improvement in quality. Also, from Figure 2, although the convergence curve of each transformer network shows the comparability to that of the FLARE2021-tran -trrinned model, 3D UX-Net further shows its capability in adapting fast convergence and enhancing the robustness of the model with finetuning. Furthermore, the qualitative representations in Figure 3 demonstrates significant improvement in preserving boundaries between neighboring organs and minimize the possibility of sver-segmentttion towards other organ regions.
6.3
ABLATION ANALYSTS
After evaluating the core performance of 3D UX-Net, we study how the different components in our designed architecture contribute such significant improvement in performance, as well as how they interact with other components. Here, both FeTA2021 and FLARE2021 are leveraged to perform ablation studies towards different modules. All ablation studies are performed with kernel size 7 x 7 7 scenario except the study of evaluating the variability of kernel size.
Comparing with Standard Convolution: We investigate the effectiveness of both standard convo- lution and depthwise convolution for initial feature extraction. With the of standard convolution, demonstrates a slight improvement with standard convolution. However, the model parameters about 3.5 times than that of using depthwise convolution, while the segmentation performance
8

---

Published as a conference paper at ICLR 2023
with depthwise convolution still demonstrates a comparable performance in both datasets. Variation of Kernel Size: From Table 3, we observe that the convolution with kernel size 7 x 7 optimally works for FeTA2021 dataset, while the segmentation performance of FLARE2021 demon- strates the best with kernel size of 13 x 13 x 13. The significant improvement of using 13 x 13 x 13 kernel for FLARE2021 may be due the larger receptive field provided to enhance the feature correspondence between multiple neighboring organs within the abdominal region. For FeTA2001 dataset, only the small infant brains are well localized as foreground and x 7 kernel demon- strates to be optimal recpetive field to extract the tissues correspondence. Adapting DCS: We found that a significant decrement is performed without using MLP for feature scaling. With the linear scaling, the performance enhanced significantly in FLARE2021, while slight improvement
demonstrated in FeTA2021. Interestingly, leveraging depthwise convolution with 1 x 1 x 1 kernel size for scaling, demonstrates slightly enhancement in performance for both FeTA2021 and FLARE2021 datasets. Also, the model parameters further drops from 56.3M to 53.0M without trading off the model performance.
7 DISCUSSION
In this work, we present a block-wise design to simulate the behavior of Swin Transformer using pure ConvNet modules. We further adapt our design generic encoder backbone into "U-Net" like architecture via skip connections for volumetric segmentation. We found that the key compo- be divided into main perspectives: 1) the sliding window
nents for improved performance strategy of computing MSA and 2) the inverted bottleneck architecture of widening the computed feature channels. The W-MSA enhances learning the feature correspondence within each window, while the SW-MSA strengthens the cross-window connections at the feature level between differ- ent non-overlapping windows. Such strategy integrates ConvNet priors into transformer networks and enlarge receptive fields for feature extraction. However, we found that the depth convolutions can demonstrate similar operations of computing MSA in Swin Transformer blocks. In depth- -iie convolutions, we convolve each input channel with a single convolutional filter and stack the con- volved outputs together, which is comparable to the patch merging layer for feature outputs in swin transformers. Furthermore, adapting the depth convolutions with LK filters demonstrates similarities with both W-MSA and SW-MSA, which learns the feature connections within a large receptive field. Our design provides similar capabilities to Swin Transformer and additionally has the advantage of reducing the number of model parameters using ConvNet modules.
Another interesting difference is the inverted bottleneck architecture. Figure shows that both Swin Transformer and some standard ConvNets have their specific bottleneck architectures (yellow dotted line). The distinctive component in swin transformer's bottleneck is maintain the channels size as four times wider than the input dimension and the spatial position of the MSA layer. We follow the inverted bottleneck architecture in Swin Transformer block and move the depthwise convolution to the top similar to the MSA layer. Instead of using linear scaling, we introduce the idea of depthwise convolution in pointwise setting scale the dense feature with wider channels. Interestingly, we found a slight improvement in performance is shown across datasets (FeTA2221: 0.872 to 0.874, FLARE2021: 0.933 to 0.934), but with less model parameters. As each encoder block only consists of two scaling layers, the limited number of scaling blocks may affect the performance to small extent. We will further investigate the scalability of linear scaling layer in 3D as the future work.
8 CONCLUSION
We introduce 3D UX-Net, the first volumetric network adapting the capabilities of hierarchical trans- former with pure ConvNet modules for medical image segmentation. We re-design the encoder blocks with depthwise convolution and projections to simulate the behavior of hierarchical trans- former. Furthermore, we adjust wiiee-wiee&e design in the encoder block and enhance the segmentation performance across different training settings. 3D UX-Net outperforms current transformer SOTAs with fewer model parameters using three challenging public datasets in both supervised training and transfer learning scenarios.

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGMENTS
This research supported by NIH Common Fund and National Institute of Diabetes, Digestive and Kidney Diseases U (Spraggins), NSF CAREER 1452485, NIH 2R01EB0061 NIH IR01EB017230 (Landman), and NIH ROINS09529. This study was in part using the resources of the Advanced Computing Center for Research and Education (ACCRE) at Vanderbilt University, Nashville, TN. The identified datasets used for the analysis described were obtained from the Re- search Derivative (RD), database of clinical and related data. The imaging dataset(s) used for the analysis described were obtained from ImageVU, a research repository of medical imaging data and image-related metadata. Image VU and RD are supported by the VICTR CTSA award (ULTR000445 from NCATS/NIH) and Vanderbilt University Medical Center institutional funding. ImageVU pilot work was also funded by PCORI (contract CDRN-1306-04866 We further thank Quan Liu, Ph.D student in Computer Science Department of Vanderbilt University, to extensively discuss the initial idea of this paper.
REFERENCES
Sara Atito, Muhammad Awais, and Josef Kittler. Sit: Self-supervised vision transformer. arXiv preprint arXiv:2 (0t::100443602,, 2021.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv: /60.06640. 2016.
Hangho Bao, Li Dong, and Furu Wei. Beit: Bert pre-training of image transformers. arXiv preprint (0Xiv:2/0600825: 2021.
Hu Cao, Yueyue Wang, Joy Chen, Dongsheng Jiang, Xiaopeng Zhang, Qi Tian, and Manning Wang. Swin-unet: Unet-like pure transformer for medical image segmentation. arXiv preprini arXiv:2105.05537, 2021.
Jieneng Chen, Yongyi Lu, Qihang Yu, Xiangde Luo, Ehsan Adeli, Yan Wang, Le Lu, Alan I. Yuille, and Yuyin Zhou. Transunet: Transformers make strong encoders for medical image segmentation. arXiv preprint arXiv:2102.04300 2021.
Özgün Çiçek, Ahmed Abdulkadir, Soeren S Lienkamp. Thomas Brox, and Olaf Ronneberger. 3d u- net: learning dense volumetric segmentation from sparse annotation. In International conference on medical image computing and
intervention, pp. 424 432. Springer, 2016.
Xiaohan Ding, Xiangyu Zhang, Ningning Ma. Jungong Han, Guiguang Ding, and Jian Sun. Repvgg: Making vgg-style convnets great again. In Proceedings of the IEEE/CVF Conference Com- puter Vision and Pattern Recognition, pp. 17733-13742, 2021.
Xiaohan Ding, Honghao Chen, Xiangyu Zhang, Kaiqi Huang, Jungong Han, and Guiguang Ding. Re-parameterizing your optimizers rather than architectures. arXiv preprint arXiv:2205. 15242, 2022a.
Xiaohan Ding, Xiangyu Zhang, Jungong Han, and Guiguang Ding. Scaling up your kernels to 31x31: Revisiting large kemel design in cnns. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 11963-11975, 2022b.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition scale. arXiv preprint arXiv:2010.11929 2020.
Yunhui Guo, Yandong Li, Liqiang Wang, and Tajana Rosing. Depthwise convolution is all you need for learning multiple visual domains. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pp. 8368-8375, 2019.
Zhimeng Han, Muwei Jian, and Gai-Ge Wang. Convunext: An efficient convolution neural network for medical image segmentation. Knowledge-Basee Systems, pp. 109512, 2022.
10

---

Published as a conference paper at ICLR 2023
Ali Hatamizadeh, Vishwesh Nath, Yucheng Tang, Dong Yang, Holger R Roth, and Daguang Xu. Swin unetr: Swin transformers for semantic segmentation of brain tumors in mri images. In International MICCAI Brainlesion Workshop, pp. 272-284. Springer, 2022a.
Ali Hatamizadeh, Yucheng Tang, Vishwesh Nath, Dong Yang, Andriy Myronenko, Bennett Land- man, Holger Roth, and Daguang Xu. Unetr: Transformers for 3d medical image segmentation. In Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision, pp- 574-584, 2022b.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked au- toencoders are scalable vision learners. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 16000-16009, 2022.
Dan Hendrycks and Kevin Gimpel. arXiv: 1ee.166660111111 2016.
Gaussian error linear units (gelus). arXiv preprint
Mu Hu, Junyi Feng, Jiashen Hua, Baisheng Lai, Jianqiang Huang, Xiaojin Gong, and Xian-Sheng Hua. Online convolutional pr-aammtterrzaaion. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp- 568-577, 2022.
Chao Huang, Hu Han, Qingsong Yao, Shankuan Zhu, and Kevin Zhou. 3d nne A 3d universal u-net for multi-domain medical image segmentation. In International Conference on Medical Image Computing and Computer-Assiste Intervention, pp. 291-299. Springer, 2019.
Yuankai Huo, Zhoubing Xu, Shunxing Bao, Camilo Bermudez, Hyconsoo Moon, Prasanna Par- vathaneni, Tamara Moyo, Michael R Savona, Albert Assad, Richard G Abramson, et al. Splenomegaly segmentation multi-modal mri using deep convolutional networks. IEEE trans- actions on medical imaging, 38(5): 185-11181111111 1196, 2018
Fabian Isensee, Paul F Jaeger, Simon AA Kohl, Jens Petersen, and Klaus H Maier-Hein. nnu- net: self-configuring method for deep learning baaning-based biomedical image segmentation. Nature methods, 18(2):203-211 2021.
Yuanfeng Ji, Haotian Bai, Jie Yang, Chongjian Ge. Ye Zhu, Ruimao Zhang, Zhen Li, Lingyan Zhang, Wanling Ma, Xiang Wan, et al. Amos: large-scale abdominal multi-organ benchmark for versatile medical image segmentation. arXiv preprint arXiv:2206.0880 2022.
Yun Jiang, Yuan Zhang, Xin Lin, Jinkun Dong, Tongtong Cheng, and Jing Liang. Swinbts: a method for 3d multimodal brain tumor segmentation using swin transformer. Brain Sciences, 12(6):797, 2022.
Ho Hin Lee, Yucheng Tang, Shunxing Bao, Richard G Abramson, Yuankai Huo, and Bennett Landman. Rap-net: Coarse-to-fine multi-organ segmentation with single random anatomical prior. In 2021 IEEE 18th International Symposium on Biomedical Imaging (ISBI), pp. 1491- 1494. IEEE, 2021.
Hao Li, Yang Nan, and Guang Yang. Lkau-net: 3d large-kernel attention- u-net for automatic mri brain tumor segmentation. In Annual Conference on Medical Image Understanding and Analysis, pp. 313-327. Springer, 2022.
Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 10012-10022, 2021.
Zhuang Liu, Hanzi Mao, Chao- Yuan Wu, Christoph Feichtenhofer. Trevor Darrell, and Saining Xie. A convnet for the 2020s. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 1 11976-1 1986, 2022.

---

Published as a conference paper at ICLR 2023
Jun Ma, Yao Zhang, Song Gu, Cheng Zhu, Cheng Ge, Yichi Zhang, Xingle An, Congcong Wang, Qiyuan Wang, Xin Liu, et al. Abdomenct- 1k: Is abdominal organ segmentation a solved problem. IEEE Transactions Pattern Analysis and Machine Intelligence, 2021.
Andriy Myronenko. 3d mri brain tumor segmentation using autoencoder regularization. In Interna- tional MICCAI Brainlesion Workshop, pp. 311 311-320. Springer, 2018.
Vinod Nair and Geoffrey E Hinton. Rectified linear units improve restricted boltzmann machines. In Leml, 2010.
Kelly Payette, Priscille de Dumast, Hamza Kebiri, Ivan Ezhov, Johannes C Paetzold, Suprosanna Shit, Asim Iqbal, Romesa Khan, Raimund Kottke, Patrice Grehten, al. An automatic multi- tissue human fetal brain segmentation benchmark using the fetal tissue annotation dataset. Scien- rific Data, 8(1):1-14, 2021.
Tim Salimans and Durk Kingma, Weight normalization: simple reparameterization accelerate training of deep neural networks. Advances in neural information processing systems, 29, 2016.
Fahad Shamshad, Salman Khan, Syed Waqas Zamir, Muhammad Haris Khan, Munawar Hayat, Fahad Shahbaz Khan, and Huazhu Fu, Transformers in medical imaging: survey, arXiv preprint arXiv:2201. 09873,, 2022.
Xiu Shu, Yunyun Yang, and Boying Wu. Adaptive segmentation model for liver et images based on neural network and level set method. Neurocomputingg 453:438-452, 2021.
Yucheng Tang, Dong Yang, Wenqi Li, Holger R Roth, Bennett Landman, Daguang Xu, Vishwesh Nath, and Ali Hatamizadeh. Sele-vupd pre-training of swin transformers for 3d medical image analysis. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 20730-20740, 2022.
Dmitry Ulyanov, Andrea Vedaldi, and Victor Lempitsky. Instance normalization: The missing in- gredient for fast stylization. arXiv preprint arXiv:1607.08022 2016.
Wenxuan Wang, Chen Chen, Meng Ding, Hong Yu, Sen Zha, and Jiangyun Li. Transbts: Mul- timodal brain tumor segmentation using transformer. In International Conference on Medical Image Computing and Computer-Assisted Intervention, pp. 109-119. Springer, 2021.
Yuxin Wu and Kaiming He. Group normalization. In Proceedings of the European conference on computer vision (ECCV), pp. 3-19, 2018
Yuxin Wu and Justin Johnson. Rethinking" batch" in batchnorm. arXiv preprint arX.v:2105.07576. 2021.
Yutong Xie, Jianpeng Zhang, Chunhua Shen, and Yong Xia. Cotr: Efficiently bridging cnn and
transformer for 3d medical image segmentation. In International conference on medical image computing and computer-assisted intervention, pp. 171-180. Springer, 2021.
Zizhao Zhang, Han Zhang, Long Zhao, Ting Chen, Sercan Ö Arik, and Tomas Pfister. Nested
hierarchical transformer: Towards accurate, data-efficient and interpretable visual understanding. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pp. 3417-3425, 2022.
Hong- Yu Zhou, Jiansen Guo, Yinghao Zhang, Lequan Yu, Liansheng Wang, and Yizhou Yu. nn- former: Interleaved transformer for volumetric segmentation. arXiv preprint 21X 2021.
Zongwei Zhou, Md Mahfuzur Rahman Siddiquee, Nima Tajbakhsh, and Jianming Liang. Unet++: Redesigning skip connections exploit multiscale features image segmentation. IEEE trans- actions medical imaging, 39(6): 118856666777 2019.
Hasib Zunair and A Ben Hamza. Sharp u-net: depthwise convolutional network for biomedical image segmentation. Computers in Biology and Medicine, 136: 2021.
12

---

Published as a conference paper at ICLR 2023
A APPENDIX
A.1 DATA PREPROCESSING & MODEL TRAINING
We apply hierarchical steps for data preprocessing: 1) intensity clipping is applied to further hance the contrast of soft tissue (FLARE2021 & AM0S2022.(min-1133 max:250)). 2) Inten- sity normalization is performed after clipping for each volume and use min-max normalization: (X X X1) to normalize the intensity value between 0 and I, where Xp denote as the Pth percentile of intensity in X. We then randomly crop sub-volumes with size 96 x 96 96 at the foreground and perform data augmentations, including rotations, intensity shifting, and scaling (scaling factor: 0.1). All training processes with 3D UX-Net are optimized with an Adam opti- mizer. We trained all models for 40000 steps using a learning rate of 0.0001 on an AVIDIA-Quadry RTX 5000 for both FeTA2021 and FLARR2021, while we perform training for AMOS2022 using NVIDIA-Quadny RTX A6000. One epoch takes approximately about minute for FeTA2021, 10 minutes for FLARE2021, and 7 minutes for AMOS2022, respectively. We further summarize all the training parameters with Table 4.
Table 4: Hyperparameters of both directly training and finetuning scenarios on three public datasets
Hyperparameters
Direct Training
Finetuning
Encoder Stage
4
Layer-wise wise Channel
48, 96, 194,384
Hidden Dimensions
768
Patch Size
96 96 x 96
No. of Sub-volumes Cropped
2
Training Steps
40000
Batch Size
2
AdamW e
1c 8
AdamW 3
(0.9, 0.999)
Peak Learning Rate
1e
Learning Rate Scheduler
N/A N/A
Factor & Patience
0.9, 10
Dropout
X
Weight Decay
0.08
Data Augmentation
Intensity Shift, Rotation, Scaling V 0.1 -30 to +30 x: 0.1, y: 0.1. 2: 0.1
Cropped Foreground
Intensity Offiet
Rotation Degree
Scaling Factor
A.2 PUBLIC DATASETS DETAILS
Table 5: Complete Overview of three public MICCAI Chanllenge Datasets
MROCAAChallerg
IrTA 2001
HARE 3P1
AMOS 2002
Imaging Modalit
Mala- Contaa CT
Malii Cont-CCttaaa
Dinnsioss
156
256
519 (37 (0.61 0.95) (0.61 -0 (us) -300 361
512 nos 768 353)
Renattion
03 0.70) x (0 63 0.70) (0:43 0.70)
(c 15 - 1.07] 10.45 1.07] :3.05 - 5.00)
Sample Size
80
200
Centhnopiral
Spleen, Right Kidayy,
Asstomical Lobel
Gry Mane (GMM) Mener (WM) Verrictes. Carisellam, Deep Okty Mater (DGM) Statassem
Spleen, Kidaey. Laver, Pancreas
Liver, Pazcreas L:fi & Right Alteral Gland (AG), Dawdesim Blzaks, Pootacaloorn
Infecor Veua
Spliss
5-Rod
Fild
Sdidatoe
A.3 FURTHER DISCUSSIONS COMPARING TO NN-UNET
In Table 1 & 2, compare proposed network with multiple CNN-based SOTA networks and the golden standard approach nn- We observe that the performance of On-UNet nearly outper- form most of the transformer state-of-the-arts in both FeTA 2021 and FLARE 2021 datasets. Such
13

---

Published as a conference paper at ICLR 2023
Table 6: Albation Studies of Adapting nn-UNet architecture the Feta 2021 and FLARE 2021 testing dataset. p 0.01, with Wilcoxon signed-rank test to all SOTA approaches, D.S: Deep Supervision)
FeTA 2021
FLARE 2021
Methods
BCF
GM
WM
Vent.
Coreb.
DGM
BS5
Meca
Spleen
Kidney
Liver
Pancreas
Moan
un-UNef Turnsco NurrnNeeeee
2021
0.883
0.775
0.930
0.868
0.920
0.880
0.840
0,570
0.971
0.966
0.976
0.792
0.926
3D UN- UN-Nerrrin
08822
0.780
0.924
0.872
0.917
0.886
0.845
0.874
0.981
0.969
0.982
0.801
0.934
UX-Net (un-UNett struet,, wio D.S.)
0.872
0.983
0.972
0,983
0.821
(mn- Lun-UNNet struct.. D.S.)
0.939
0.974
improvement may mainly contribute to its innovation of coffigoofignrttoonn training strategies and en- sembling outputs as postprocessing technique, while the network used in nn-UNet is only the plain 3D U-Net architecture. To further characterize the ability our proposed network, we further sub- stitute the plain 3D U-Net architecture with our proposed 3D UX-Net and adapt the self-configuring hyperparameters for training. We demonstrate a significant improvement of performance in FeTA 2021 and FLARE 2021 datasets with mean organ Dice from 0.874 0.881 and from 0.934 to 0.944 respectively, as shown Table 6. To further investigate the difference in the network architecture, observed that the convolution blocks in nn- un-UNet leverage the combination of instance normaliza- tion and leakyReLU.. Such design allows to normalize channel-wise feature independently and mix the channel context with small kernel convolutional layers. our design, we provide alternative thought of extracting channel-wise features independently with depthwise convolution and mix the channel information during the downsampling layer only. Therefore, layer normalization is lever- aged in our scenario and we want to further enhance the feature correspondence with large receptive field across channels efficiently. Furthermore, we found that the deep supervision strategy on- UNet, which compute an auxilary loss with each stages' intermediate output, also demonstrates its effectiveness to further improve the performance (FeTA 2021: from 0.876 to 0.881; FLARE 2021: from 0.940 to 0.944).
For the training scenarios, instead of using the proposed initial learning rate 0.01, we reduce the initial learning rate to 0.002 train with 150 epochs (40000 steps x 150 epochs) for FLARE 2021 and 850 epochs (40000 steps =e 850 epochs) for FeTA 2021 respectively, with the batch size of 2. For the finetuning scenario with AMOS 2022, we only train the nn-UNet model with 250 epochs (40000 steps R 250 epochs), instead of the default settings (1000 epochs) to ensure the fair network comparison with similar steps.
A.4 FURTHER DISCUSSIONS ON TRAINING AND INFERENCE EFFICIENCY
Table Albation Studies of Optimizing 3D U-XNet architecture the Feta 2021 and FLARE 2021 testing dataset. (SD: Stage Depth, HDim: Hidden Dimension in the Bottleneck Layer.)
FeTA2021 FLARE2021
Methods
#Params (M)
FLOPs (G)
Mean Dice
nn-UNet SwinUNETR
31.2M
743.3G
0.870
0.926
62.2M
328.4G
0.867
0.929
SD: 2,2,2, HDim: 768
53.0M
639.4G
0.874
0.934
SD: 2,2,8,2, HDim: 384
32.1M
536.8G
0.873
0.932
Apart from the advantage of quantitative performance, we further leverage the LK depthwise con- volutions reduce the model parameters from 62. 22.2M to 53.0M, compared to SwinUNETR Table 3. However, although the training efficiency of 3D UX-Net is already better than nn- nn-UNet (FLOPs: 743.3G to 639.4G), we observed that the FLOPs of 3D UX-Net still remains at a high value. In- spired by the architectures of both Swin Transformer Liu et al. (2021) and ConvNeXt Liu et al. 2022) used the natural image domain, we further remove the bottleneck layer (ResNet block with 768 channels) and increase the block depth of stage 3 (e.g., g.. 8 blocks). Such optimized design further significantly reduces both the model parameters (from 53.0M to 32. M.. nn-UNet: 31.2M) and FLOPs (from 639.4G 536.10 1G, nn-UNet: 743.3M), while preserving the performance (shown
14

---

Published as a conference paper at ICLR 2023
in Table 7). Additional validation studies is needed to investigate the effectiveness of both MLP and pointwise DCS, and optimizing 3D UX-Net architecture, which will be the next steps of our future work. Another observation in Table is the subtle differences in model parameters between kernel size of 3 x 3 3 and 7 x x 7. We found that the increase both model parameters and FLOPs is also attributed to the design of decoder network. Our decoder block design further add 3D ResNet block after the transpose convolution to further resample and mix the channel context, instead of directly perform transpose convolution in nn-UNet. A efficient block design in decoder network is demanded to be further investigated and using depthwise convolution may be another potential solution to reduce the low efficiency burden.
To further reduce the burden of low training and inference efficiency, pe-aaamtttrrzztionn of LK convolutional blocks may be another promising direction to focus. Prior works have demonstrated to scale few convolutional blocks with LK size (31 x 31) and propose the idea of parallel branches with small kernels for residual shortcuts Ding et al. (2022b) 2021 2022a). The parallel branch can then be mutually converted through equivalent transformation of parameters. For example. a branch of 1 X 1 convolution and a branch of 7 x 7 convolution, can be transferred into a single branck of 7 X convolution Ding et al. (022T. Furthermore, Hu et al. proposed online convolutional parameterization (OREPA) to leverage a linear scaling at each branch to diversify the optimization directions, instead of applying non-linear normalization after convolution layer Hu al.) (2022). Also, stack of small kernels are leveraged to generate similar receptive field of view as LKs with better training and inference efficiency. The effectiveness of leveraging small kernels stack and multiple parallel branches design will be further investigated as another directions of our future work.
15


Reviews:
[{"review_id": "k-JVkPNySvV", "reviewer": "Reviewer_Ud7h", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "Volumetric segmentation of medical images is still an open problem. This work proposes a new U-net type segmentation network that uses large receptive field depth-wise convolution operations to emulate vision transformers (ViTs). Additionally, patch-wise feature propagation and 1x1 convolutions are introduced as means of reducing the compute compared to ViTs. Experiments on multiple volumetric datasets and relevant CNN and ViT baseline models are reported with small to no improvements. ", "summary_of_the_review": "This work proposes a volumetric segmentation method with a motivation of bridging the performance gap between ViTs and CNNs, by adding depth-wise convolutions with large kernels and couple of other tricks. The premise of the work that ViTs have surpassed CNNs, and CNNs need to emulate ViTs is flawed. Experimental evaluation shows small or no improvements, even with orders of magnitude more compute/parameters compared to other baseline methods. ", "strength_and_weaknesses": "**Strengths:**\n\n* The authors propose a volumetric segmentation method, that is evaluated on multiple relevant datasets. The choice of different architectural components in this work are explained clearly. \n\n* Ablation studies are reported and model hyperparameters are reported clearly.\n\n* Experiments on multiple datasets and several relevant baselines are conducted. \n\n**Weaknesses:**\n\n* **Motivation for the model**: This work aims to reintroduce CNNs to match the performance of Vision transformers (ViTs) for volumetric medical image segmentation. This is strange to me, as the chronology of how these models were developed is somehow warped. This work claims to revisit CNNs for volumetric segmentation, and improve them so that they can emulate ViTs -- which were already mimicking CNNs. This logic of arguments is circular to me. I found it very hard to subscribe to the primary motivation of this work. Further, and more fundamentally, the claim that ViTs have surpassed CNNs in volumetric segmentation is highly exaggerated (first sentence in abstract). Going by even the results reported in this work this claim does not hold up.\n\n* **Experimental evaluation**: Results in Table 1 which have the main results are reported in single repeats. For the scale of performance improvements that is claimed, this could be just random, as the difference in Dice accuracy is very small (first or second decimal in most cases). Reading them as improvements based on single model runs is highly problematic. These small differences could simply be overcome with different initialisations. Also, the number of parameters between the current model and a CNN baseline like nn-Unet is double. The claims of superior performance under these conditions is not convincing.\n\n* **Additional insights**: Unfortunately, there are no new insights gained from this work. How does the use of the depth-wise convolution with large kernels influence segmentation across different tasks. What is the reason for these improvements? What type of errors that the other ViT or CNN models make are alleviated with this new architecture? If ViTs were capable of modelling longe range interactions and CNNs introduced image-specific inductive biases, how does the new model bridge this gap? How can this be evaluated?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}, {"review_id": "cZF_PdFOqD", "reviewer": "Reviewer_ERdC", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors present an extension to convolutional networks (ConvNets) tailored to simulate the behavior of hierarchical transformers, aiming to provide an alternative to visional transformers (ViT) with a significantly reduced number of trainable parameters. The architectural variations in the proposed model focused on volumetric depth-wise convolutions with large kernel size for feature extraction that enables larger global receptive field. They then provide experimental evidence using 3 datasets that demonstrated an overall improvement in performance of 3D semantic segmentation in medical imaging.", "summary_of_the_review": "In the current stage of scientific research in medical image segmentations, after the immersion of Visual transformers, this study provides a relevant point of view, aiming to simulate a convolutional neural network alternative with comparable performance and simplified architecture. The study is well applied and communicated in this manuscript and would be a good contribution to the field and ICLR.", "strength_and_weaknesses": "Strength:\nThe manuscript is well written and constructed, the figures and tables are coherent.\nThe provided context in the literature is adequate.\nThe experimental validation is well applied and provides evidence to support the authors conclusions.\n\nWeaknesses:\nWhile the authors provide an extensive literature review to provide the right context to their research, it would be beneficial if they could explicitly outline in paragraph 2.2 the differences of the proposed approach to the previously published models. Particularly, relating to the main inspiration of this paper, e.g. Liu et al. 2022, in the context of novelty should be better clarified.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "ZeHngXCvI5", "reviewer": "Reviewer_8E6F", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper explores the necessity of incorporating transformer elements (self-attention, swin module) into 3D segmentation networks. It demonstrates SOTA performance on three challenge datasets (FLARE, FeTA, AMOS) with comparable model size and computational FLOPS. It is heavily inspired Liu's \"A ConvNet for the 2020s\" CVPR paper, but nonetheless makes a meaningful contribution for practical 3D applications that could be further enhanced by a slightly better ablation study. ", "summary_of_the_review": "A recommend acceptance, subject to moderate improvements in terms of the comprehensiveness of the ablation study. Both the practical impact and the methodological insights (revisiting larger kernels and depth-separable convolutions in contrast to transformer modules) are relevant.", "strength_and_weaknesses": "**Strengths**\n+ the paper is well written, easy to follow and addresses a relevant task in 3D/medical image analysis. \n+ the experimental section contains numerous direct comparisons with different versions of Swin/Vision-Transformers and demonstrates both faster convergence and higher accuracy of the proposed 3D UX-Net\n+ the authors also evaluate pre-train/fine-tune scenarios and reach SOTA performance on AMOS22\n+ within the method a new module: \"depthwise convolutional scaling (DCS)\" is introduced that performs slightly better than MLPs\n\nthe work gives (yet another) reason to consider claimed advantages of transformers with much care, since e.g. the increased receptive field of using 7x7x7 convolutions together with reasonable practical design choices alone can outperform those in a simpler Conv-only model\n\n**Weaknesses**\n- while closely following the narrative of ConvNeXt (which I liked) the authors perform there ablation study in the reverse direction: ie. in Tab 3 they start from the proposed model and **remove** individual parts, while ConvNeXt starts from a plain ResNet and **adds** the transformer-inspired changes. This is unfortunate, because for 3D segmentation there is indeed one prominent framework the nnUNet that uses a plain Conv-only model and has to date won nearly every medical segmentation challenge (incl. AMOS22). This is also evident from Tab. 1 where the nnUNet outperforms every single transformer approach and is only beaten by the proposed method (in light of these facts I would also strongly recommend to rephrase/remove the statement \"To our best knowledge, this is the first 3D ConvNet architecture that competes favorably with transformers SOTA in volumetric segmentation tasks.)  A somewhat more elaborate discussion on the differences of the nnUNet (which e.g. also doesn't use BatchNorm) would be of important value. Importantly: the kernel-size ablation should in addition be extended by 3x3x3 and 5x5x5!\n- I am also missing a more detailed discussion about runtimes for training and inference, since the raw FLOPs numbers are not always very favourable for UX-Net and depth-separable convolutions are additionally prone to lower efficiency (when comparing similar FLOP numbers). On the other hand the nnUNet is known to be fairly resource intensive for training, so any method that reduces this burden would be of great impact!\nAs a minor comment (since this was published after the ICLR submission deadline): one could include the nnUNet results for AMOS22 from https://arxiv.org/pdf/2208.10791.pdf in the final version ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster