[논문 리뷰] Learning Multi-axis Representation in Frequency Domain for Medical Image Segmentation

Jiacheng Ruan, Jingsheng Gao외 2명|arXiv (Cornell University)|2023. 12. 28.

Advanced Neural Network Applications인용 수 4

한 줄 요약

이 논문은 새로운 다축 외부 가중치 블록(Multi-axis External Weights block)을 통해 다차원 주파수 도메인 표현을 학습함으로써 의료 영상 분할 성능을 향상시키는 U-Net 기반 아키텍처인 MEW-UNet을 제안한다. 3개의 공간 및 채널 축을 따라 2D 이산 푸리에 변환(2D DFT)을 적용하고 주파수 도메인에서 학습 가능한 가중치로 특징을 조절함으로써 더 풍부한 전역 및 국소적 맥락을 포착한다. 이 모델은 Synapse, ACDC, ISIC17, ISIC18 등 네 가지 의료 영상 데이터셋에서 최신 기술(SOTA) 성능을 달성하였으며, ISIC18에서 평균 Dice 스코어 91.00%를 기록하였다.

ABSTRACT

Recently, Visual Transformer (ViT) has been extensively used in medical image segmentation (MIS) due to applying self-attention mechanism in the spatial domain to modeling global knowledge. However, many studies have focused on improving models in the spatial domain while neglecting the importance of frequency domain information. Therefore, we propose Multi-axis External Weights UNet (MEW-UNet) based on the U-shape architecture by replacing self-attention in ViT with our Multi-axis External Weights block. Specifically, our block performs a Fourier transform on the three axes of the input features and assigns the external weight in the frequency domain, which is generated by our External Weights Generator. Then, an inverse Fourier transform is performed to change the features back to the spatial domain. We evaluate our model on four datasets, including Synapse, ACDC, ISIC17 and ISIC18 datasets, and our approach demonstrates competitive performance, owing to its effective utilization of frequency domain information.

연구 동기 및 목표

• 기존 비전 트랜스포머 기반 모델이 공간 도메인 모델링에만 집중하고 주파수 도메인 정보를 간과하는 한계를 해결하기 위해.
• 다양한 축을 통해 주파수 도메인 신호를 활용하여 의료 영상 분할에서 전역 맥락 모델링을 향상시키기 위해.
• 주파수 도메인 표현과 국소 공간 특징을 효과적으로 융합할 수 있는 학습 가능한 메커니즘을 개발하기 위해.
• 다축 주파수 도메인 분석이 단일 축 분석보다 더 구분력 있는 신호 패턴을 제공하며, 특히 미세한 병리학적 구조에 대해 유의미한 성능 향상을 이끌어내는지 검증하기 위해.

제안 방법

• 다축 외부 가중치 블록(MEWB)은 높이-너비, 채널-너비, 채널-높이의 세 축을 각각 독립적으로 2D 이산 푸리에 변환(2D DFT)하여 다축 주파수 표현을 추출한다.
• 외부 가중치 생성기(External Weights Generator)를 통해 학습 가능한 외부 가중치를 생성하고, 이를 주파수 도메인 특징에 적용하여 전역 맥락 정보를 조절한다.
• 네 번째 브랜치는 깊이분리형(DW) 합성곱을 적용하여 국소 공간 세부 정보를 유지함으로써 국소 및 전역 표현 간의 보완적 학습을 보장한다.
• MEWB는 U-Net 인코더-디코더 아키텍처 내 비전 트랜스포머의 자기주의 모듈을 대체하여 MEW-UNet 모델을 구성한다.
• 가중치 조절 후 주파수 도메인 특징 맵을 역으로 이산 푸리에 변환(inverse DFT)을 통해 공간 도메인으로 변환함으로써 엔드 투 엔드 학습이 가능하도록 한다.
• 표준 평가 지표인 Dice 스코어 및 mIoU를 사용하여 네 가지 공개 의료 영상 분할 데이터셋에서 모델을 엔드 투 엔드로 학습시켰다.

실험 결과

연구 질문

• RQ1 다축 주파수 도메인 표현 학습이 단일 축 또는 공간 도메인 전용 접근 방식에 비해 의료 영상 분할 성능 향상에 기여하는가?
• RQ2 주파수 도메인에 학습 가능한 외부 가중치를 통합할 경우, 장거리 의존성과 맥락 정보를 포착하는 데 모델의 능력이 어떻게 향상되는가?
• RQ3 분할 작업에서 주파수 도메인 특징과 국소 공간 특징 간의 상대적 기여도는 어떠하며, 상호작용 방식은 무엇인가?
• RQ4 제안된 MEW-UNet이 다양한 데이터셋에서 기존 최고 성능 모델을 초월하는가?

주요 결과

• MEW-UNet은 ISIC18 데이터셋에서 평균 Dice 스코어 91.00%를 기록하여 기존 최고 성능 모델인 MT-UNet(90.43%) 및 Swin-Unet(88.07%)를 능가하였다.
• 제거 실험 결과, DW 합성곱 브랜치를 제거할 경우 mIoU가 0.61% 감소하여 국소 특징 학습의 중요성을 확인하였다.
• 외부 가중치 생성기 내부 표현 블록(IRB) 없이 무작위 가중치만 사용할 경우 mIoU가 1.31% 감소하여 학습 가능한 가중치 생성 메커니즘의 필수성을 입증하였다.
• 세 주파수 축과 DW 브랜치를 모두 포함한 전체 MEWB는 ISIC18에서 81.90% mIoU를 기록하여 단일 축 또는 DW 없음 기반 모델보다 뚜렷한 성능 향상을 보였다.
• 시각화 결과, MEW-UNet은 MT-UNet 및 TransFuse에 비해 더 정밀하고 완전한 경계를 생성하였으며, 특히 작은 또는 비정상적인 병변을 가진 도전적인 케이스에서 유의미한 개선을 보였다.
• 주파수 도메인 신호 강도 분석 결과, 다축 DFT는 다양한 조직 유형 간에 겹치지 않는 신호 패턴을 드러내어 설계 선택의 구분력 향상에 기여함을 검증하였다.