[논문 리뷰] DAE-Former: Dual Attention-guided Efficient Transformer for Medical Image Segmentation

Reza Azad, René Arimond외 3명|arXiv (Cornell University)|2022. 12. 27.

Advanced Neural Network Applications인용 수 5

한 줄 요약

DAE-Former는 공간 및 채널 차원 전반의 장거리 의존성을 모델링하기 위해 효율적 공간 주의(efficient spatial attention)와 역행렬(채널) 주의(transpose attention)을 도입한 계산적으로 효율적이고 컨볼루션 없는 U-Net 유사 Transformer 아키텍처를 제안한다. 이는 사전 훈련 없이도 심장 및 피부 병변 분할 벤치마크에서 최신 기술 수준(SOTA) 성능을 달성하며, 향상된 특징 융합과 국소화를 위해 고유한 스킵 커넥션 크로스 주의(ScC) 모듈을 활용한다.

ABSTRACT

Transformers have recently gained attention in the computer vision domain due to their ability to model long-range dependencies. However, the self-attention mechanism, which is the core part of the Transformer model, usually suffers from quadratic computational complexity with respect to the number of tokens. Many architectures attempt to reduce model complexity by limiting the self-attention mechanism to local regions or by redesigning the tokenization process. In this paper, we propose DAE-Former, a novel method that seeks to provide an alternative perspective by efficiently designing the self-attention mechanism. More specifically, we reformulate the self-attention mechanism to capture both spatial and channel relations across the whole feature dimension while staying computationally efficient. Furthermore, we redesign the skip connection path by including the cross-attention module to ensure the feature reusability and enhance the localization power. Our method outperforms state-of-the-art methods on multi-organ cardiac and skin lesion segmentation datasets without requiring pre-training weights. The code is publicly available at https://github.com/mindflow-institue/DAEFormer.

연구 동기 및 목표

• 고해상도 의료 영상에서 표준 자기주의(self-attention)의 제곱형 계산 복잡도를 해결한다.
• 의료 영상 분할에서 장거리 공간적 및 구조적 의존성을 포착하는 데 한계가 있는 컨볼루션 네트워크(CNNs)의 한계를 극복한다.
• 컨볼루션 기반 모델에 의존하지 않고도 높은 효율성과 강력한 국소화 능력을 유지하는 순수 트랜스포머 기반 U-Net 아키텍처를 설계한다.
• 에코더와 디코더 간의 특징 융합을 향상시키기 위해 스킵 커넥션 크로스 주의(ScC) 모듈을 도입하여 표현 재사용성과 분할 정확도를 향상시킨다.
• 사전 훈련된 가중치가 필요 없이 다중 기관 및 피부 병변 분할 작업에서 최신 기술 수준의 성능를 달성한다.

제안 방법

• 컨볼루션 연산을 제거하기 위해 겹치는 4×4 패치를 사용해 입력 영상을 시퀀스 토큰으로 토크나이즈한다.
• 세 개의 스택된 블록을 갖는 계층적 에코더-디코더 구조를 구현하며, 패치 병합 및 확장 레이어를 통해 공간 차원을 압축하고 확장하면서 채널 차원을 두 배로 증가시킨다.
• 이중 주의 블록을 도입하여 두 단계로 구성한다: (1) 정규화된 쿼리 및 키프로 공간 중요도를 계산하는 효율적 주의(efficient attention)로 O(N²) 복잡도를 O(d²N)으로 감소시키고, (2) 키프 및 밸류 행렬을 전치하여 채널 간 의존성을 모델링하는 전치 주의(transpose attention).
• 소프트맥스 정규화된 쿼리 및 키프를 사용해 효율적 주의 메커니즘을 적용하고, 값과의 행렬 곱셈을 통해 전역적 맥락 벡터를 생성함으로써 계산 비용을 감소시키면서도 높은 표현 능력을 확보한다.
• 에코더와 디코더 경로의 특징을 공간 및 채널 표현을 모두 고려해 주의를 통해 융합하는 스킵 커넥션 크로스 주의(ScC) 모듈을 설계하여 특징 재사용성과 국소화 정확도를 향상시킨다.
• 최종 단계에서 선형 투영 레이어를 사용해 정제된 특징에서 분할 맵을 생성한다.

실험 결과

연구 질문

• RQ1 순수 트랜스포머 기반 U-Net 아키텍처가 계산 효율성을 유지하면서도 의료 영상 분할에서 최신 기술 수준 성능를 달성할 수 있는가?
• RQ2 자기주의 메커니즘을 어떻게 재구조화하여 제곱형 복잡도를 유발하지 않고도 공간적 및 채널별 의존성을 효율적으로 모델링할 수 있는가?
• RQ3 스킵 커넥션 크로스 주의(ScC) 모듈을 도입함으로써 트랜스포머 기반 분할 네트워크에서 특징 융합과 국소화가 얼마나 향상되는가?
• RQ4 순수 이중 주의 메커니즘(효율적 공간 주의 + 전치 채널 주의)에 기반한 모델이 사전 훈련 없이도 CNN-Transformer 하이브리드 및 기존 순수 트랜스포머보다 의료 영상 분할에서 승리할 수 있는가?
• RQ5 제안된 아키텍처가 심장 및 피부 병변 분할과 같은 다양한 의료 영상 작업에 어떻게 일반화되는가?

주요 결과

• DAE-Former는 사전 훈련된 가중치가 필요 없이 다기관 심장 분할 및 피부 병변 분할 데이터셋에서 최신 기술 수준 성능를 달성한다.
• 모델은 강력한 일반화 능력을 보이며, 심장 및 피부 병변 분할 벤치마크에서 기존 최신 기술 수준 방법들을 능가한다.
• 이중 주의 메커니즘—효율적 공간 주의 이후 전치 채널 주의—는 계산 비용을 줄이며 공간 및 채널 차원 전반의 장거리 의존성을 효과적으로 모델링할 수 있다.
• 스킵 커넥션 크로스 주의(ScC) 모듈은 에코더 및 디코더 특징 간의 다중 모odal 주의를 가능하게 하여 특징 재사용성과 국소화 정확도를 크게 향상시킨다.
• 컨볼루션 기반 모델 없이 순수 트랜스포머 설계와 효율적 주의를 통해 고해상도 의료 영상에 적합한 높은 효율성을 유지한다.
• 제거 분석(ablation study)은 이중 주의 메커니즘과 ScC 모듈이 성능 향상에 기여하며 설계 선택의 타당성을 입증한다.