[논문 리뷰] Multiclass MRI Brain Tumor Segmentation using 3D Attention-based U-Net

Maryann M. Gitonga|arXiv (Cornell University)|2023. 05. 10.

Advanced Neural Network Applications인용 수 5

한 줄 요약

이 논문은 뇌 종양 다중 분류 분할을 향상시키기 위해 세 가지 MRI 모odal리티(T2-FLAIR, T1CE, T2)를 하나의 4차원 입력으로 융합하는 3D 어텐션 기반 U-Net 아키텍처를 제안한다. 스위프 커넥션에 3D 채널 및 공간 어텐션 모듈을 통합함으로써 종양 영역에 대한 집중을 높이고 배경 조직을 억제하며, BraTS 2021 테스트 세트에서 0.9864의 Dice 계수를 달성하여 이전 최고 성능 방법을 초월한다.

ABSTRACT

This paper proposes a 3D attention-based U-Net architecture for multi-region segmentation of brain tumors using a single stacked multi-modal volume created by combining three non-native MRI volumes. The attention mechanism added to the decoder side of the U-Net helps to improve segmentation accuracy by de-emphasizing healthy tissues and accentuating malignant tissues, resulting in better generalization power and reduced computational resources. The method is trained and evaluated on the BraTS 2021 Task 1 dataset, and demonstrates improvement of accuracy over other approaches. My findings suggest that the proposed approach has potential to enhance brain tumor segmentation using multi-modal MRI data, contributing to better understanding and diagnosis of brain diseases. This work highlights the importance of combining multiple imaging modalities and incorporating attention mechanisms for improved accuracy in brain tumor segmentation.

연구 동기 및 목표

• 3D U-Net 아키텍처에 어텐션 메커니즘을 통합하여 다중 모odal MRI에서 다중 영역 뇌 종양 분할 정확도를 향상시키는 것.
• 학습 중 불필요한 배경 특징을 억제함으로써 계산 오버헤드를 줄이고 일반화 능력을 향상시키는 것.
• T2-FLAIR, T1CE, T2 볼륨을 하나의 입력 텐서로 스택하여 한 번의 분할로 다중 모달 MRI 스캔을 처리할 수 있도록 하는 것.
• α=0.7, β=0.3로 설정된 Tversky 손실 함수를 사용하여 종양 분할에서의 클래스 불균형 문제를 해결하는 것.
• 3D 어텐션 게이트가 작은 복잡한 종양 하위 영역의 특징 표현을 향상시키는 데 효과적인지 입증하는 것.

제안 방법

• 모델는 특징 추출을 위한 수축 경로와 정밀 분할을 위한 확장 경로를 갖는 3D U-Net이며, 스위프 커넥션으로 연결된다.
• 스위프 커넥션에 3D 어テン션 게이트를 삽입하여 채널 및 공간 중요도에 따라 특징 맵을 동적으로 가중한다. 이는 쿼리(g) 및 키(x) 벡터를 사용한다.
• 어텐션 게이트는 3D 컨볼루션, ReLU 활성화 함수, 그리고 시그모이드 함수를 통해 학습 가능한 어텐션 맵을 계산하여 특징 반응을 조절한다.
• 입력은 세 가지 비표준 MRI 모달리티(T2-FLAIR, T1CE, T2)를 하나의 입력 볼륨으로 융합한 4차원 텐서 형태(3M × L × W × S)이다. 이는 더 풍부한 특징 표현을 가능하게 한다.
• 네트워크는 Tversky 손실(α=0.7, β=0.3)을 사용하여 정밀도와 재현율을 균형 있게 유지하고, BraTS 2021에서 Dice 계수를 평가 지표로 사용하여 훈련한다.
• 두 번의 훈련 시도를 실시: 하나는 ReLU, 다른 하나는 LeakyReLU를 사용하여 죽은 ReLU 문제를 완화하고, LeakyReLU가 더 우수한 성능을 보였다.

Figure 1: 3D Attention U-Net Architecture

실험 결과

연구 질문

• RQ1 기본 3D U-Net 및 기타 최고 성능 모델에 비해 3D 어텐션 기반 U-Net이 다중 클래스 뇌 종양 분할 정확도를 향상시킬 수 있는가?
• RQ2 세 가지 MRI 모달리티(T2-FLAIR, T1CE, T2)를 하나의 입력 볼륨으로 융합함으로써 특징 표현과 분할 성능가 향상되는가?
• RQ3 3D 어텐션 메커니즘이 훈련 및 추론 중 불필요한 배경 특징을 억제함으로써 계산 부담을 어느 정도 줄이는가?
• RQ4 α=0.7, β=0.3로 설정된 Tversky 손실 함수는 뇌 종양 분할 작업에서의 클래스 불균형 문제를 어느 정도 효과적으로 다루는가?
• RQ5 ReLU를 LeakyReLU로 교체함으로써 이 3D 의료 영상 분할 설정에서 모델 수렴 및 최종 분할 성능이 향상되는가?

주요 결과

• 제안된 3D 어텐션 U-Net은 에포크 127에서 BraTS 2021 검증 세트에서 0.9562의 Dice 계수를 기록하여 베이스라인 ReLU 모델을 능가했다.
• 테스트 세트에서 모델은 0.9864의 Dice 계수를 달성하여 동일한 데이터셋에서 표준 3D U-Net이 보고한 0.92를 크게 초월했다.
• LeakyReLU 활성화 함수를 사용한 모델은 ReLU보다 더 우수한 성능을 보였으며, 테스트 세트에서 Tversky 손실 0.0136, Dice 계수 0.9864를 기록했다.
• 어텐션 메커니즘이 배경 조직을 효과적으로 억제하고 종양 하위 영역(괴사, 부종, 강화 종양)에 대한 집중도를 향상시켰다.
• 스택된 다중 모달 입력(T2-FLAIR, T1CE, T2)은 더 풍부한 공간적 및 강도 기반 특징을 제공하여 단일 프로퍼그레이션 내에서 더 정확한 분할을 가능하게 했다.
• 3D 어텐션 게이트의 통합은 업샘플링 과정에서 관련 특징을 우선시함으로써 모델 일반화 능력을 향상시키고 계산 부하를 감소시켰다.

Figure 2: Visual representation of 3D Attention mechanism