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

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

Medical Image Segmentation Techniques인용 수 9

한 줄 요약

이 논문은 3D MRI 영상 전체를 잘라내지 않고 처리하는 순수 볼륨형 트랜스포머 아키텍처인 VT-UNet을 제안한다. 인코더는 계층적 자기주의를, 디코더는 병렬 자기주의 및 상호주의를 사용하며, 푸리에 위치 인코딩을 적용한다. BraTS 데이터셋에서 최신 기술 수준(SOTA) 성능을 달성하면서도 MRI 아티팩트에 대한 강건성은 향상시키며, 계산 복잡도는 낮게 유지한다.

ABSTRACT

We propose a Transformer architecture for volumetric segmentation, a challenging task that requires keeping a complex balance in encoding local and global spatial cues, and preserving information along all axes of the volume. Encoder of the proposed design benefits from self-attention mechanism to simultaneously encode local and global cues, while the decoder employs a parallel self and cross attention formulation to capture fine details for boundary refinement. Empirically, we show that the proposed design choices result in a computationally efficient model, with competitive and promising results on the Medical Segmentation Decathlon (MSD) brain tumor segmentation (BraTS) Task. We further show that the representations learned by our model are robust against data corruptions. \href{https://github.com/himashi92/VT-UNet}{Our code implementation is publicly available}.

연구 동기 및 목표

• 2D 슬라이스 기반 처리를 방지함으로써 3D 의료 영상 분할에서 전체 볼륨 공간적 맥락을 유지하는 데 도전한다.
• 계산 효율성을 유지하면서도 모든 3D 축을 따라 장거리 의존성을 포착할 수 있는 순수 트랜스포머 기반 UNet 아키텍처를 설계한다.
• 디코더에서 병렬 주의 메커니즘을 통해 특징 정제를 향상시켜 종양 분할의 경계 정의를 개선한다.
• 운동, 가로등, 스팀트 아티팩트와 같은 일반적인 MRI 아티팩트에 대한 모델 강건성을 향상시킨다.
• 3D에서 자기주의 메커니즘이 CNN 기반 및 하이브리드 모델보다 정확도와 강건성 면에서 우월할 수 있음을 입증한다.

제안 방법

• 인코더는 3D 볼륨에서 국소적이고 전역적 맥락 신호를 동시에 포착하기 위해 계층적이고 이동된 창 방식의 두 개의 연속 자기주의 레이어를 사용한다.
• 디코더는 공유 쿼리 프로젝션을 통해 병렬 자기주의 및 상호주의 메커니즘을 적용하여 업샘플링 과정에서 전역 맥락을 유지한다.
• 푸리에 위치 인코딩을 디코딩 과정에 통합하여 공간적 인덕티브 바이어스를 향상시키고 특징 표현을 개선한다.
• 병렬 주의 헤드에서의 특징를 융합하기 위해 볼록 조합 전략을 적용하여 표현 품질을 향상시킨다.
• 학습은 AdamW를 사용하며, 초기 학습률은 1e-4로 설정하고, 데이터 증강(회전, 노이즈, 블러, 감마 조정)을 적용하며, ImageNet-22K에서 미리 훈련된 Swin-T 가중치를 초기화에 사용한다.
• 아키텍처는 2D 슬라이스 분해 없이 MSD BraTS 데이터셋의 전체 3D MRI 볼륨에서 평가된다.

실험 결과

연구 질문

• RQ1 순수 볼륨형 트랜스포머 아키텍처가 저비용 계산 복잡도를 유지하면서도 CNN 기반 및 하이브리드 모델보다 3D 종양 분할에서 승승을 거두는가?
• RQ2 디코더에서 병렬 자기주의 및 상호주의가 3D 분할에서 경계 정확도와 특징 정제를 어떻게 향상시키는가?
• RQ3 푸리에 위치 인코딩이 3D 트랜스포머 기반 분할 네트워크에서 공간 모델링을 얼마나 향상시키는가?
• RQ4 제안된 아키텍처가 운동, 가로등, 스파이크 아티팩트와 같은 일반적인 MRI 아티팩트에 대해 향상된 강건성을 보여주는가?
• RQ5 완전히 주의 기반 3D UNet 설계가 컨volutional 인덕티브 바이어스에 의존하지 않고도 최신 기술 수준의 성능을 달성할 수 있는가?

주요 결과

• VT-UNet-B는 전체 종양(WT) 클래스에서 91.9의 최고 DSC(Dice Similarity Coefficient)를 기록하여 BraTS 데이터셋에서 모든 SOTA 방법을 능가했다.
• WT 클래스에 대해 Hausdorff Distance(HD95)가 3.43으로, 비교된 모든 방법 중에서 가장 낮아 경계 정확도가 뛰어났다.
• VT-UNet는 MRI 아티팩트에 대한 뛰어난 강건성을 보였다: 스파이크 아티팩트에서 DSC 86.6, 운동 아티팩트에서 DSC 85.8을 기록하여, 모든 손상 유형에서 nnFormer를 능가했다.
• 아블레이션 연구 결과, 푸리에 위치 인코딩, 볼록 조합, 병렬 주의 메커니즘의 조합이 성능 향상에 크게 기여했으며, 이 모든 구성 요소를 제거할 경우 DSC가 1.78점 감소함을 확인했다.
• VT-UNet-S와 VT-UNet-B는 각각 3.84 및 3.43 GFLOPs의 낮은 FLOPs를 기록하여, nnFormer(4.05 GFLOPs)를 포함한 대부분의 SOTA 모델보다 낮은 계산 복잡도를 유지하면서도 더 높은 정확도를 확보했다.
• 정성적 결과에서는 VT-UNet이 더 선명한 종양 경계와 복잡한 구조물(예: 종양 주변 부종, 강화 종양)의 더 정확한 분할을 제공하는 것으로 나타났다.