Source-linked AI summary
MedNeXt: Transformer-driven Scaling of ConvNets for Medical Image Segmentation
Saikat Roy, Gregor Koehler, Constantin Ulrich, Michael Baumgartner, Jens Petersen, Fabian Isensee, Paul F. Jaeger, Klaus Maier-Hein
TL;DR
Medical image segmentation lacks the abundant high-quality annotations that help Transformers realize their benefits, motivating architectures better suited to data-scarce settings. MedNeXt addresses this with a scalable, Transformer-inspired fully ConvNeXt 3D network using residual resampling, UpKern, and compound scaling, achieving state-of-the-art performance across four CT and MRI tasks.
Problem
Medical image segmentation often lacks the large, high-quality annotated datasets needed to maximize Transformer performance benefits.
Method
MedNeXt is a scalable fully ConvNeXt 3D encoder-decoder with residual up/downsampling blocks, UpKern initialization, and compound scaling of depth, width, and kernel size.
Results
MedNeXt achieves state-of-the-art performance across four CT and MRI segmentation tasks with dataset sizes ranging from 30 to 1251 samples.
Takeaways & Limitations
MedNeXt is presented as a scalable modernization of standard convolution blocks for medical image segmentation on limited datasets.
Abstract
from arXiv · showhide
There has been exploding interest in embracing Transformer-based architectures for medical image segmentation. However, the lack of large-scale annotated medical datasets make achieving performances equivalent to those in natural images challenging. Convolutional networks, in contrast, have higher inductive biases and consequently, are easily trainable to high performance. Recently, the ConvNeXt architecture attempted to modernize the standard ConvNet by mirroring Transformer blocks. In this work, we improve upon this to design a modernized and scalable convolutional architecture customized to challenges of data-scarce medical settings. We introduce MedNeXt, a Transformer-inspired large kernel segmentation network which introduces - 1) A fully ConvNeXt 3D Encoder-Decoder Network for medical image segmentation, 2) Residual ConvNeXt up and downsampling blocks to preserve semantic richness across scales, 3) A novel technique to iteratively increase kernel sizes by upsampling small kernel networks, to prevent performance saturation on limited medical data, 4) Compound scaling at multiple levels (depth, width, kernel size) of MedNeXt. This leads to state-of-the-art performance on 4 tasks on CT and MRI modalities and varying dataset sizes, representing a modernized deep architecture for medical image segmentation. Our code is made publicly available at: https://github.com/MIC-DKFZ/MedNeXt.
1 Introduction
MedNeXt modernizes ConvNeXt for data-scarce medical image segmentation by combining a fully ConvNeXt 3D encoder-decoder with residual resampling, UpKern initialization, and compound scaling. It achieves state-of-the-art performance across four CT and MRI segmentation tasks spanning datasets of 30 to 1251 samples.
- Architecture: MedNeXt is a scalable, fully ConvNeXt 3D encoder-decoder network designed for medical image segmentation.The architecture is presented as a modernization of ConvNets for limited medical datasets.
- Architecture: Residual Inverted Bottlenecks replace regular upsampling and downsampling blocks to preserve contextual richness during resampling.The modified residual connection is reported to improve gradient flow during training.
- Kernel scaling: UpKern iteratively increases kernel size by initializing large-kernel MedNeXts with trained, upsampled small-kernel networks.The technique is introduced to prevent performance saturation on large-kernel models.
- Compound scaling: MedNeXt supports orthogonal scaling of width, receptive field, and depth through channels, kernel size, and layer count.This compound scaling design targets multiple network dimensions simultaneously.
- Results: MedNeXt achieves state-of-the-art performance against Transformer-based, convolutional, and large-kernel baselines on four CT and MRI tasks.The evaluated datasets range from 30 to 1251 samples and include organ and tumor segmentation.
2 Proposed Method
The proposed method embeds ConvNeXt blocks throughout a 3D U-Net-like encoder-decoder, including residual upsampling and downsampling blocks. It addresses large-kernel saturation with UpKern and scales depth, width, and receptive field together through multiple MedNeXt configurations.
- Fully ConvNeXt 3D Segmentation Architecture: MedNeXt uses ConvNeXt as the building block in a 3D U-Net-like macro architecture.The architecture extends ConvNeXt blocks to upsampling and downsampling layers.
- Fully ConvNeXt 3D Segmentation Architecture: The depthwise convolution uses k × k × k kernels and GroupNorm, while the expansion and compression layers use 1×1×1 convolutions.The expansion layer uses GELU activation and separates width scaling from receptive-field scaling.
- Residual up/downsampling blocks: MedNeXt extends inverted bottlenecks to resampling by placing strided or transposed convolutions in the first depthwise layer.This design allows upsampling and downsampling blocks to participate in ConvNeXt scaling.
- UpKern: UpKern initializes larger-kernel networks from compatible pretrained smaller-kernel networks by trilinearly upsampling convolutional kernels.Unchanged tensors, including normalization layers, are copied directly.
- Compound Scaling of Depth, Width and Receptive Field: Compound scaling varies block count, expansion ratio, and kernel size to jointly scale depth, width, and receptive field.The study evaluates four MedNeXt configurations and kernel sizes k = {3, 5}.
3 Experimental Design
The experiments evaluate MedNeXt across four organ and tumor segmentation tasks using a common supervised framework and strong convolutional, transformer-based, hybrid, and ConvNeXt baselines.
- Baselines: The baselines include nnUNet, four convolution-transformer hybrids, nnFormer, and 3D-UX-Net, with TransUNet as the only 2D network.The transformer-based models include UNETR, SwinUNETR, TransBTS, TransUNet, and nnFormer.
- Datasets and tasks: The study evaluates BTCV, AMOS22, KiTS19, and BraTS21, covering organ and tumor segmentation across CT and MRI datasets.The datasets include 30, 200, and 210 CT volumes for BTCV, AMOS22, and KiTS19, respectively.
4 Results and Discussion
MedNeXt’s architectural modifications improve segmentation performance in ablations, and its scalable variants outperform or match strong baselines across four heterogeneous tasks and official test sets.
- Ablation results: Residual Inverted Bottlenecks are functionally enabling: removing them causes considerably worse segmentation performance.The proposed resampling blocks are intended to preserve semantic richness during upsampling and downsampling.
- Ablation results: UpKern improves 5 × 5 × 5 MedNeXt performance on BTCV and AMOS22, whereas large-kernel training from scratch makes large- and small-kernel performance indistinguishable.The large-kernel benefit is attributed to combining UpKern with larger kernels rather than simply doubling the training schedule.
- Cross-validation results: MedNeXt with 3 × 3 × 3 kernels achieves state-of-the-art five-fold cross-validation performance against every baseline on all four datasets.Depth and width scaling provide strong performance across organs and tumors, CT and MRI, and training sets ranging from 18 to 1000 samples.
- Cross-validation results: With UpKern and 5 × 5 × 5 kernels, MedNeXt improves over its small-kernel networks comprehensively for organ segmentation and more selectively for tumor segmentation.The broader improvement occurs on BTCV and AMOS22, while tumor-task gains are more limited on KiTS19 and BraTS21.
- Test-set results: On official test sets, MedNeXt exceeds nnUNet on BTCV, AMOS22, KiTS19, and BraTS21, including AMOS22 leaderboard rank 1 with DSC 91.77 and NSD 84.00.Reported test-set results include BTCV DSC 88.76 and HD95 15.34, KiTS19 DSC 91.02, and BraTS21 DSC 88.01 and HD95 10.69.
5 Conclusion
MedNeXt addresses limited-data medical segmentation with a scalable, Transformer-inspired, fully ConvNeXt 3D architecture and reports state-of-the-art results across four challenging tasks.
- Conclusion: MedNeXt is presented as a scalable, Transformer-inspired, fully ConvNeXt 3D segmentation architecture customized for limited medical image datasets.The paper positions compound scaling as a modernization of standard convolution blocks for deep medical segmentation networks.
- Conclusion: MedNeXt demonstrates state-of-the-art performance across four challenging tasks against seven strong baselines.The conclusion describes the architecture as an effective modernization of standard convolution blocks for medical image segmentation.
6 Supplementary Material
The supplementary material provides concise UpKern pseudocode and configuration information for MedNeXt models, including computational cost and training time.
- UpKern algorithm: The 3D UpKern procedure can be translated to 2D by replacing trilinear interpolation with bilinear interpolation.The technique is presented as general beyond the 3D architecture used in this work.
- UpKern algorithm: UpKern initialization is summarized in 20 lines of PyTorch pseudocode for large-kernel MedNeXt models.The algorithm is described as an upsampled-kernel initialization procedure.
- UpKern algorithm: The pseudocode iterates over model state-dictionary keys, copies compatible pretrained entries, and applies upsampled-kernel initialization otherwise.Common keys are detected before loading the resulting state dictionary into the network.
- Model configuration: Table 3 reports MedNeXt configurations, GFLOPs, and training time for kernel sizes 3 and 5.GFLOPs are measured for one 128 × 128 × 128 patch, and training time covers one epoch with batch size 2 and 250 batches.
7 Tabulation of Results on 4 Test Sets
Table 4 tabulates results on public validation sets for BTCV, AMOS22, KiTS19, and BraTS21. All datasets use the same MedNeXt-L ensemble and nnUNet 3d_fullres configuration.
- Results are reported for four datasets: BTCV, AMOS22, KiTS19, and BraTS21.
- The evaluation uses a single 5-fold ensemble of MedNeXt-L with 5 × 5 × 5 kernels initialized using UpKern.
- The experiments use nnUNet's 3d_fullres mode without cascade mode.