[논문 리뷰] MISSFormer: An Effective Medical Image Segmentation Transformer

Xiaohong Huang, Zhifang Deng외 2명|arXiv (Cornell University)|2021. 09. 15.

Advanced Neural Network Applications참고 문헌 40인용 수 170

한 줄 요약

MISSFormer은 Enhanced Transformer Block과 Enhanced Transformer Context Bridge를 갖춘 계층적 U자형 트랜스포머를 처음부터 학습시키고, Synapse 및 ACDC 데이터셋에서 의학 영상 분할의 최첨단 성능을 달성한다.

ABSTRACT

The CNN-based methods have achieved impressive results in medical image segmentation, but they failed to capture the long-range dependencies due to the inherent locality of the convolution operation. Transformer-based methods are recently popular in vision tasks because of their capacity for long-range dependencies and promising performance. However, it lacks in modeling local context. In this paper, taking medical image segmentation as an example, we present MISSFormer, an effective and powerful Medical Image Segmentation tranSFormer. MISSFormer is a hierarchical encoder-decoder network with two appealing designs: 1) A feed-forward network is redesigned with the proposed Enhanced Transformer Block, which enhances the long-range dependencies and supplements the local context, making the feature more discriminative. 2) We proposed Enhanced Transformer Context Bridge, different from previous methods of modeling only global information, the proposed context bridge with the enhanced transformer block extracts the long-range dependencies and local context of multi-scale features generated by our hierarchical transformer encoder. Driven by these two designs, the MISSFormer shows a solid capacity to capture more discriminative dependencies and context in medical image segmentation. The experiments on multi-organ and cardiac segmentation tasks demonstrate the superiority, effectiveness and robustness of our MISSFormer, the experimental results of MISSFormer trained from scratch even outperform state-of-the-art methods pre-trained on ImageNet. The core designs can be generalized to other visual segmentation tasks. The code has been released on Github: https://github.com/ZhifangDeng/MISSFormer

연구 동기 및 목표

• 의료 영상 분할에서 장거리 의존성 모델링의 한계를 CNN이 가진 한계로 인식한다.
• 정확한 분할을 위한 위치 자유화된 계층형 U자형 트랜스포머를 제안한다.
• 로컬 및 글로벌 컨텍스트를 포착하기 위해 Enhanced Transformer Block과 Enhanced Transformer Context Bridge를 설계한다.
• 의료 데이터셋에서 from scratch로 MISSFormer를 평가하고, 견고성과 처음부터 학습 가능한 특성을 입증한다.

제안 방법

• Feature discrimination 및 로컬/글로벌 컨텍스트 통합을 개선하기 위해 Enhanced Mix-FFN이라는 피드포워드 네트워크를 재설계했다.
• LayerNorm, Efficient Self-Attention, Enhanced Mix-FFN을 결합한 Enhanced Transformer Block으로 장거리 정보와 로컬 정보를 더 낮은 복잡도로 모델링한다.
• 다중 스케일 피처를 플랫화된 다중 수준 토큰을 연결하고 Enhanced Transformer Block으로 처리하는 Enhanced Transformer Context Bridge를 통해 피처를 융합한다.
• 중첩된 4x4 패치, 패치 병합/확장, 스킵 연결이 있는 U자형 아키텍처의 계층적 인코더-디코더를 제시한다.
• 의료 데이터셋에서 처음부터 학습하는 학습 방식이며, 일반적인 데이터 증가를 활용하고 SGD 옵티마이저와 다항 학습률 정책을 사용한다.

실험 결과

연구 질문

• RQ1 MISSFormer가 Synapse 및 ACDC 데이터셋에서 from scratch로 학습된 의학 영상 분할 방법들보다 최첨단 성능을 달성할 수 있는가?
• RQ2 Enhanced Transformer Block 및 Context Bridge가 기존 Transformer/MLP 기반 접근법에 비해 판별력과 컨텍스트 모델링을 향상시키는가?
• RQ3 Enhanced Transformer Context Bridge를 통한 다중 스케일 피처 융합이 분할 정확도와 경계 구분에 어떤 영향을 주는가?
• RQ4 Enhanced Mix-FFN의 다양한 스킵 연결 및 재귀적 단계가 수렴성과 성능에 어떤 영향을 미치는가?

주요 결과

• MISSFormer는 Synapse 및 ACDC 데이터셋에서 최첨단 성능을 달성하며 종종 ImageNet 사전 학습 없이도 우수한 성능을 보인다.
• 간단한 Enhanced Mix-FFN과 재귀적 스킵 연결이 기본 SegFormer 기반 변형들보다 학습 안정성과 분할 정확도를 개선한다.
• 변성 실험에서 Enhanced Transformer Context Bridge와 다중 스케일 융합을 갖춘 MISSFormer가 Dice-Sørensen 계수(DSC)에서 유의미한 이득과 경계 구분 개선을 보인다.
• 다중 스케일 브리지가 없는 MISSFormer_S는 다중 스케일 정보 통합의 이점을 보여주는 MISSFormer와 비교해 저조한 성능을 보인다.
• MISSFormer는 어려운 사례에서도 강한 경계 예측과 강인한 성능을 나타내며, Synapse에서 항목별 결과가 경쟁력 있거나 우수하고 ACDC에서 전반적으로 견고한 성능을 보인다.