[논문 리뷰] TransBTS: Multimodal Brain Tumor Segmentation Using Transformer

Wenxuan Wang, Chen Chen외 4명|arXiv (Cornell University)|2021. 03. 07.

Advanced Neural Network Applications참고 문헌 23인용 수 123

한 줄 요약

TransBTS는 3D CNN 인코더–Transformer 아키텍처를 도입하여 3D MRI에서 로컬 및 글로벌 특징을 모델링하고 BraTS 2019 및 2020에서 최첨단 방법과 경쟁력 있는 성능을 달성합니다.

ABSTRACT

Transformer, which can benefit from global (long-range) information modeling using self-attention mechanisms, has been successful in natural language processing and 2D image classification recently. However, both local and global features are crucial for dense prediction tasks, especially for 3D medical image segmentation. In this paper, we for the first time exploit Transformer in 3D CNN for MRI Brain Tumor Segmentation and propose a novel network named TransBTS based on the encoder-decoder structure. To capture the local 3D context information, the encoder first utilizes 3D CNN to extract the volumetric spatial feature maps. Meanwhile, the feature maps are reformed elaborately for tokens that are fed into Transformer for global feature modeling. The decoder leverages the features embedded by Transformer and performs progressive upsampling to predict the detailed segmentation map. Extensive experimental results on both BraTS 2019 and 2020 datasets show that TransBTS achieves comparable or higher results than previous state-of-the-art 3D methods for brain tumor segmentation on 3D MRI scans. The source code is available at https://github.com/Wenxuan-1119/TransBTS

연구 동기 및 목표

• 부피 MRI에서 로컬 3D 컨텍스트와 글로벌 장거리 의존성을 모두 모델링하는 동기를 제시한다.
• 뇌종양 분할을 위한 3D CNN–Transformer 인코더–디코더 네트워크(TransBTS)를 제안한다.
• BraTS 2019와 2020에서의 효과를 입증하고 해부 실험(ablation)을 통해 아키텍처 선택을 분석한다.

제안 방법

• 다중 모달 MRI로부터 컴팩트한 로컬 3D 특징 맵을 추출하기 위해 3D CNN 인코더를 사용한다.
• 인코더 특징을 학습 가능한 위치 임베딩을 추가하여 Transformer 호환 토큰으로 펼치고 투영한다.
• 토큰을 Transformer 인코더를 통해 처리하여 공간 및 깊이 차원에 걸친 글로벌 장거리 의존성을 캡처한다.
• Transformer 출력물을 다시 4D 특징 맵으로 매핑하고 진행식 업샘플링 및 스킵 연결이 있는 3D CNN 디코더를 적용하여 전체 해상도 분할을 수행한다.
• 데이터 증강 및 Dice 기반 손실과 함께 BraTS 데이터에서 처음부터 학습하며, 선택적으로 테스트 시 증강(TTA)을 적용한다.
• BraTS 2019 및 2020 데이터셋에서 최첨단 3D 분할 방법과 비교한다.

실험 결과

연구 질문

• RQ1 Transformer 기반 글로벌 컨텍스트 모델을 3D CNN 프레임워크에 효과적으로 통합하여 부피 뇌종양 분할을 달성할 수 있는가?
• RQ2 로컬 3D 컨텍스트와 글로벌 의존성을 공동으로 모델링하면 BraTS 데이터셋에서 Dice 점수와 Hausdorff 거리의 개선이 있는가?
• RQ3 Transformer 깊이, 임베딩 차원, 토큰 시퀀스 길이, 스킵 연결 배치가 분할 성능에 미치는 영향은 무엇인가?

주요 결과

데이터 세트	방법	ET Dice (%)	WT Dice (%)	TC Dice (%)	ET HD (mm)	WT HD (mm)	TC HD (mm)
BraTS2019	3D U-Net	—	—	—	—	—	—
BraTS2019	V-Net	—	—	—	—	—	—
BraTS2019	KiU-Net	—	—	—	—	—	—
BraTS2019	Attention U-Net	—	—	—	—	—	—
BraTS2019	Wang et al.	—	—	—	—	—	—
BraTS2019	Li et al.	—	—	—	—	—	—
BraTS2019	Frey et al.	—	—	—	—	—	—
BraTS2019	Myronenko et al.	—	—	—	—	—	—
BraTS2019	TransBTS w/o TTA	ET 78.36	WT 88.89	TC 81.41	ET 5.908	WT 7.599	TC 7.584
BraTS2019	TransBTS w/ TTA	ET 78.93	WT 90.00	TC 81.94	ET 3.736	WT 5.644	TC 6.049
BraTS2020	TransBTS w/o TTA	ET 78.50	WT 89.00	TC 81.36	ET 16.716	WT 6.469	TC 10.468
BraTS2020	TransBTS w/ TTA	ET 78.73	WT 90.09	TC 81.73	ET 17.947	WT 4.964	TC 9.769

• TransBTS는 BraTS 2019 검증에서 여러 3D 기초선과 비교해 경쟁력 있거나 우수한 Dice 점수와 더 낮은 Hausdorff 거리를 달성한다(예: ET 78.93, WT 90.00, TC 81.94 with TTA).
• TransBTS w/ TTA는 TTA를 적용했을 때 Dice에서 w/o TTA보다 ET, WT, TC에서 우수하며 HD가 크게 감소한다(예: ET HD 3.736 mm).
• BraTS 2020 검증에서 TransBTS는 w/o TTA 및 w/ TTA 모두 ET과 TC에서 약 78.5–78.7의 Dice 점수, WT에서 약 89–90의 점수를 달성하고 HD는 여러 기초선 대비 향상되며(예: ET HD ~16.7–17.9 mm).
• TransBTS의 경량 변형(더 적은 Transformer 레이어와 더 작은 FFN)은 여전히 강한 Dice 점수를 유지하면서 파라미터 및 FLOPs를 각각 약 54%와 38% 감소시킨다.
• 절제 연구는 더 긴 토큰 시퀀스(OS=8)와 U-Net 스타일 블록에서의 스킵 연결이 다른 구성에 비해 성능을 크게 향상시킨다.