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

In plain words

Medical scans in 3D, such as brain MRIs or abdominal CTs, need every voxel labelled by organ or tissue. The strongest recent models for this borrow from transformers, which look at the whole scan through attention. That works well but is heavy: many parameters and slow training.

This paper asks whether a plain convolutional network can behave like those transformers. Its trick is to use very large 3D filters, applied one channel at a time, so that each output sees a wide neighbourhood at low cost. It also copies a transformer habit of widening the channels inside each block and then narrowing them again, and it replaces some normalisation and activation layers with simpler ones. The result is a U-shaped network that the authors say matches or beats a transformer baseline on three public benchmarks while using fewer parameters. The gains are small, and the authors report them themselves.

TL;DR

3D UX-Net is a volumetric ConvNet for 3D medical image segmentation. Its encoder copies the block layout of the Swin Transformer: a large depthwise convolution in place of window attention, and a per-channel expand-and-compress step in place of the MLP. The authors report mean Dice gains over SwinUNETR on FeTA 2021, FLARE 2021 and AMOS 2022 fine-tuning, with fewer parameters but more FLOPs. Use it when you need a pure-CNN 3D segmenter that a transformer baseline can be compared against. The official code is public, but several training settings differ from the paper (see Paper vs code).

Problem

Hierarchical transformers such as Swin and SwinUNETR set the state of the art for 3D medical segmentation. Their advantage is attributed to a large receptive field from windowed self-attention and to a large parameter count §1. Global self-attention over dense 3D volumes is too expensive, so windowed variants are used §1. Existing 3D ConvNets rely on small kernels, and the few large-kernel 3D designs reported cost in FLOPs and speed §2. The paper's question is whether large-kernel depthwise convolutions can give the same receptive-field effect as shifted-window attention at lower parameter cost §1 §3. The assumptions are three public volumetric datasets, supervised training, and fixed 96×96×96 sub-volumes §5 tab. 4.

Method

The network is a U-shaped encoder-decoder. The encoder has four stages, each with two large-kernel blocks; the resolution halves between stages, from H/2 to H/16 §4.1. The decoder is a convolutional U-Net-style path with skip connections §4.2.

Three design choices come from the Swin Transformer block §3:

1. Patch projection. A 7×7×7 convolution with stride 2 maps the input to 48 channels, in place of a linear patch embedding §4.1. The kernel is fixed at 7 in the code uxnet_encoder.py.
2. Large-kernel depthwise convolution (DWC). Each channel is convolved with its own 7×7×7 kernel, starting from 7 and varied in ablations §3 §4.1. The paper argues this stands in for windowed attention §3 §7.
3. Inverted bottleneck with depthwise convolutional scaling (DCS). Each channel is expanded to four times its width and compressed back by 1×1×1 operations that act on each channel independently. This adds no cross-channel mixing and keeps the parameter count down §3 §4.1.

The block recursion is written in eq. 1. In prose: each block applies a DWC with a residual connection, then a DCS with a residual connection, with layer normalisation (LN) before each. Stage-to-stage downsampling uses a 2×2×2 convolution with stride 2, which exchanges information across channels, in place of an MLP §4.1.

# One large-kernel block: two residual sub-layers  [§4.1] [eq. 1]
# Input:  z_prev  (B, C, H, W, D)
# Output: z       (B, C, H, W, D)

def ux_block(z_prev):
    # Depthwise 7x7x7 convolution after LN, with residual  [§3] [eq. 1]
    hat_z = depthwise_conv_7(layer_norm(z_prev)) + z_prev
    # Per-channel expand to 4C, then compress to C, with residual  [§3] [eq. 1]
    # (no channel mixing in either step)  [§4.1]
    z = channel_scale_4x(layer_norm(hat_z)) + hat_z
    return z

Encoder stages and downsampling:

# Encoder: four stages of two blocks each  [§4.1]
# Channels per stage: 48, 96, 192, 384  [tab. 4]

def encoder(volume):
    # Patch projection: 7x7x7 conv, stride 2, to 48 channels  [§4.1]
    x = conv3d_k7_stride2(volume)
    for stage in range(4):
        # Downsample by 2 with a 2x2x2 conv, stride 2, before stages 2-4  [§4.1]
        if stage > 0:
            x = conv3d_k2_stride2(x)
        # Two large-kernel blocks per stage  [§4.1]
        x = ux_block(ux_block(x))
        skips[stage] = x
    return skips

