[논문 리뷰] The Fully Convolutional Transformer for Medical Image Segmentation

Athanasios Tragakis, Chaitanya Kaul외 2명|arXiv (Cornell University)|2022. 06. 01.

Radiomics and Machine Learning in Medical Imaging인용 수 7

한 줄 요약

이 논문은 의료 영상 분할을 위한 새로운, 컴act한, 완전 컨볼루션형 트랜스포머 아키텍처인 풀리 컨볼루션형 트랜스포머(Fully Convolutional Transformer, FCT)를 제안한다. FCT는 장거리 의존성과 계층적 특징을 포착하기 위해 컨볼루션형 어텐션과 다중 해상도 확장 컨볼루션형 워이드포커스 모듈을 결합한다. FCT는 전이 학습 없이도 여러 데이터셋에서 최신 기준 성능을 달성하며, 기존 모델보다 딱 4.4% 높은 딱지 스코어를 기록한다. 또한 nnUNet나 TransUNet와 같은 경쟁 모델 대비 최대 5배 적은 파라미터를 사용한다.

ABSTRACT

We propose a novel transformer model, capable of segmenting medical images of varying modalities. Challenges posed by the fine grained nature of medical image analysis mean that the adaptation of the transformer for their analysis is still at nascent stages. The overwhelming success of the UNet lay in its ability to appreciate the fine-grained nature of the segmentation task, an ability which existing transformer based models do not currently posses. To address this shortcoming, we propose The Fully Convolutional Transformer (FCT), which builds on the proven ability of Convolutional Neural Networks to learn effective image representations, and combines them with the ability of Transformers to effectively capture long-term dependencies in its inputs. The FCT is the first fully convolutional Transformer model in medical imaging literature. It processes its input in two stages, where first, it learns to extract long range semantic dependencies from the input image, and then learns to capture hierarchical global attributes from the features. FCT is compact, accurate and robust. Our results show that it outperforms all existing transformer architectures by large margins across multiple medical image segmentation datasets of varying data modalities without the need for any pre-training. FCT outperforms its immediate competitor on the ACDC dataset by 1.3%, on the Synapse dataset by 4.4%, on the Spleen dataset by 1.2% and on ISIC 2017 dataset by 1.1% on the dice metric, with up to five times fewer parameters. Our code, environments and models will be available via GitHub.

연구 동기 및 목표

• 기존 트랜스포머 기반 모델이 UNet의 세밀한 특징 학습 능력을 갖추지 못한 점을 해결하기 위해.
• 의료 영상에서 장거리 의미론적 의존성과 계층적 글로벌 맥락을 효과적으로 포착할 수 있는 완전 컨볼루션형 트랜스포머를 개발하기 위해.
• 전이 학습 없이도 최신 기준 성능을 초월하는 경량형, 정확도 높고 견고한 모델을 만들기 위해.
• 위치 인코딩과 패치 토큰화 없이도 컨볼루션 연산과 자기어텐션 메커니즘을 효과적으로 통합하기 위해.

제안 방법

• FCT는 새로운 완전 컨볼루션형 트랜스포머 레이어를 사용하며, 이는 두 가지 주요 구성 요소로 이루어져 있다: 컨볼루션형 어텐션 모듈과 워이드포커스 모듈.
• 컨볼루션형 어텐션 모듈은 깊이 분리형 컨볼루션을 사용하여 겹치는 공간 인식 패치를 생성함으로써, 표준 패치 임베딩을 대체하고 위치 인코딩이 필요 없도록 한다.
• 워이드포커스 모듈은 세 개의 브랜치에서 증가하는 확장률을 가진 다중 해상도 확장 컨볼루션을 적용하여 계층적 국소-글로벌 특징을 포착한다.
• 모델은 두 단계로 입력을 처리한다: 먼저 컨볼루션형 어텐션 모듈을 통해 장거리 의미론적 의존성을 학습하고, 그 다음 워이드포커스 모듈을 통해 계층적 글로벌 표현을 구축한다.
• 엔코더에서 디코더로의 스킵 커넥션은 UNet 설계를 따르며, 공간적 세부 정보를 유지하고 특징 전파를 향상시킨다.
• 표준 데이터 증강 기법을 사용해 엔드 투 엔드로 학습하고, 겹치는 예측을 평균화하기 위해 후처리를 수행하는 슬라이딩 윈도우 추론을 사용해 추론을 수행한다.

실험 결과

연구 질문

• RQ1 전이 학습 없이도 완전 컨볼루션형 트랜스포머 아키텍처가 기존 CNN 및 트랜스포머 기반 모델보다 의료 영상 분할 성능에서 뛰어나게 할 수 있는가?
• RQ2 컨볼루션형 어텐션과 다중 해상도 확장 컨볼루션의 통합이 장거리 및 세밀한 특징을 포착하는 데 성능에 어떤 영향을 미치는가?
• RQ3 의료 영상 분할에서 워이드포커스 모듈의 최적 구성(브랜치 수와 확장률)은 무엇인가?
• RQ4 스킵 커넥션은 제안된 FCT 아키텍처에서 성능 향상에 얼마나 기여하는가?
• RQ5 경량 트랜스포머 기반 모델이 기존 모델보다 훨씬 적은 파라미터로 최신 기준 성능을 달성할 수 있는가?

주요 결과

• ACDC 데이터셋에서 FCT는 전이 학습 없이도 딱지 스코어 1.3% 향상으로 모든 기존 트랜스포머 및 CNN 기반 모델을 압도한다.
• Synapse 데이터셋에서 FCT는 직접 경쟁자보다 딱지 스코어 4.4% 향상되었고, nnUNet 대비 파라미터 수가 5배 적었다.
• 간장 데이터셋에서 FCT는 다음으로 우수한 모델 대비 딱지 스코어 1.2% 향상되어 다양한 해부학적 구조에서의 견고함을 입증했다.
• ISIC 2017 데이터셋에서 FCT는 가장 가까운 경쟁자보다 딱지 스코어 1.1% 높게 기록하여 다양한 영상 모odalities에서의 일반화 능력을 확인했다.
• ACDC Post-2017-MICCAI-Challenge 온라인 테스트 세트에서 FCT는 단 3170만 개의 파라미터로 대규모 앙상블 모델과 nnUNet를 모두 능가하는 새로운 최신 기준 성능을 확립했다.
• 절단 실험 결과, 스킵 커넥션과 최적의 워이드포커스 모듈 구성(선형 증가하는 확장률을 가진 세 개의 브랜치)이 최고 성능을 내기 위해 필수적임을 확인했다.