[논문 리뷰] Attention Swin U-Net: Cross-Contextual Attention Mechanism for Skin Lesion Segmentation

Ehsan Khodapanah Aghdam, Reza Azad외 2명|arXiv (Cornell University)|2022. 01. 01.

Cutaneous Melanoma Detection and Management인용 수 6

한 줄 요약

이 논문은 피부 병변 분할을 위한 개선된 Swin U-Net 아키텍처인 Att-SwinU-Net을 제안한다. 이는 스킵 커넥션에 새로운 크로스컨텍스추얼 어텐션 메커니즘을 통합하여, 표준 연결 방식을 대체하는 학습 가능한 어텐션 모듈을 도입함으로써 성능을 향상시킨다. 이 모델은 여러 기준 데이터셋에서 최신 기술 수준(SOTA) 성능을 달성하며, CNN 기반 및 Transformer 기반 모델을 모두 능가한다.

ABSTRACT

Melanoma is caused by the abnormal growth of melanocytes in human skin. Like other cancers, this life-threatening skin cancer can be treated with early diagnosis. To support a diagnosis by automatic skin lesion segmentation, several Fully Convolutional Network (FCN) approaches, specifically the U-Net architecture, have been proposed. The U-Net model with a symmetrical architecture has exhibited superior performance in the segmentation task. However, the locality restriction of the convolutional operation incorporated in the U-Net architecture limits its performance in capturing long-range dependency, which is crucial for the segmentation task in medical images. To address this limitation, recently a Transformer based U-Net architecture that replaces the CNN blocks with the Swin Transformer module has been proposed to capture both local and global representation. In this paper, we propose Att-SwinU-Net, an attention-based Swin U-Net extension, for medical image segmentation. In our design, we seek to enhance the feature re-usability of the network by carefully designing the skip connection path. We argue that the classical concatenation operation utilized in the skip connection path can be further improved by incorporating an attention mechanism. By performing a comprehensive ablation study on several skin lesion segmentation datasets, we demonstrate the effectiveness of our proposed attention mechanism.

연구 동기 및 목표

• 의료 영상 분할을 위한 U-Net 아키텍처에서 국소 수용 영역의 한계를 해결하기 위해.
• 표준 연결 방식을 어텐션 기반 메커니즘으로 대체하여 스킵 커넥션 내 특징 재사용성을 향상시키기 위해.
• 복잡한 피부 병변의 분할을 향상시키기 위해 스위프트 U-넷에서 장거리 의존성 모델링을 강화하기 위해.
• 하이브리드 어텐션 메커니즘을 사용하여 피부 병변 분할 기준 데이터셋에서 최신 기술 수준의 성능을 달성하기 위해.

제안 방법

• 스킵 커넥션에 이중 수준의 어텐션 메커니즘을 도입: 첫째, 인코더 블록에서 중요한 토큰을 강조하기 위해 어텐션 가중치를 전달함; 둘째, 서로 다른 스케일 간의 특징을 정교화하기 위해 크로스컨텍스추얼 어텐션 메커니즘을 적용함.
• 제안된 어텐션 모듈을 Swin U-Net의 스킵 커넥션 경로에 통합하여, 표준 연결 방식을 학습 가능한 어텐션 기반 특징 융합 방식으로 대체함.
• 공간적 및 채널 방향의 의존성을 별도로 모델링하는 듀얼 브랜치 어텐션 메커니즘을 도입하여 특징 표현을 향상시킴.
• 인코더 및 디코더 경로에서 국소적 및 전역적 특징 모델링을 위해 스위프트 트랜스포머의 시프트된 윈도우 메커니즘을 활용함.
• 어텐션 모듈의 위치, 입력 해상도, 모델 스케일, 구성 요소 제거 등의 영향을 평가하기 위해 아블레이션 스터디를 수행함.
• 표준 분할 손실 함수를 사용하여 여러 피부 병변 데이터셋에서 모델을 엔드 투 엔드로 훈련함.

실험 결과

연구 질문

• RQ1 스킵 커넥션에 어텐션 메커니즘이 도입될 경우, U-Net 아키텍처에서 특징 재사용성과 분할 정확도가 향상되는가?
• RQ2 제안된 크로스컨텍스추얼 어텐션 메커니즘은 스킵 커넥션에서 표준 연결 방식보다 어떻게 비교되는가?
• RQ3 피부 병변 분할을 위한 스킵 커넥션 내 어텐션 모듈의 최적의 수와 배치는 무엇인가?
• RQ4 제안된 방법은 기존의 CNN 기반 및 Transformer 기반 모델보다 피부 병변 분할 기준 데이터셋에서 우수한 성능을 내는가?
• RQ5 공간 어텐션 및 채널 어텐션 구성 요소가 각각 최종 성능에 기여하는 정도는 어떠한가?

주요 결과

• ISIC2017에서 Att-SwinU-Net은 0.9240의 Dice Similarity Coefficient(DSC)를 기록하여 이전 SOTA 방법인 TMU-Net을 0.0076 이상 앞서며 우월성을 입증하였다.
• ISIC2018에서는 DSC 0.9105, PH2에서는 DSC 0.9504를 기록하여 모든 데이터셋에서 일관된 우수성을 보였다.
• 아블레이션 스터디 결과, 스킵 커넥션에 세 개의 어텐션 모듈을 사용할 경우 최고의 성능을 기록하였으며, DSC는 0.8987에서 0.9240으로 상승하였다.
• 공간 어텐션 모듈 또는 크로스컨텍스추얼 어텐션 모듈을 제거할 경우, DSC가 각각 0.0039와 0.0044 감소하여, 각각의 기여도를 확인하였다.
• ISIC2017에서 정확도 0.9656을 기록하여 높은 분할 일관성과 낮은 가짜 양성 결과를 보였다.
• 정성적 결과 분석에서 TransUNet 및 Swin U-Net 대비 더 매끄럽고 정밀한 분할 마스크를 제공하였으며, 특히 부족 분할 아티팩트를 줄이는 데 효과적이었다.