Decoder:

# Decoder: U-shaped, convolutional  [§4.2]

def decoder(skips, volume):
    # Two post-normalised 3x3x3 convs with instance norm per stage  [§4.2]
    # Each stage upsamples with a transpose conv, concatenates the
    # skip from the stage above, then applies a residual block  [§4.2]
    # Final: residual block with 1x1x1 conv and softmax  [§4.2]
    return probabilities

The paper describes encoder and decoder normalisation as instance normalisation in the decoder and LN in the encoder §4.2 §3.

Hyperparameters

Stated values come from tab. 4 and app. A.1 unless marked otherwise. "Code" shows what the official repository does; where it differs, see Paper vs code.

Results

Author-reported. All numbers are mean Dice on the held-out test folds unless stated. Setup: five-fold cross-validation on FeTA 2021 and FLARE 2021 with training from scratch; AMOS 2022 is fine-tuned from the best FLARE 2021 fold §5 app. A.2.

FeTA 2021, seven tissues tab. 1:

FLARE 2021, four organs tab. 1:

AMOS 2022 fine-tuning, fifteen organs tab. 2:

The paper marks the 3D UX-Net means with an asterisk (p < 0.01, Wilcoxon signed-rank test against all listed state-of-the-art methods) tab. 1 tab. 2.

Ablations tab. 3, FeTA / FLARE mean Dice:

nn-UNet structure with 3D UX-Net blocks tab. 6, FeTA / FLARE: plain 0.874 / 0.934 becomes 0.876 / 0.940 without deep supervision and 0.881 / 0.944 with it. The authors trained these at a learning rate of 0.002 and gave epoch counts instead of steps app. A.3.

Efficiency variant tab. 7: depth 2,2,8,2 with hidden size 384 has 32.1M parameters and 536.8G FLOPs, with mean Dice 0.873 (FeTA) and 0.932 (FLARE). Default: 53.0M, 639.4G, 0.874 and 0.934.

Qualitative figures: fig. 4 shows boundary and organ examples. fig. 3 shows validation Dice curves on the three datasets.

Paper vs code

Official code: github.com/MASILab/3DUX-Net at commit 14ea46b. The repository has no LICENSE file; the README says MIT README.md.

• Layer normalisation placement. Eq. 1 applies LN before the DWC and before the DCS, so each block has two LN layers eq. 1. The code applies a single LN after the DWC and none before it uxnet_encoder.py.
• Activation between the scaling layers. Eq. 1 shows no activation in DCS. The code puts GELU between the two per-channel 1×1×1 layers uxnet_encoder.py.
• Layer scale. The code multiplies each block output by a learnable per-channel scale initialised to 1e-6 uxnet_encoder.py uxnet_encoder.py. Eq. 1 has no such term.
• Kernel size is not configurable. The 7×7×7 depthwise kernel is written into the block uxnet_encoder.py. The kernel-size ablation in tab. 3 therefore needs code edits, not a flag.
• Learning-rate scheduler. tab. 4 lists ReduceLROnPlateau. In both training scripts the scheduler is commented out, so the learning rate is constant main_train.py main_finetune.py.
• Weight decay. tab. 4 gives 0.08. The AdamW call passes no weight decay main_train.py, so the PyTorch default applies.
• Batch size and crops. tab. 4 gives batch size 2 and 2 crops for direct training. The default batch size is 1 and the TRAINING.md commands use 1 main_train.py TRAINING.md. Fine-tuning uses 2 crops in the commands, but tab. 4 gives 1 TRAINING.md.
• Intensity window. app. A.1 gives −175 to 250 for FLARE and AMOS. The code uses −125 to 275 load_datasets_transforms.py load_datasets_transforms.py.
• Rotation range. tab. 4 gives ±30°. The code uses ±π/15 (about ±12°) for FeTA load_datasets_transforms.py and ±π/30 (about ±6°) for FLARE and AMOS load_datasets_transforms.py load_datasets_transforms.py.
• Validation metric. The paper's per-organ means exclude background, since the FeTA mean covers seven columns. The code's Dice metric includes the background channel main_train.py. Checkpoint selection therefore uses a different mean from the one reported in tab. 1.
• Inference overlap. The README evaluation command uses overlap 0.7 README.md. The script default is 0.5 test_seg.py. The paper does not state the setting.
• Five-fold splits. The paper reports five-fold cross-validation §5. The code reads imagesTr and imagesVal folders load_datasets_transforms.py and does not implement the fold split. KFold is imported but unused load_datasets_transforms.py. The split files are not in the repository.
• Reported FLARE number. The README lists 0.936 for 3D UX-Net on FLARE 2021 "(latest)" README.md. The paper reports 0.934 tab. 1.
• Training-time metric. The training loop uses a sliding window with batch size 2 and no overlap argument main_train.py.
• Extra normalisation not in the paper. The code applies LN before each stage-to-stage downsampling uxnet_encoder.py. It also passes each stage output through a further channels-first LN before the decoder uses it uxnet_encoder.py. The paper describes neither.
• Other networks in the repo. The repository also contains baseline code for other networks. Only the UX-Net path was checked against the paper.

