[논문 리뷰] U-Net Transformer: Self and Cross Attention for Medical Image Segmentation

Olivier Petit, Nicolas Thome외 2명|arXiv (Cornell University)|2021. 03. 10.

Radiomics and Machine Learning in Medical ImagingMedicine인용 수 18

한 줄 요약

이 논문은 다중 헤드 자기주의(multi-head self-attention) 및 교차주의(cross-attention) 메커니즘을 U-Net 기반 아키텍처에 통합하여, 저대비 및 복잡한 장기의 영상 분할 성능을 향상시키기 위해 U-Transformer를 제안한다. 장거리 문맥적 의존성 모델링과 주의 지도형 스킵 연결을 통한 공간적 세부 사항 정밀화를 통해 U-Transformer는 복잡한 장기, 예를 들어 췌장과 같은 기관에서 최신 기술(SOTA) 성능을 달성하며, 복부 CT 데이터셋에서 U-Net과 Attention U-Net보다 최대 3.4% 높은 Dice 스코어를 기록한다.

ABSTRACT

Medical image segmentation remains particularly challenging for complex and low-contrast anatomical structures. In this paper, we introduce the U-Transformer network, which combines a U-shaped architecture for image segmentation with self- and cross-attention from Transformers. U-Transformer overcomes the inability of U-Nets to model long-range contextual interactions and spatial dependencies, which are arguably crucial for accurate segmentation in challenging contexts. To this end, attention mechanisms are incorporated at two main levels: a self-attention module leverages global interactions between encoder features, while cross-attention in the skip connections allows a fine spatial recovery in the U-Net decoder by filtering out non-semantic features. Experiments on two abdominal CT-image datasets show the large performance gain brought out by U-Transformer compared to U-Net and local Attention U-Nets. We also highlight the importance of using both self- and cross-attention, and the nice interpretability features brought out by U-Transformer.

연구 동기 및 목표

• U-Net이 저대비 및 복잡한 해부학적 구조에서 장거리 문맥적 의존성과 공간 관계를 포착하는 데 한계가 있음을 해결하기 위해.
• 췌장, 담낭, 위와 같이 크기가 작고 시각적으로 모호한 기관의 분할 정확도를 향상시키기 위해.
• 자기주의 및 교차주의 메커니즘을 U-Net 아키텍처에 통합하여 전역적 문맥을 모델링하고 공간적 세부 사항을 정밀화하기 위해.
• 주의 맵 시각화와 분석 연구를 통해 모델의 해석 가능성성을 향상시키기 위해.
• 자기주의 및 교차주의의 상호보완적 이점을 통합된 트랜스포머 기반 U-Net 아키텍처에서 검증하기 위해.

제안 방법

• 전체 이미지의 전역적 문맥 상호작용을 모델링하기 위해 인코더의 끝단에 다중 헤드 자기주의(MHSA)를 통합한 U-Transformer 네트워크를 도입한다.
• 고수준 의미적 특징과 고해상도 특징을 융합하기 위해 스킵 연결에 다중 헤드 교차주의(MHCA)를 활용하여 세밀한 공간 복구를 가능하게 한다.
• MHSA 및 MHCA 모듈에서 공간적 위치 정보를 유지하기 위해 학습 가능한 위치 인코딩을 사용하여 주의 집중 정밀도를 향상시킨다.
• 다중 수준의 교차주의를 위해 디코더 단계의 다양한 수준에 MHCA 모듈을 스택하여 점진적인 예측 정밀화를 수행한다.
• MHSA는 특징 간 장거리 의존성을 포착하고, MHCA는 특징 업샘플링 중 비의미적 노이즈를 필터링하는 이중 브랜치 주의 메커니즘을 구현한다.
• 기본적으로 잔차 블록과 스킵 연결을 갖춘 표준 U-Net 백본을 사용하며, 합성곱 레이어를 대체하는 대신 주의 모듈을 통합한다.

실험 결과

연구 질문

• RQ1 U-Net의 제한된 수신 영역을 초월해 장거리 문맥적 의존성을 모델링함으로써, 인코더에 자기주의를 통합하면 복잡한 기관의 분할 성능 향상에 기여할 수 있는가?
• RQ2 스킵 연결에서 교차주의를 적용함으로써 비의미적 특징을 필터링하고 위치 정보를 통합함으로써 공간적 세부 사항 복구 성능이 향상되는가?
• RQ3 자기주의 및 교차주의 메커니즘이 개별적으로나 조합적으로 적용되었을 때 분할 성능 및 해석 가능성 측면에서 어떻게 비교되는가?
• RQ4 위치 인코딩과 다중 수준 주의의 영향은 디코더 경로에서 교차주의 성능에 어떤 영향을 미치는가?
• RQ5 두 주의 메커니즘의 통합이 저대비 및 복잡한 형태를 가진 기관을 포함한 다양한 기관에 대한 일반화 능력을 향상시키는가?

주요 결과

• IMO 데이터셋에서 췌장에 대해 U-Transformer는 DSC(Dice similarity coefficient) 73.10%를 기록하여 U-Net 대비 +3.4%, Attention U-Net 대비 +4.45% 향상된 성능을 보였다.
• TCIA 췌장 데이터셋에서 U-Transformer는 DSC 80.65%를 기록하여 U-Net(76.35%) 및 Attention U-Net(77.23%)를 상당한 격차로 앞서며 성능을 뛰어넘었다.
• U-Transformer에서 자기주의 및 교차주의의 조합은 각각 독립적으로 사용했을 때보다 열등한 결과를 내며, 이는 두 메커니즘이 상호보완적임을 입증한다.
• 위치 인코딩은 TCIA에서 MHCA 성능을 +1.7점 향상시키고, IMO에서는 +0.6점 향상시켜 공간 정렬에 있어 중요성을 입증한다.
• 다중 수준의 교차주의 적용은 TCIA에서 +1.8점, IMO에서 +0.6점 향상시켜 계층적 주의 정밀화가 분할 정확도 향상에 기여함을 보였다.
• Hausdorff 거리 결과도 DSC 스코어와 동일한 추세를 보이며, U-Transformer가 예측 아티팩트를 감소시키고 경계 정밀도를 향상시킴을 시사한다.