[논문 리뷰] 3D UX-Net: A Large Kernel Volumetric ConvNet Modernizing Hierarchical Transformer for Medical Image Segmentation

Ho Hin Lee, Shunxing Bao외 2명|arXiv (Cornell University)|2022. 09. 29.

Medical Imaging and Analysis인용 수 99

한 줄 요약

본 논문은 대형 커널 깊이 방향 컨볼루션을 갖춘 순수 ConvNet인 3D UX-Net을 제시하며, 볼륨형 의학 이미지 분할을 위해 계층적 트랜스포머 동작을 모방한다. 매개변수가 더 적은 상태에서 공개 데이터셋에서 최첨단 결과를 달성한다.

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. 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 size (e.g. starting from $7 imes7 imes7$) 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.

연구 동기 및 목표

• 성능과 모델 규모의 균형을 맞춘 효율적인 3D 분할 백본의 필요성을 동기 부여한다.
• 대형 커널 깊이 방향 컨볼루션을 사용하여 계층적 트랜스포머 동작을 모방하는 경량화된 부피형 ConvNet를 제안한다.
• 분할 정확도를 유지하거나 향상시키면서 매개변수 수와 정규화 요구 사항을 줄인다.
• 감독학습 및 전이 학습 설정에서 공개 뇌 및 복부 데이터셋에 대한 강력한 실증 결과를 보여준다.

제안 방법

• 큰 커널 크기(LK)를 갖는 볼륨 깊이 방향 컨볼루션을 도입하여 큰 수용 영역을 시뮬레이션한다.
• 트랜스포머 블록의 MLP를 포인트와이즈 깊이 컨볼루션으로 대체하여 매개변수 수를 줄이면서 특징을 확장한다(DCS).
• 깊이 방향 컨볼루션을 이용한 역병목(inverted bottleneck) 설계를 사용하여 계층 간 채널 특징을 확장하고 축소한다.
• 인코더 블록에서 배치 정규화를 레이어 정규화로 대체하고 GELU 활성화를 사용한다.
• 4단계 인코더를 구성하되 각 단계에 두 개의 LK 블록을 배치하고, 이어서 스킵 연결이 있는 ConvNet 기반 U-형 디코더를 사용한다.

실험 결과

연구 질문

• RQ1 대형 커널 깊이 방향 컨볼루션을 갖춘 순수 ConvNet이 트랜스포머 기반 3D 분할의 성능에 근접하거나 이를 능가할 수 있는가?
• RQ2 제안된 3D UX-Net이 더 적은 매개변수와 축소된 정규화 계층으로 비교 가능한 혹은 더 나은 정확도를 달성하는가?
• RQ3 공개 부피 데이터셋에서 감독 학습과 전이 학습에서 3D UX-Net의 성능은 어떠한가?
• RQ4 데이터셋 전반에 걸친 분할 성능에 대한 커널 크기와 깊이 방향 스케일링(DCS)의 영향은 무엇인가?

주요 결과

• 3D UX-Net은 FeTA2021에서 SwinUNETR보다 우수합니다(0.874 대 0.867 Dice) 및 FLARE2021에서(0.934 대 0.929 Dice).
• AMOS2022 전이 학습에서 3D UX-Net은 Dice 0.900를 달성하며, 최적의 트랜스포머 기준선 대비 2.27% 향상을 보인다.
• 결과 분석 연구는 커널 크기와 깊이 방향 스케일링이 성능에 영향을 미치며, LK 크기가 대략 7x7x7에서 13x13x13 범위에서 데이터셋당 최적 이득을 제공한다.
• FeTA2021에서 빠르게 수렴하고 AMOS2022에서 강건한 전이 학습 동작을 보여준다.
• 모델은 여러 트랜스포머 기반 기준선과 비교하여 매개변수 수가 더 적은(53.0M) 경쟁력 있는 Dice 점수를 달성한다.