No discrepancies were found in the 7×7×7 stem, the 2×2×2 stride-2 downsampling, or the four stage widths of 48, 96, 192 and 384.

Pitfalls

• Compute is higher than the parameter count suggests. 3D UX-Net has fewer parameters than SwinUNETR (53.0M against 62.2M) but about twice the FLOPs (639.4G against 328.4G) tab. 1. The abstract reports only the parameter and accuracy gains. Training and inference time may be the real constraint.
• Per-organ means do not reproduce the printed AMOS average. The fifteen 3D UX-Net values in tab. 2 average to 0.898, not 0.900. The FLARE four-organ values average to 0.933, not 0.934. The 2.27% gain quoted in §6.2 is computed from 0.900. The paper does not say how the averages were formed.
• FeTA baseline value. The abstract puts SwinUNETR at 0.867 on FeTA. §6.1 quotes the FeTA improvement as 0.870 to 0.874. In tab. 1, 0.870 is nn-UNet, not SwinUNETR.
• FLARE case count. §5 states 511 CTs for FLARE 2021. tab. 5 totals 361 cases (272 + 69 + 20).
• Wrong figure references in the text. §6.1 cites Figure 2 for convergence curves and Figure 3 for qualitative results. In the compiled paper, Figure 3 shows validation curves and Figure 4 shows the qualitative examples fig. 3 fig. 4.
• Efficiency table versus text. app. A.4 says FLOPs fall to 536.1G and gives nn-UNet as "743.3M". tab. 7 shows 536.8G and 743.3G.
• Undefined learning rate in the nn-UNet comparison. app. A.3 says the "proposed" initial learning rate is 0.01, but tab. 4 gives 1e-4. The 0.002 rate and epoch counts in that section are not explained and do not match the step counts (40000 steps is stated as about 150 epochs for FLARE).
• Steps as epochs. app. A.3 equates 40000 steps with 150 epochs for FLARE at batch size 2. The number of epochs depends on how many crops are drawn per case, which the text does not give.
• Hyperparameter defaults hide the choices. The code uses one-crop-per-sample defaults and never sets the weight decay. Running it with defaults gives a different optimisation from tab. 4.
• Pretrained weights. AMOS fine-tuning requires FLARE-trained weights. The README links Google Drive files for the trained weights README.md README.md. The weights' licence is not stated in the repository.
• Dataset licences. The three challenge datasets carry their own terms. This card does not cover them.
• Rounding. Several printed means differ from the mean of their columns by 0.001 to 0.002 (for example SwinUNETR FeTA, 0.868 against 0.867 printed). Treat third-decimal comparisons between methods with care.

Lineage

• Builds on: Swin Transformer and its use in SwinUNETR §1 §3; ConvNeXt, whose large-kernel 2D depthwise design the paper adapts §2 §3; nn-UNet, used as a baseline and as a training recipe in app. A.3. The paper cites these without arXiv identifiers, and no cards for them exist yet.
• Compared against: SwinUNETR, nnFormer, UNETR, TransBTS, nn-UNet and 3D U-Net tab. 1 tab. 2.
• Superseded by: not known from content already in this wiki.