[論文レビュー] MedNeXt: Transformer-driven Scaling of ConvNets for Medical Image Segmentation

Saikat Roy, Gregor Koehler+6 more|arXiv (Cornell University)|Mar 17, 2023

Radiomics and Machine Learning in Medical Imaging被引用数 17

ひとこと要約

MedNeXt は、残差反転ボトルネックと UpKern カーネルアップサンプリングを用いた完全 ConvNeXt 3D エンコーダ-デコーダアーキテクチャを提案し、CTおよびMRIデータセット全体で最先端の医用画像分割を達成します。

ABSTRACT

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.

研究の動機と目的

• データ不足の条件下で、完全な ConvNeXt 3D エンコーダ-デコーダが Transformer ベースおよび大カーネルのベースラインを医用画像分割で上回ることを示す。
• アップ/ダウンサンプリングにおける残差反転ボトルネックを導入し、スケールを横断して意味的豊かさを保持する。
• UpKern を開発: 限られたデータで大カーネルの訓練飽和を緩和するカーネルアップサンプリング初期化手法。
• 深さ、幅、受容野にわたる複合スケーリングを適用し、タスクとモダリティ全体で性能を最適化する。

提案手法

• ConvNeXt ブロックのみで構成された 4 層の Encoder-Decoder MedNeXt アーキテクチャを使用する。
• Up/Downsampling パスに残差反転ボトルネックを実装して勾配フローと意味的保存を改善する。
• UpKern を導入: 三次線形 upsampling ベースの初期化で、から訓練をやり直すことなくカーネルサイズを拡大する。
• MedNeXt の構成 (S, B, M, L) にわたり、深さ、幅、カーネルサイズを共にスケールさせる複合スケーリングを採用する。
• 4データセットで5分割交差検証を用いて訓練し、nnUNet および複数のTransformer/大カーネルベースラインと比較する。

((a)) MedNeXt macro and block architecture

実験結果

リサーチクエスチョン

• RQ1 データが限定的な医用画像タスクにおいて、完全な ConvNeXt 3D 分割ネットワークは Transformer ベースおよび大カーネル手法を超えることができるか？
• RQ2 残差反転ボトルネックと UpKern 初期化は、医用分割における大カーネルの性能と訓練安定性を改善するか？
• RQ3 深さ、幅、受容野にわたる複合スケーリングは、サイズの異なる CT および MRI データセット全体で一貫した利得を生み出すか？

主な発見

Network	BTCV DSC	BTCV SDC	AMOS22 DSC	AMOS22 SDC	KiTS19 DSC	KiTS19 SDC	BraTS21 DSC	BraTS21 SDC	AVG DSC	AVG SDC
nnUNet Baselines	83.56	86.07	88.88	91.70	89.88	86.88	91.23	90.46	88.39	88.78
UNETR	75.06	75.00	81.98	82.65	84.10	78.05	89.65	88.28	82.36	81.00
TransUNet	76.72	76.64	85.05	86.52	80.82	72.90	89.17	87.78	82.94	80.96
TransBTS	82.35	84.33	86.52	88.84	87.03	83.53	90.66	89.71	86.64	86.60
nnFormer	80.76	82.37	84.20	86.38	89.09	85.08	90.42	89.83	86.12	85.92
SwinUNETR	80.95	82.43	86.83	89.23	87.36	83.09	90.48	89.56	86.41	86.08
3D-UX-Net	80.76	82.30	87.28	89.74	88.39	84.03	90.63	89.63	86.77	86.43
MedNeXt-S kernel:3	83.90	86.60	89.03	91.97	90.45	87.80	91.27	90.46	88.66	89.21
MedNeXt-B	84.01	86.77	89.14	92.10	91.02	88.24	91.30	90.51	88.87	89.41
MedNeXt-M	84.31	87.34	89.27	92.28	90.78	88.22	91.57	90.78	88.98	89.66
MedNeXt-L	84.57	87.54	89.58	92.62	90.61	88.08	91.57	90.81	89.08	89.76
MedNeXt-S kernel:5	83.92	86.80	89.27	92.26	90.08	87.04	91.40	90.57	88.67	89.17
MedNeXt-B	84.23	87.06	89.38	92.36	90.30	87.40	91.48	90.70	88.85	89.38
MedNeXt-M	84.41	87.48	89.58	92.65	90.87	88.15	91.49	90.67	89.09	89.74
MedNeXt-L	84.82	87.85	89.87	92.95	90.71	87.85	91.46	90.73	89.22	89.85

• MedNeXt 系は、ベースラインと比較して4つのデータセット（BTCV、AMOS22、KiTS19、BraTS21）で最先端の性能を達成する。
• UpKern 初期化とカーネル 5×5×5 の MedNeXt-L は、公開テストセットで nnUNet を上回る（DSC スコア；BTCV 88.76、AMOS22 91.77、KiTS19 91.02、 BraTS21 88.01）。
• アブレーションにより、残差反転ボトルネックが標準リサンプリングより著しく性能を向上させることを示す。
• UpKern 初期化により大カーネルネットワーク（5×5×5）が小カーネルモデルを上回る。大カーネルをゼロから訓練すると劣る。
• 深さ、幅、カーネルサイズの複合スケーリングは、データセット全体でさらなる改善をもたらす。
• 5分割 CV では、3×3×3 または 5×5×5 カーネルを用いた MedNeXt バリアントが、全データセットで全ベースラインを上回る。

((b)) UpKern Initialization