[논문 리뷰] Swin Deformable Attention Hybrid U-Net for Medical Image Segmentation

Lichao Wang, Jiahao Huang외 2명|arXiv (Cornell University)|2023. 02. 28.

Radiomics and Machine Learning in Medical ImagingMedicine인용 수 3

한 줄 요약

이 논문은 의료 영상 분할 정확도와 해석 가능성 향상을 위해 스위프트 변형 다중 헤드 자기주의(Swin Deformable Multi-Head Self-Attention, SDMSA)를 병렬 컨볼루션 브랜치와 융합한 하이브리드 U-Net 아키텍처인 SDAH-UNet을 제안한다. 심장판막이나 종양과 같은 주요 영역에 정밀하게 집중하는 동적이고 변형 가능한 주의를 가능하게 하여 계산의 불필요한 중복을 줄이고, 변형 점을 통해 시각적 설명을 제공함으로써 ACDC 및 BraTS2020 데이터셋에서 최신 기술 수준(SOTA) 성능을 달성하며, 좌심실 분할에서 92.23%의 DSC를 기록한다.

ABSTRACT

Medical image segmentation is a crucial task in the field of medical image analysis. Harmonizing the convolution and multi-head self-attention mechanism is a recent research focus in this field, with various combination methods proposed. However, the lack of interpretability of these hybrid models remains a common pitfall, limiting their practical application in clinical scenarios. To address this issue, we propose to incorporate the Shifted Window (Swin) Deformable Attention into a hybrid architecture to improve segmentation performance while ensuring explainability. Our proposed Swin Deformable Attention Hybrid UNet (SDAH-UNet) demonstrates state-of-the-art performance on both anatomical and lesion segmentation tasks. Moreover, we provide a direct and visual explanation of the model focalization and how the model forms it, enabling clinicians to better understand and trust the decision of the model. Our approach could be a promising solution to the challenge of developing accurate and interpretable medical image segmentation models.

연구 동기 및 목표

• 의료 영상 분할을 위한 하이브리드 CNN-Transformer 모델에서의 해석 가능성 부족 문제를 해결하기 위해.
• 스위프트 트랜스포머 프레임워크 내에서 표준 다중 헤드 자기주의의 계산 불필요 중복을 줄이기 위해 변형 가능한 샘플링을 도입함으로써.
• SDMSA를 통한 전역적 구조 모델링과 병렬 컨볼루션 브랜치를 통한 국소적 텍스처 학습을 융합하여 분할 정확도를 향상시키기 위해.
• 변형 점을 통해 모델 주의의 내재적이고 시각적인 설명을 제공하여 임상의가 모델 결정을 해석할 수 있도록 하기 위해.
• 설명 가능한 집중 기능을 통해 높은 성능와 임상적 신뢰성 간의 균형을 이루는 모델을 개발하기 위해.

제안 방법

• 제안된 SDAH-UNet은 세 개의 인코더 및 디코더 블록을 포함한 U-Net 아키텍처를 사용하며, 각 블록은 변형 가능한 주의와 컨볼루션 특징을 융합한 새로운 SDAPC 블록을 포함한다.
• SDAPC 블록은 불규칙한 해부학적 구조(예: 좌심실 벽)에 집중하기 위해 동적이고 비균일한 샘플링 점을 학습하는 스위프트 변형 다중 헤드 자기주의(SDMSA) 모듈을 통합한다.
• 병렬 컨볼루션 브랜치는 국소적 텍스처 세부 정보를 유지하고, 전역 주의 메커니즘을 보완하는 인덕티브 바이어스를 제공한다.
• 소형 커널과 GELU 활성화 함수를 사용하는 컨볼루션 임bedding 모듈을 통해 공간 해상도와 수신 영역을 유지한다.
• 디컨볼루션 확장 모듈을 통해 특징을 업샘플링하여 마스크 예측을 수행하며, 고해상도 출력을 유지한다.
• 변형 점, 주의 점수 히트맵, 기울기 기반 시각화(SEG-GRAD-CAM)를 통해 모델의 해석성을 향상시켰으며, 변형 점은 목표물에 정밀하게 집중하는 것으로 나타났다.

실험 결과

연구 질문

• RQ1 스위프트 트랜스포머 기반 아키텍처에서 변형 가능한 주의가 복잡하고 비정형적인 구조를 가진 의료 영상에서의 분할 정확도 향상에 기여하는가?
• RQ2 변형 가능한 주의와 컨볼루션 브랜치의 융합이 계산 불필요 중복을 줄이면서 특징 학습을 향상시키는가?
• RQ3 주의 메커니즘 내 변형 점이 기존 기울기 기반 방법보다 더 정확하고 강력한 모델 집중 설명을 제공하는가?
• RQ4 제안된 하이브리드 아키텍처가 해부학적 및 병변 분할 작업 모두에서 기존 최신 기술 수준 모델을 초월하는가?
• RQ5 모델의 주의 메커니즘이 임상적 신뢰와 이해를 지원하는 방식으로 시각적으로 해석될 수 있는가?

주요 결과

• SDAH-UNet은 ACDC 데이터셋에서 좌심실 분할에서 92.23%의 Dice Similarity Coefficient(DSC)를 달성하여 최신 기술 수준 모델을 능가했다.
• SDAPC 블록을 네 개 사용했을 경우 ACDC 데이터셋에서 좌심실 분할에서 96.91%의 DSC를 기록하였으며, HD95는 1.22 mm였다.
• 제거 실험 결과, 표준 블록을 SDAPC 블록으로 대체함으로써 평균 DSC가 약 1.2% 향상되었다.
• SDAPC 블록에서 SDMSA 또는 컨볼루션 브랜치를 제거하면 DSC가 약 2% 감소하여 양 구성 요소의 기여도가 확인되었다.
• 변형 점은 주의 점수 히트맵이나 기울기 기반 방법보다 더 명확하고 강력한 시각적 설명을 제공하였으며, 종양 경계나 강화 영역을 더 잘 포착하는 데 유리했다.
• 모델은 해부학적(ACDC) 및 병변(BraTS2020) 분할 작업 모두에서 뛰어난 성능을 보이며 일반화 능력을 입증했다.