[논문 리뷰] MedNeXt: Transformer-driven Scaling of ConvNets for Medical Image Segmentation

Saikat Roy, Gregor Koehler외 6명|arXiv (Cornell University)|2023. 03. 17.

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-층 인코더-디코더 MedNeXt 아키텍처를 사용한다.
• 업/다운샘플링 경로에 잔여 반전 병목을 구현하여 그래디언트 흐름과 의미 보존을 개선한다.
• UpKern 도입: 3선형 업샘플링 기반 초기화를 통해 처음부터 재학습 없이 커널 크기를 확장한다.
• MedNeXt 구성(S, B, M, L) 전반에 걸쳐 깊이, 폭 및 커널 크기를 함께 확장하기 위한 복합 스케일링을 적용한다.
• 네 개 데이터셋에서 5-폴드 교차 검증으로 학습하고 nnUNet 및 다수의 Transformer/대형 커널 baselines 와 비교한다.

((a)) MedNeXt macro and block architecture

실험 결과

연구 질문

• RQ1 제한된 데이터에서 의학 영상 작업에서 완전 ConvNeXt 3D 분할 네트워크가 Transformer 기반 및 대형 커널 방법을 능가할 수 있는가?
• RQ2 잔여 반전 병목 및 UpKern 초기화가 대형 커널의 성능 및 학습 안정성을 향상시키는가?
• RQ3 깊이, 폭 및 수용 영역에 걸친 복합 스케일링이 다양한 크기의 CT 및 MRI 데이터셋에서 지속적인 이점을 가져오는가?

주요 결과

네트워크	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 변형은 네 가지 데이터셋(BTCV, AMOS22, KiTS19, BraTS21)에서 기준선과 비교하여 최첨단 성능을 달성한다.
• 커널 5×5×5 및 UpKern 초기화를 갖춘 MedNeXt-L 이 공개 테스트 세트에서 nnUNet 을 능가한다 (DSC 점수; BTCV 88.76, AMOS22 91.77, KiTS19 91.02, BraTS21 88.01).
• 소거실험에서 잔여 반전 병목이 표준 리샘플링에 비해 성능을 크게 향상시킴을 보여준다.
• UpKern 초기화로 대형 커널 네트워크(5×5×5)가 소형 커널 모델을 능가하도록 가능하게 하며, 대형 커널로의 무작정 재학습은 성능이 떨어진다.
• 깊이, 폭 및 커널 크기에 걸친 복합 스케일링이 데이터셋 전반에서 추가 개선을 가져온다.
• 5-폴드 교차 검증에서 MedNeXt 변형이 3×3×3 또는 5×5×5 커널을 사용할 때 모든 데이터셋에서 모든 기준선을 능가한다.

((b)) UpKern Initialization