[논문 리뷰] A Volumetric Transformer for Accurate 3D Tumor Segmentation

Himashi Peiris, Munawar Hayat외 3명|arXiv (Cornell University)|2021. 11. 26.

Advanced Neural Network Applications참고 문헌 64인용 수 25

한 줄 요약

이 논문은 뇌종양 세그멘테이션에서 최고 성능을 기록한 VT-UNet을 제안한다. 이는 U-형 인코더-디코더 아키텍처를 가진 계산적으로 효율적인 3D 볼륨 트랜스포머로, 인코더에서 이중 자기주의성 어텐션 레이어를 사용해 局소 및 전역 공간적 특징을 포착하고, 디코더에서 병렬 이동 윈도우 기반 자기 및 상호 어텐션과 푸리에 위치 인코딩을 활용해 종양 경계를 정밀하게 보정한다. 모델은 BraTS 2021, 췌장 및 간 데이터셋에서 최고 성능을 달성하여 강력한 일반화 능력과 데이터 오염에 대한 강인성을 입증한다.

ABSTRACT

This paper presents a Transformer architecture for volumetric medical image segmentation. Designing a computationally efficient Transformer architecture for volumetric segmentation is a challenging task. It requires keeping a complex balance in encoding local and global spatial cues, and preserving information along all axes of the volumetric data. The proposed volumetric Transformer has a U-shaped encoder-decoder design that processes the input voxels in their entirety. Our encoder has two consecutive self-attention layers to simultaneously encode local and global cues, and our decoder has novel parallel shifted window based self and cross attention blocks to capture fine details for boundary refinement by subsuming Fourier position encoding. Our proposed design choices result in a computationally efficient architecture, which demonstrates promising results on Brain Tumor Segmentation (BraTS) 2021, and Medical Segmentation Decathlon (Pancreas and Liver) datasets for tumor segmentation. We further show that the representations learned by our model transfer better across-datasets and are robust against data corruptions. \href{https://github.com/himashi92/VT-UNet}{Our code implementation is publicly available}.

연구 동기 및 목표

• 3D 볼륨 의료 영상 세그멘테이션을 위한 계산적으로 효율적인 트랜스포머 설계의 과제를 해결하기 위해.
• 볼륨 데이터의 세 축 모두에서 국소 및 전역 공간적 의존성을 균형 있게 유지하기 위해.
• 새로운 어텐션 메커니즘을 활용해 종양 세그멘테이션의 경계 보정을 향상시키기 위해.
• 다양한 의료 영상 데이터셋 간의 표현 이동성과 데이터 오염에 대한 강인성을 향상시키기 위해.

제안 방법

• 모델은 패치 없이 전체 3D 볼륨 입력을 처리하는 U-형 인코더-디코더 아키텍처를 채택한다.
• 인코더는 국소 및 전역 공간적 맥락을 동시에 인코딩하기 위해 연속된 두 개의 자기주의성 어텐션 레이어를 사용한다.
• 디코더는 세부 정보 학습을 향상시키기 위해 병렬 이동 윈도우 기반 자기 및 상호 어텐션 블록을 도입한다.
• 계산 비용 증가 없이 위치 모델링을 향상시키기 위해 푸리에 위치 인코딩을 통합한다.
• 모든 축을 따라 공간 정보를 유지하면서도 효율성을 유지하도록 아키텍처를 설계한다.
• 표준 세그멘테이션 손실 함수를 사용해 BraTS 2021, 췌장 및 간 데이터셋에서 엔드 투 엔드로 모델을 훈련시킨다.

실험 결과

연구 질문

• RQ1 트랜스포머 기반 아키텍처가 계산적으로 효율적인 同시에 3D 종양 세그멘테이션에서 높은 정확도를 달성할 수 있는가?
• RQ2 이 모델은 볼륨 의료 영상에서 국소 및 전역 공간적 의존성을 얼마나 효과적으로 포착할 수 있는가?
• RQ3 제안된 어텐션 메커니즘이 종양 세그멘테이션의 경계 분할에 얼마나 향상시키는가?
• RQ4 학습된 표현은 다양한 의료 영상 데이터셋 간에 얼마나 잘 일반화되며, 데이터 오염 상황에서도 얼마나 잘 작동하는가?
• RQ5 푸리에 위치 인코딩 통합이 계산 오버헤드 없이 성능 향상에 기여하는가?

주요 결과

• 제안된 VT-UNet은 뇌종양 세그멘테이션에서 BraTS 2021 데이터셋에서 최고 성능을 기록한다.
• 모델은 의료 세그멘테이션 디카테론 데이터셋에서 췌장 및 간 종양 세그멘테이션에서 향상된 성능을 보이며 강력한 일반화 능력을 입증한다.
• 기본 모델 대비 학습된 표현이 더 나은 이동성(transferability)을 보인다.
• 다양한 데이터 오염에 대해 강인성을 보이며 신뢰할 수 있는 특징 학습 능력을 입증한다.
• 병렬 이동 윈도우 어텐션과 푸리에 위치 인코딩의 조합은 유의미한 계산 비용 증가 없이 경계 보정을 향상시킨다.
• 엔드 투 엔드로 전체 3D 볼륨을 처리하면서도 계산 효율성을 유지한다.