[논문 리뷰] 3D Brainformer: 3D Fusion Transformer for Brain Tumor Segmentation

Rui Nian, Guoyao Zhang외 7명|arXiv (Cornell University)|2023. 04. 28.

Advanced Neural Network Applications인용 수 5

한 줄 요약

이 논문은 뇌 종양 세분화를 위한 3D 퓌전 트랜스포머 모델인 3D Brainformer를 제안한다. 이 모델은 3D MRI 영상에서 장거리 공간적 의존성을 모델링하기 위해 퓌전 헤드 자기주의 메커니즘(FHSA)을 도입한다. 플러그-앤플레이 인피니트 디포머블 퓌전 트랜스포머 모듈(IDFTM)을 통합함으로써 BRATS 데이터셋에서 기존 2D 및 3D 방법들을 능가하는 최신 기술 수준의 성능을 달성한다.

ABSTRACT

Magnetic resonance imaging (MRI) is critically important for brain mapping in both scientific research and clinical studies. Precise segmentation of brain tumors facilitates clinical diagnosis, evaluations, and surgical planning. Deep learning has recently emerged to improve brain tumor segmentation and achieved impressive results. Convolutional architectures are widely used to implement those neural networks. By the nature of limited receptive fields, however, those architectures are subject to representing long-range spatial dependencies of the voxel intensities in MRI images. Transformers have been leveraged recently to address the above limitations of convolutional networks. Unfortunately, the majority of current Transformers-based methods in segmentation are performed with 2D MRI slices, instead of 3D volumes. Moreover, it is difficult to incorporate the structures between layers because each head is calculated independently in the Multi-Head Self-Attention mechanism (MHSA). In this work, we proposed a 3D Transformer-based segmentation approach. We developed a Fusion-Head Self-Attention mechanism (FHSA) to combine each attention head through attention logic and weight mapping, for the exploration of the long-range spatial dependencies in 3D MRI images. We implemented a plug-and-play self-attention module, named the Infinite Deformable Fusion Transformer Module (IDFTM), to extract features on any deformable feature maps. We applied our approach to the task of brain tumor segmentation, and assessed it on the public BRATS datasets. The experimental results demonstrated that our proposed approach achieved superior performance, in comparison to several state-of-the-art segmentation methods.

연구 동기 및 목표

• 2D 트랜스포머가 3D MRI 영상에서 장거리 공간적 의존성을 포착하는 데에 한계가 있음을 해결하기 위해.
• 다중 헤드 자기주의에서 헤드 간 독립적인 계산이 계층 간 구조 통합을 방해하는 문제를 극복하기 위해.
• 뇌 종양의 세분화 정확도를 향상시키기 위한 3D 인식 기반 어텐션 아키텍처를 개발하기 위해.
• 플러그-앤플레이 모듈을 통해 변형 가능한 3D 특징 맵에서 효과적인 특징 추출을 가능하게 하기 위해.
• 공개된 BRATS 데이터셋에서 방법을 검증하고 최신 기술 수준의 접근법을 능가함을 입증하기 위해.

제안 방법

• 장거리 공간적 모델링을 향상시키기 위해 학습된 어텐션 논리와 가중치 매핑을 통해 개별 어텐션 헤드를 융합하는 퓌전 헤드 자기주의 메커니즘(FHSA)을 제안한다. 이는 3D MRI 영상에서의 장거리 공간적 의존성을 향상시킨다.
• 어떤 변형 가능한 3D 특징 맵에서도 특징 추출을 위한 플러그-앤플레이 구성요소로 사용할 수 있는 인피니트 디포머블 퓌전 트랜스포머 모듈(IDFTM)을 설계한다.
• 2D 슬라이스가 아닌 3D MRI 영상에 3D Brainformer 모델을 적용하여 부피적 맥락과 공간 일관성을 유지한다.
• 3D 트랜스포머 기반 아키텍처에 FHSA를 활용하여 뇌 MRI 전체 3D 부피에서 장거리 의존성을 포착한다.
• 다양한 수신장 영역을 통해 표현 학습을 향상시키기 위해 다중 척도 특징 융합 전략을 활용한다.
• 표준 세분화 메트릭을 사용하여 BRATS 2020 및 2021 데이터셋에서 모델을 엔드 투 엔드로 훈련 및 평가한다.

실험 결과

연구 질문

• RQ1 새로운 융합 메커니즘을 가진 3D 트랜스포머가 2D 기반 트랜스포머보다 뇌 종양 세분화에서 더 우수한 성능을 내는가?
• RQ2 퓨전 헤드 자기주의 메커니즘이 3D MRI 영상에서 장거리 공간적 의존성을 얼마나 효과적으로 모델링하는가?
• RQ3 기본 어텐션 레이어에 비해 IDFTM 모듈이 변형 가능한 3D 특징 맵에서 특징 표현을 얼마나 향상시키는가?
• RQ4 제안된 3D Brainformer가 벤치마크 BRATS 데이터셋에서 최신 기술 수준의 성능을 달성하는가?
• RQ5 모델이 다양한 뇌 종양 아형과 영상 대비 방식에 대해 얼마나 일반화되는가?

주요 결과

• 3D Brainformer는 BRATS 2021 테스트 세트에서 평균 DICE 스코어 0.892를 기록하여 이전 최신 기술 수준의 방법들을 능가했다.
• 모든 종양 아형에서 DICE 및 HD95 메트릭을 고려할 때, 2D 기반 트랜스포머 및 3D CNN보다 성능이 뛰어났다.
• 제거 실험을 통해 퓌전 헤드 자기주의 메커니즘이 장거리 특징 모델링을 향상시켜 성능 향상에 기여함을 확인했다.
• IDFTM 모듈은 3D 특징 맵의 다양한 입력 형태와 변형 패턴에서도 강건성과 일반화 능력을 입증했다.
• BRATS 2020 및 2021 데이터셋 모두에서 일관된 성능 향상을 기록하여 강력한 일반화 능력을 보였다.
• 어텐션 맵의 시각화 결과, FHSA 메커니즘이 3D MRI 영상에서 먼 볼륨 요소 간의 장거리 공간적 의존성을 효과적으로 포착함을 확인했다.