[논문 리뷰] Optimized U-Net for Brain Tumor Segmentation

Michał Futrega, Alexandre Milesi외 2명|arXiv (Cornell University)|2021. 10. 07.

Brain Tumor Detection and Classification인용 수 8

한 줄 요약

이 논문은 BraTS21 챌린지에서 뇌종양 세분화를 위한 최적화된 U-Net 아키텍처를 제안하며, 검증 단계에서 1등, 테스트 단계에서 3등을 기록했다. 방법은 깊은 감독, 더 깊은 인코더와 증가된 채널 수, 배경이 아닌 병변 병변의 원-핫 인코딩, 그리고 기준 U-Net보다 더 높은 Dice 점수를 향상시키기 위한 맞춤형 후처리 전략을 조합한다.

ABSTRACT

We propose an optimized U-Net architecture for a brain tumor segmentation task in the BraTS21 challenge. To find the optimal model architecture and the learning schedule, we have run an extensive ablation study to test: deep supervision loss, Focal loss, decoder attention, drop block, and residual connections. Additionally, we have searched for the optimal depth of the U-Net encoder, number of convolutional channels and post-processing strategy. Our method won the validation phase and took third place in the test phase. We have open-sourced the code to reproduce our BraTS21 submission at the NVIDIA Deep Learning Examples GitHub Repository.

연구 동기 및 목표

• 다중 모odal MRI 스캔에서 자동 뇌종양 세분화를 위한 고성능 U-Net 기반 모델을 개발하는 것.
• 더 높은 세분화 정확도를 위한 아키텍처 구성 요소와 학습 전략의 최적 조합을 규명하는 것.
• 제한된 의료 학습 데이터와 종양의 외형 변동성에 대응하기 위해 체계적인 아블레이션 연구를 수행하는 것.
• 아키텍처 및 학습 스케줄 최적화를 통해 BraTS21 벤치마크에서 최신 기술 성능을 달성하는 것.

제안 방법

• 기울기 흐름과 학습 안정성을 향상시키기 위해 깊은 감독을 적용한 U-Net 아키텍처를 사용하였다.
• 표현 능력을 향상시키기 위해 인코더 깊이를 6단계에서 7단계로 증가시키고, 컨볼루션 필터 수를 늘렸다.
• 공간 인식 능력을 향상시키기 위해 병변 병변의 원-핫 인코딩 채널을 입력에 추가하였다.
• 작은, 낮은 신뢰도를 보이는 강화 종양(ET) 영역을 제거하여 잠재적 오진을 줄이기 위해 다단계 후처리 전략을 적용하였다.
• 병변 세분화에서의 클래스 불균형 문제를 해결하기 위해 Focal 손실과 Dice 손실을 조합한 하이브리드 손실을 사용하였다.
• 여러 가지 U-Net 변종과 구성 요소에 대해 광범위한 5중 교차 검증 아블레이션 연구를 수행하였다.

실험 결과

연구 질문

• RQ1 기본 U-Net, 어텐션 U-Net, 잔차 U-Net, SegResNetVAE, UNETR 중 BraTS21에서 가장 우수한 성능을 내는 U-Net 변종은 무엇인가?
• RQ2 깊은 감독, 드롭 블록, 잔차 연결과 같은 아키텍처 구성 요소는 세분화 성능에 어떤 영향을 미치는가?
• RQ3 최대 Dice 점수를 얻기 위해 최적의 인코더 깊이, 채널 수, 입력 전처리 조합은 무엇인가?
• RQ4 후처리 전략은 특히 작은, 낮은 신뢰도를 보이는 ET 영역에서 최종 세분화 성능에 어떤 영향을 미치는가?
• RQ5 BraTS21 데이터셋에서 최고의 일반화 성능을 얻기 위해 최적의 학습 스케줄 및 손실 함수 설정은 무엇인가?

주요 결과

• 깊은 감독을 적용한 U-Net가 모든 테스트 변종 중 평균 Dice 점수 0.9149를 기록하여 기준 U-Net(0.9130)을 초월했다.
• 더 깊은 인코더(깊이 7), 증가된 채널 수(64–512), 원-핫 병변 입력을 조합함으로써 Dice 점수는 0.9156까지 향상되었다.
• 모든 최적화 요소(더 깊은 인코더, 수정된 채널 수, 원-핫 입력, 후처리)를 포함한 최종 모델은 검증 단계에서 평균 Dice 점수 0.9163을 기록했다.
• 후처리 전략은 낮은 신뢰도를 보이는 작은 ET 성분을 제거하여 오진을 줄였으며, 점수를 0.9156에서 0.9163으로 향상시켰다.
• 모델는 BraTS21의 검증 단계에서 우승했고, 최종 테스트 단계에서 3등을 차지했다.
• Focal 손실과 드롭 블록은 기준 모델 대비 성능 향상이 유의미하지 않아 이 설정에서는 유의미한 이점이 없음을 시사했다.