[논문 리뷰] Dual Cross-Attention for Medical Image Segmentation

Gorkem Can Ates, Prasoon P. Mohan외 1명|arXiv (Cornell University)|2023. 03. 30.

Radiomics and Machine Learning in Medical Imaging참고 문헌 74인용 수 4

한 줄 요약

이 논문은 U-Net 기반 의료 영상 분할 모델에서 다중 스케일 인코더 특징 간의 채널별 및 공간적 종속성을 교차 attention을 통해 순차적으로 모델링함으로써 스킵 커넥션 성능을 향상시키는 경량 attention 모듈인 이중 교차 attention(Dual Cross-Attention, DCA)을 제안한다. DCA는 모델 복잡도를 크게 증가시키지 않으면서도 인코더와 디코더 특징 간의 의미적 격차를 좁힘으로써 다섯 개인 기준 데이터셋에서 일관된 딱지 점수 향상을 이끌어내며, 최대 2.05% 향상됨.

ABSTRACT

We propose Dual Cross-Attention (DCA), a simple yet effective attention module that is able to enhance skip-connections in U-Net-based architectures for medical image segmentation. DCA addresses the semantic gap between encoder and decoder features by sequentially capturing channel and spatial dependencies across multi-scale encoder features. First, the Channel Cross-Attention (CCA) extracts global channel-wise dependencies by utilizing cross-attention across channel tokens of multi-scale encoder features. Then, the Spatial Cross-Attention (SCA) module performs cross-attention to capture spatial dependencies across spatial tokens. Finally, these fine-grained encoder features are up-sampled and connected to their corresponding decoder parts to form the skip-connection scheme. Our proposed DCA module can be integrated into any encoder-decoder architecture with skip-connections such as U-Net and its variants. We test our DCA module by integrating it into six U-Net-based architectures such as U-Net, V-Net, R2Unet, ResUnet++, DoubleUnet and MultiResUnet. Our DCA module shows Dice Score improvements up to 2.05% on GlaS, 2.74% on MoNuSeg, 1.37% on CVC-ClinicDB, 1.12% on Kvasir-Seg and 1.44% on Synapse datasets. Our codes are available at: https://github.com/gorkemcanates/Dual-Cross-Attention

연구 동기 및 목표

• U-Net 기반 아키텍처에서 의료 영상 분할을 위한 인코더와 디코더 특징 간의 의미적 격차를 해소하기 위해.
• 다중 스케일 인코더 특징 간의 장거리 종속성을 모델링함으로써 스킵 커넥션 성능을 향상시키기 위해.
• 계산 오버헤드를 증가시키지 않으면서도 특징 융합 성능을 향상시키는 경량형 플러그 앤 플레이 attention 모듈을 설계하기 위해.
• DCA의 유효성을 다양한 U-Net 변종과 기준 데이터셋에서 검증하기 위해.

제안 방법

• CCA는 다중 스케일 인코더 특징의 채널 토큰 간에 교차 attention을 적용하여 전반적인 채널별 종속성을 포착한다.
• 이후 SCA는 특징 맵 내의 공간 토큰 간에 교차 attention을 적용하여 공간적 종속성을 모델링한다.
• CCA 및 SCA 모듈은 CCA-SCA 순서로 순차적으로 적용되며, 특징 맵은 2D 평균 풀링을 통해 패치 임베딩을 처리하고, 1×1 깊이 분리 컨볼루션을 통해 투영을 수행한다.
• 향상된 인코더 특징은 업샘플링되어 해당 디코더 레이어에 스킵 커넥션을 통해 연결되며, 기존의 표준 연결 방식을 대체한다.
• DCA 블록은 스킵 커넥션을 갖는 모든 U-Net 기반 아키텍처에 쉽게 통합될 수 있도록 설계되었다.
• 제거 실험에서는 융합 전략(합산, 연결, 순차적)과 패치 임베딩 방법(평균 풀링 대비 컨볼루션)을 비교하였다.

실험 결과

연구 질문

• RQ1 채널과 공간 종속성을 순차적으로 모델링하는 이중 attention 메커니즘이 U-Net 기반 분할에서 특징 표현을 향상시키는가?
• RQ2 채널 및 공간 교차 attention의 적용 순서(CCA-SCA 대비 SCA-CCA) 중 어느 것이 최적의 성능을 낳는가?
• RQ3 채널 및 공간 attention 출력을 융합하는 데 있어 어떤 융합 전략(합산, 연결, 순차적)이 가장 우수한가?
• RQ4 성능 및 파rameter 효율성 측면에서 평균 풀링과 컨볼루션 기반 패치 임베딩 방식은 어떻게 비교되는가?

주요 결과

• CCA-SCA 순차적 배치가 가장 우수한 성능을 기록하였으며, 개별 모듈 대비 GlaS에서 0.79% 향상되고, MoNuSeg에서 0.16% 향상됨.
• 순차적 융합(CCA-then-SCA) 전략이 합산 및 연결 전략을 모두 초월하여, GlaS 데이터셋에서 딱지 점수 2.05% 향상됨.
• 패치 임베딩에 2D 평균 풀링을 사용할 경우 컨volutional 패치 임베딩보다 성능이 뛰어나며, GlaS에서 0.14% 높은 딱지 점수를 기록하고 파rameter 수가 감소함.
• DCA는 여섯 개인 U-Net 변종에서 모두 분할 성능 향상을 이끌어내었으며, MoNuSeg에서 2.74%, Synapse에서 1.44%, CVC-ClinicDB에서 1.37% 향상됨.
• DCA 모듈은 낮은 계산 비용을 유지하며, 파arameter 수가 약간 증가(예: U-Net 대비 8.75M 대비 8.64M)하지만 일관된 성능 향상을 제공함.
• 이 방법은 일반화 능력이 뛰어나 GlaS, MoNuSeg, CVC-ClinicDB, Kvasir-Seg, Synapse 등 다양한 데이터셋에서 뚜렷한 성능 향상을 보임.