[논문 리뷰] AgileFormer: Spatially Agile Transformer UNet for Medical Image Segmentation

Peijie Qiu, Yang Jin외 3명|arXiv (Cornell University)|2024. 03. 29.

COVID-19 diagnosis using AI인용 수 11

한 줄 요약

AgileFormer는 변형 가능한 패치 임베딩, 공간적으로 동적 자기 주의, 및 다중 스케일 변형 위치 인코딩을 ViT-UNet에 도입하여 세 가지 데이터셋에서 의료 영상 분할 분야의 최첨단 성능과 뛰어난 확장성을 달성한다.

ABSTRACT

In the past decades, deep neural networks, particularly convolutional neural networks, have achieved state-of-the-art performance in a variety of medical image segmentation tasks. Recently, the introduction of the vision transformer (ViT) has significantly altered the landscape of deep segmentation models. There has been a growing focus on ViTs, driven by their excellent performance and scalability. However, we argue that the current design of the vision transformer-based UNet (ViT-UNet) segmentation models may not effectively handle the heterogeneous appearance (e.g., varying shapes and sizes) of objects of interest in medical image segmentation tasks. To tackle this challenge, we present a structured approach to introduce spatially dynamic components to the ViT-UNet. This adaptation enables the model to effectively capture features of target objects with diverse appearances. This is achieved by three main components: extbf{(i)} deformable patch embedding; extbf{(ii)} spatially dynamic multi-head attention; extbf{(iii)} deformable positional encoding. These components were integrated into a novel architecture, termed AgileFormer. AgileFormer is a spatially agile ViT-UNet designed for medical image segmentation. Experiments in three segmentation tasks using publicly available datasets demonstrated the effectiveness of the proposed method. The code is available at \href{https://github.com/sotiraslab/AgileFormer}{https://github.com/sotiraslab/AgileFormer}.

연구 동기 및 목표

• segmentation 작업에서 모양과 크기가 서로 다른 의료 영상 타겟의 이질성 addressed.
• 공간적으로 동적 ViT-UNet을 개발하여 위치화 및 전역 맥락 처리 개선.
• 표현 학습 강화를 위한 변형 패치 임베딩 및 변형 위치 인코딩 제안.

제안 방법

• 고정 패치 임베딩을 변형 가능 컨볼루션 기반의 변형 가능 패치 임베딩으로 대체.
• deformable 다중-헤드 주의와 이웃 다중-헤드 주의를 교대로 수행하는 공간적으로 동적 자기 주의 블록을 사용.
• 자기 주의에서 불규칙한 격자를 모델링하기 위한 다중 스케일 변형 위치 인코딩 도입.
• skip 연결 및 깊은 감독을 갖춘 U-Net 스타일 인코더-디코더를 채택.
• 깊은 감독 여부에 따라 AgileFormer-T와 AgileFormer-B를 비교.
• Synapse(2D 다기관), ACDC(2D 심장), Decathlon 뇌종양(3D) 데이터셋에서 평가.

실험 결과

연구 질문

• RQ1 공간적으로 동적 구성요소가 ViT-UNet이 의료 영상의 이질적인 기관 모양과 크기를 더 잘 다루게 할 수 있는가?
• RQ2 변형 가능 패치 임베딩과 다중 스케일 변형 위치 인코딩이 고정 창(window)을 사용하는 ViT-UNet보다 분할 정확도와 확장성을 개선하는가?
• RQ3 다양한 데이터셋에서 다른 SOTA ViT-UNet과 비교해 모델 크기에 따른 AgileFormer의 성능 확장이 어떻게 나타나는가?
• RQ4 깊은 감독이 AgileFormer의 분할 성능에 미치는 영향은 무엇인가?
• RQ5 제안된 구성요소가 2D 및 3D 의료 분할 작업에서 어떻게 성능을 발휘하는가?

주요 결과

• AgileFormer는 세 가지 데이터셋(Synapse, ACDC, Decathlon)에서 최근 SOTA 방법을 능가했다.
• AgileFormer-T에서 AgileFormer-B로의 모델 확장은 중간 수준의 매개변수 및 FLOPs 증가와 함께 주목할 만한 DSC 이득을 초래한다.
• 제거 가능 패치 임베딩 및 MS-DePE를 통한 성능 향상에 대한 제거 연구가 확인되었다.
• 공간적으로 동적 주의(DMSA/NMSA) 및 변형 가능 구성요소가 다양한 객체 형태/크기에 대해 향상된 처리를 제공한다.
• 깊은 감독은 DSC 개선과 함께 성능을 더욱 향상시킨다.
• 3D 실험에서 AgileFormer-T+DS가 3D ViT-UNets와의 경쟁력 있는 결과를 보인다.