[논문 리뷰] ScaleFormer: Revisiting the Transformer-based Backbones from a Scale-wise Perspective for Medical Image Segmentation

Huimin Huang, Shiao Xie외 5명|arXiv (Cornell University)|2022. 07. 29.

Advanced Neural Network Applications인용 수 13

한 줄 요약

ScaleFormer는 의료 영상 분할을 위한 새로운 트랜스포머 기반 백본을 제안하며, 이중 브랜치 아키텍처를 통해 스케일별 과제를 해결한다: 각 스케일 내에서 국소적 및 전역적 특징을 융합하기 위한 경량화된 이중축 다중헤드 자기주의(DA-MSA)를 통한 내부 스케일 모델링과, 공간 인식 트랜스포머를 통한 상호 스케일 모델링을 통해 크로스 스케일 특징 상호작용을 향상시킨다. 이는 다양한 의료 영상 분할 벤치마크에서 최신 기술(SOTA) 성능을 달성한다.

ABSTRACT

Recently, a variety of vision transformers have been developed as their capability of modeling long-range dependency. In current transformer-based backbones for medical image segmentation, convolutional layers were replaced with pure transformers, or transformers were added to the deepest encoder to learn global context. However, there are mainly two challenges in a scale-wise perspective: (1) intra-scale problem: the existing methods lacked in extracting local-global cues in each scale, which may impact the signal propagation of small objects; (2) inter-scale problem: the existing methods failed to explore distinctive information from multiple scales, which may hinder the representation learning from objects with widely variable size, shape and location. To address these limitations, we propose a novel backbone, namely ScaleFormer, with two appealing designs: (1) A scale-wise intra-scale transformer is designed to couple the CNN-based local features with the transformer-based global cues in each scale, where the row-wise and column-wise global dependencies can be extracted by a lightweight Dual-Axis MSA. (2) A simple and effective spatial-aware inter-scale transformer is designed to interact among consensual regions in multiple scales, which can highlight the cross-scale dependency and resolve the complex scale variations. Experimental results on different benchmarks demonstrate that our Scale-Former outperforms the current state-of-the-art methods. The code is publicly available at: https://github.com/ZJUGiveLab/ScaleFormer.

연구 동기 및 목표

• 기존 비전 트랜스포머에서 각 스케일 내 국소적 및 전역적 특징이 효과적으로 융합되지 않는 내부 스케일 한계를 해결하기 위해.
• 크기와 형태가 다양한 객체에 대해 특히 부족한 크로스 스케일 특징 상호작용으로 인해 발생하는 상호 스케일 과제를 해결하기 위해.
• 특징 계층 구조 전반에 걸쳐 스케일별 의존성을 명시적으로 모델링하여 의료 영상 분할의 표현 학습을 향상시키기 위해.
• 작은 객체의 신호 전파를 향상시키고 스케일 변동에 대해 강건한 경량화된 효과적인 백본을 개발하기 위해.

제안 방법

• CNN 기반 국소적 특징와 트랜스포머 기반 전역적 특징을 경량화된 이중축 다중헤드 자기주의(DA-MSA) 메커니즘을 사용해 융합하는 스케일별 내부 스케일 트랜스포머를 도입한다.
• DA-MSA에서 행과 열 방향의 어텐션을 활용하여 양 방향에서 장거리 의존성을 독립적으로 캡처함으로써 국소-전역 특징 융합을 향상시킨다.
• 다양한 스케일 간에 대응 영역을 주시하여 크로스 스케일 상호작용을 명시적으로 모델링하는 공간 인식 상호 스케일 트랜스포머를 설계한다.
• 다양한 인코더 스테이지의 특징 맵을 상호 스케일 모듈의 입력으로 사용하여 계층적 특징 정제를 가능하게 한다.
• 학습 안정성 향상과 기울기 흐름 개선을 위해 잔차 연결과 레이어 정규화를 적용한다.
• 의료 영상 데이터셋에서 표준 분할 손실(예: dice 손실)을 사용해 전체 모델을 엔드 투 엔드로 훈련한다.

실험 결과

연구 질문

• RQ1 작은 객체의 표현 학습을 향상시키기 위해 각 스케일 내에서 국소적 및 전역적 특징을 어떻게 효과적으로 융합할 수 있는가?
• RQ2 크로스 스케일 특징 상호작용은 의료 영상에서 크기와 형태가 다양한 객체를 처리하는 데 모델의 능력을 어느 정도 향상시킬 수 있는가?
• RQ3 경량화된 스케일 인식 트랜스포머 모듈은 계산 비용을 크게 증가시키지 않으면서 성능 향상에 기여할 수 있는가?
• RQ4 제안된 내부 스케일 및 상호 스케일 설계는 의료 영상 분할에서 표준 비전 트랜스포머 백본과 비교해 어떻게 성능을 냅니다?

주요 결과

• ScaleFormer는 nnU-Net 및 기타 표준 데이터셋을 포함한 다양한 의료 영상 분할 벤치마크에서 최신 기술(SOTA) 성능을 달성한다.
• 제거 실험을 통해 내부 스케일 및 상호 스케일 모듈이 성능 향상에 크게 기여하며, 특히 작은 객체에서 두드러진다.
• 더 나은 국소-전역 특징 융합 덕분에 작은 크기 및 비정형 형태의 구조에서 분할 정확도가 향상된다.
• 공간 인식 상호 스케일 모듈은 크로스 스케일 의존성을 효과적으로 캡처하여 스케일 계층 전반에 걸쳐 더 강건한 특징 표현을 이끈다.
• 기존 트랜스포머 기반 백본과 비교해 경쟁적인 추론 속도와 파rameter 효율성을 유지한다.