[논문 리뷰] Multi-dimension Transformer with Attention-based Filtering for Medical Image Segmentation

Wentao Wang, Xi Xiao외 6명|arXiv (Cornell University)|2024. 05. 20.

Brain Tumor Detection and ClassificationNeuroscience인용 수 3

한 줄 요약

이 논문은 병변 영역의 특징 표현을 향상시키기 위해 패치 임bedding 단계에서 주의 기반 노이즈 필터링을 통합하고, 공간적 및 채널 차원을 모두 고려한 자기주도 주의 메커니즘을 확장한 다차원 트랜스포머인 MDT-AF를 제안한다. 이는 저 SNR(신호 대 잡음 비율)에 대한 강건성과 풍부한 특징 학습 능력을 향상시켜 세 가지 공개 벤치마크에서 최신 기술 수준(SOTA) 성능을 달성한다.

ABSTRACT

The accurate segmentation of medical images is crucial for diagnosing and treating diseases. Recent studies demonstrate that vision transformer-based methods have significantly improved performance in medical image segmentation, primarily due to their superior ability to establish global relationships among features and adaptability to various inputs. However, these methods struggle with the low signal-to-noise ratio inherent to medical images. Additionally, the effective utilization of channel and spatial information, which are essential for medical image segmentation, is limited by the representation capacity of self-attention. To address these challenges, we propose a multi-dimension transformer with attention-based filtering (MDT-AF), which redesigns the patch embedding and self-attention mechanism for medical image segmentation. MDT-AF incorporates an attention-based feature filtering mechanism into the patch embedding blocks and employs a coarse-to-fine process to mitigate the impact of low signal-to-noise ratio. To better capture complex structures in medical images, MDT-AF extends the self-attention mechanism to incorporate spatial and channel dimensions, enriching feature representation. Moreover, we introduce an interaction mechanism to improve the feature aggregation between spatial and channel dimensions. Experimental results on three public medical image segmentation benchmarks show that MDT-AF achieves state-of-the-art (SOTA) performance.

연구 동기 및 목표

• 의료 영상에서 낮은 신호 대 잡음 비율이 특징 학습과 분할 정확도에 악영향을 미치는 문제를 해결한다.
• 의료 영상 분할에 핵심적인 복잡한 공간적 및 채널 별 상관관계를 포착하는 데에 시각 트랜스포머의 표현 능력을 향상시킨다.
• 표준 자기주도 주의 메커니즘이 의료 영상에서 장거리 및 다차원 특징을 모델링하는 데에 한계를 보이는 문제를 해결한다.
• 고도화된 학습 전략에 의존하지 않고도 특징 품질과 분할 정밀도를 향상시키는 강건한 종단 간 프레임워크를 개발한다.

제안 방법

• 자신주도 주의 기반 필터링 메커니즘을 패치 임bedding 단계에 도입하여 거친 특징을 정제하고, 거친 것에서부터 세련된 방식으로 노이즈를 억제한다.
• 패치 임bedding 모듈을 재설계하여 자기주도 주의 계산 이전에 관련 없거나 노이즈가 많은 특징을 필터링하는 데 도움이 되는 주의 가중치를 생성한다.
• 자기주도 주의 메커니즘을 공간적, 채널적, 그리고 공간-채널 조합 차원으로 확장하여 더 풍부한 특징 표현을 가능하게 한다.
• 공간적 특징과 채널적 특징 간의 특징 상호작용 및 집합을 수행하는 다차원 트랜스포머 블록을 구현하여 분류 능력을 향상시킨다.
• 에코더에서 추출한 계층적 특징을 융합하기 위해 다중 수준의 MLP 디코더를 사용하여 정확한 의미 분할 마스크 예측을 수행한다.
• 주의 기반 필터링과 다차원 자기주도 주의 메커니즘을 통합된 인코더-디코더 아키텍처에 통합하여 종단 간 학습을 구현한다.

실험 결과

연구 질문

• RQ1 패치 임bedding 단계에서 주의 기반 필터링이 저 SNR 의료 영상에서 특징 품질 향상과 노이즈에 대한 강건성을 향상시키는가?
• RQ2 공간적 및 채널적 차원으로 확장된 자기주도 주의 메커니즘이 의료 영상 분할에서 특징 표현과 성능 향상에 기여하는가?
• RQ3 제안된 MDT-AF 프레임워크는 다양한 의료 영상 벤치마크에서 기존의 시각 트랜스포머 및 CNN 기반 방법보다 얼마나 효과적인가?
• RQ4 다차원 자기주도 주의 메커니즘이 의료 영상에서 복잡한 해부학적 구조를 얼마나 잘 모델링하는가?

주요 결과

• Lung X-ray 데이터셋에서 MDT-AF는 DSC(재현율 유사도 계수) 97.17%를 기록하여 기준 모델 대비 +0.33% 향상되었다.
• 피부 병변 데이터셋에서 MDT-AF는 DSC 94.25%를 기록하여 기준 모델 대비 +0.09% 향상되었다.
• Kvasir-SEG 데이터셋에서 MDT-AF는 DSC 93.38%를 기록하여 기준 모델 대비 +1.11% 향상되었다.
• 제거 실험 결과, 주의 기반 필터링과 다차원 자기주도 주의 구성 요소가 성능 향상에 크게 기여하는 것으로 확인되었다.
• 모든 MDT-AF 구성 요소를 포함한 모델은 필터링 기능이 없는 변형(+1.11% DSC 향상, Kvasir-SEG 기준)이나 다차원 주의 기능이 없는 변형(+0.78% DSC 향상, Kvasir-SEG 기준)보다 성능이 뛰어나다.
• 시각적 비교 결과, MDT-AF는 특히 노이즈가 많은 내 endoscopic 영상에서 더 선명한 경계와 더 적은 오분류를 생성하는 것으로 나타났다.