[논문 리뷰] TransDeepLab: Convolution-Free Transformer-based DeepLab v3+ for Medical Image Segmentation

Reza Azad, Moein Heidari외 5명|arXiv (Cornell University)|2022. 01. 01.

Advanced Neural Network Applications인용 수 5

한 줄 요약

TransDeepLab는 병변 이미지 분할을 위한 컨볼루션 없는 순수 트랜스포머 기반의 DeepLabv3+ 아키텍처를 제안하며, 이는 이동된 창문을 가진 계층적 Swin-Transformer와 새로운 다중 척도 특징 융합 모듈(SSPP)을 활용하여 장거리 의존성을 모델링하고 경계 국소화를 향상시킨다. 이는 모델 복잡도를 크게 줄인 21.14M 파라미터로 다수의 병변 이미지 기준에서 최신 기술 수준의 성능을 달성하며, CNN 기반 및 하이브리드 트랜스포머-CNN 모델을 모두 능가한다.

ABSTRACT

Convolutional neural networks (CNNs) have been the de facto standard in a diverse set of computer vision tasks for many years. Especially, deep neural networks based on seminal architectures such as U-shaped models with skip-connections or atrous convolution with pyramid pooling have been tailored to a wide range of medical image analysis tasks. The main advantage of such architectures is that they are prone to detaining versatile local features. However, as a general consensus, CNNs fail to capture long-range dependencies and spatial correlations due to the intrinsic property of confined receptive field size of convolution operations. Alternatively, Transformer, profiting from global information modelling that stems from the self-attention mechanism, has recently attained remarkable performance in natural language processing and computer vision. Nevertheless, previous studies prove that both local and global features are critical for a deep model in dense prediction, such as segmenting complicated structures with disparate shapes and configurations. To this end, this paper proposes TransDeepLab, a novel DeepLab-like pure Transformer for medical image segmentation. Specifically, we exploit hierarchical Swin-Transformer with shifted windows to extend the DeepLabv3 and model the Atrous Spatial Pyramid Pooling (ASPP) module. A thorough search of the relevant literature yielded that we are the first to model the seminal DeepLab model with a pure Transformer-based model. Extensive experiments on various medical image segmentation tasks verify that our approach performs superior or on par with most contemporary works on an amalgamation of Vision Transformer and CNN-based methods, along with a significant reduction of model complexity. The codes and trained models are publicly available at https://github.com/rezazad68/transdeeplab

연구 동기 및 목표

• 병변 이미지 분할에서 CNN의 장거리 공간적 의존성과 맥락적 상관관계를 포착하는 데 한계가 있음을 해결하기 위해.
• 밀도 예측 작업에서 표준 비전 트랜스포머의 높은 계산 비용과 낮은 수준의 특징 학습 능력을 극복하기 위해.
• 기존 DeepLabv3+ 아키텍처의 순수 트랜스포머 기반 대체품을 개발하여 성능은 유지하면서 모델 복잡도를 감소시키기 위해.
• 트랜스포머 인코더-디코더 프레임워크 내에서 새로운 SSPP 모듈을 통해 계층적 특징 표현과 다중 척도 맥락 모델링을 통합하기 위해.

제안 방법

• 딥랩v3+의 CNN 기반 인코더를 이동된 창문을 사용한 자기주도 어텐션을 통해 국소적이고 전역적인 특징 학습을 가능하게 하는 계층적 Swin-Transformer 인코더로 대체한다.
• Swin-Transformer의 다양한 스테이지 간의 교차 어텐션을 통한 다중 척도 특징 융합 메커니즘을 도입하여 맥락 표현을 향상시킨다.
• 다양한 속도에서 Swin-Transformer의 여러 스테이지에서 특징을 집계하는 새로운 Atrous Spatial Pyramid Pooling(ASPP) 변형인 SSPP를 활용하여 다중 척도 맥락을 포착한다.
• 경량 디코더 헤드와 스킵 연결을 사용하여 분할 맵을 정밀하게 개선하고 세밀한 세부 정보를 유지한다.
• Swin-Transformer의 계층적이고 국소적으로 연결된 설계를 활용하여 이차형 어텐션 복잡도를 감소시켜 학습 효율성을 최적화한다.
• 다양한 병변 이미지 데이터셋에서 교차 엔트로피 손실과 Dice 손실을 사용하여 엔드 투 엔드로 학습하며, 구성 요소 기여도를 검증하기 위한 추론 분석을 실시한다.

실험 결과

연구 질문

• RQ1 순수 트랜스포머 기반 아키텍처가 모델 복잡도를 줄이며 병변 이미지 분할에서 CNN 기반 모델을 능가할 수 있는가?
• RQ2 수정된 ASPP 모듈(SSPP)을 갖춘 계층적 Swin-Transformer와의 통합이 병변 이미지 분할에서 다중 척도 맥락 모델링을 어떻게 향상시키는가?
• RQ3 트랜스포머 스테이지 간의 교차 어텐션 기반 특징 융합 전략이 간단한 연결 또는 원소별 연산 대비 장거리 의존성 학습을 어떻게 향상시키는가?
• RQ4 딥랩v3+의 CNN 인코더를 트랜스포머 아키텍처로 대체함으로써 복잡한 해부학적 구조에서 분할 정확도가 얼마나 향상되는가?
• RQ5 제안된 방법은 하이브리드 트랜스포머-CNN 모델과 기존의 CNN 기반 기준 모델에 비해 효율성과 성능 면에서 어떻게 비교되는가?

주요 결과

• TransDeepLab는 Synapse 데이터셋에서 평균 Dice 스코어 85.68을 기록하여 베이스라인 CNN 기반 인코더(75.89)를 능가하고, 여러 기준에서 최신 기술 수준의 방법과 동등하거나 이를 초월한다.
• 모델 파라미터 수를 21.14M으로 크게 줄여, Swin-UNet(27.17M), DeepLabv3(54.70M), Trans-UNet(105M)보다 뛰어난 모델 효율성을 입증한다.
• 추론 분석 결과, 교차 어텐션 기반 특징 융합 전략은 기본 연결 융합 대비 Dice 스코어에서 4.27% 향상된 것으로 확인되었다.
• 이중 수준의 SSPP 모듈을 사용할 경우 최고의 Dice 스코어(80.16)를 기록하여 다중 척도 맥락 포착과 계산 비용 사이의 최적 균형을 이룬다.
• 시각화 결과, 강화된 장거리 의존성 모델링 덕분에 복잡한 해부학적 영역에서 더 정확한 경계 예측을 제공함을 확인할 수 있었다.
• 이중 수준의 SSPP를 사용한 경우 Synapse 데이터셋에서 Hausdorff 거리가 21.25로 나타나, 다른 설정에 비해 국소화 정밀도가 향상됨을 나타낸다.