[논문 리뷰] MEW-UNet: Multi-axis representation learning in frequency domain for medical image segmentation

Jiacheng Ruan, Mingye Xie외 3명|arXiv (Cornell University)|2022. 10. 25.

Advanced Neural Network Applications인용 수 15

한 줄 요약

MEW-UNet은 3개의 공간축을 따라 2D 푸리에 변환을 활용하여 전역 주파수 도메인 표현을 학습하는 새로운 다축 외부 가중치 블록(Multi-axis External Weights block)을 제안한다. 이는 U-Net 아키텍처에서 자기주의 주의(self-attention)를 대체하기 위해 국소적 특징을 위해 깊이 분리형 합성곱을 조합한다. 모델은 Synapse 데이터셋에서 HD95 기준으로 MT-UNet보다 10.15mm 향상된 최신 기술 수준의 성능을 달성한다.

ABSTRACT

Recently, Visual Transformer (ViT) has been widely used in various fields of computer vision due to applying self-attention mechanism in the spatial domain to modeling global knowledge. Especially in medical image segmentation (MIS), many works are devoted to combining ViT and CNN, and even some works directly utilize pure ViT-based models. However, recent works improved models in the aspect of spatial domain while ignoring the importance of frequency domain information. Therefore, we propose Multi-axis External Weights UNet (MEW-UNet) for MIS 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 feature and assigns the external weight in the frequency domain, which is generated by our 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 and achieve state-of-the-art performances. In particular, on the Synapse dataset, our method outperforms MT-UNet by 10.15mm in terms of HD95. Code is available at https://github.com/JCruan519/MEW-UNet.

연구 동기 및 목표

• 기존 의료 영상 분할 모델이 공간 도메인 표현에만 집중하면서 주파수 도메인 정보를 간과하는 한계를 해결하기 위해.
• 다축 푸리에 변환을 통해 주파수 도메인 특징을 통합하여 U-Net 기반 아키텍처의 전역적 맥락 모델링을 향상시키기 위해.
• 주파수 도메인 주의와 깊이 분리형 합성곱을 통합한 블록을 통해 전역적 및 국소적 특징 학습의 균형을 맞추기 위해.
• 비전 트랜스포머에서 표준 자기주의 주의에 의존하지 않고도 표현 학습을 향상시키는 주파수 인식 주의 메커니즘을 개발하기 위해.
• 주파수 도메인 불변성과 다중 척도 맥락 모델링을 통합하여 여러 의료 영상 기준에서 최신 기술 수준의 분할 성능를 달성하기 위해.

제안 방법

• 다축 외부 가중치(Multi-axis External Weights, MEW) 블록은 입력 특징 맵을 채널 차원에 따라 네 개의 브랜치로 분할한다.
• 세 개의 브랜치는 높이-너비, 채널-너비, 채널-높이 축을 따라 각각 2D 이산 푸리에 변환(2D DFT)을 적용하여 특징을 주파수 도메인으로 매핑한다.
• 학습 가능한 외부 가중치가 주파수 도메인에 적용되어 주파수 성분을 조절함으로써 전역 맥락 모델링을 향상시킨다.
• 가중치가 적용된 주파수 특징을 다시 공간 도메인으로 변환하기 위해 역 2D DFT를 적용한다.
• 네 번째 브랜치는 국소적 공간 세부 정보를 유지하기 위해 깊이 분리형 합성곱을 적용한다.
• 모든 네 브랜치의 출력을 연결하고 원본 입력에 잔차 연결(residual connection)을 적용하여 다축 외부 가중치 블록(Multi-axis External Weights Block, MEWB)을 구성한다.

실험 결과

연구 질문

• RQ1 다축 주파수 도메인 표현 학습이 공간 도메인 모델링을 초월하여 의료 영상 분할 성능 향상에 기여할 수 있는가?
• RQ2 다양한 축을 따라 2D DFT를 사용한 주파수 도메인 학습이 모델의 전역 맥락 파악 능력에 어떤 영향을 미치는가?
• RQ3 통합된 블록 내에서 주파수 도메인 주의와 국소적 합성곱을 조합할 경우 분할 성능에 기여하는 바는 무엇인가?
• RQ4 제안된 MEW-UNet이 기존의 하이브리드 CNN-ViT 및 순수 ViT 모델보다 의료 영상 분할 과제에서 우월한 성능을 보일 수 있는가?
• RQ5 주파수 도메인에 학습 가능한 외부 가중치를 포함할 경우 모델의 일반화 능력과 성능에 어떤 영향을 미치는가?

주요 결과

• MEW-UNet은 ISIC17, ISIC18, Synapse, ACDC의 네 가지 공개 의료 영상 분할 데이터셋에서 최신 기술 수준의 성능를 달성한다.
• Synapse 데이터셋에서 MEW-UNet은 MT-UNet 대비 HD95를 10.15mm 향상시키고 DSC를 0.33% 향상시켜 경계 정렬 능력이 뛰어나다는 것을 입증한다.
• ACDC 데이터셋에서 MEW-UNet은 DSC 점수 91.00%를 기록하여 MT-UNet(90.43%)과 Swin-Unet(88.07%)를 모두 초월한다.
• 단계적 실험 결과, 세 축 모두에 주파수 도메인 연산을 추가할수록 mIoU와 DSC가 점진적으로 향상되며, 전체 MEWB는 ISIC18에서 mIoU 81.90%, DSC 90.05%를 달성한다.
• 단계적 실험에서 GroupNorm을 BatchNorm으로 대체할 경우 성능 저하가 발생하여, 이 아키텍처에 있어 GroupNorm이 더 효과적임을 확인한다.
• Synapse 데이터셋에 대한 시각화 결과, MEW-UNet은 특히 장기 경계 주변에서 더 선명하고 정확한 분할 마스크를 생성하는 것으로 나타났다.