[논문 리뷰] ScribFormer: Transformer Makes CNN Work Better for Scribble-based Medical Image Segmentation

Zihan Li, Zheng Yuan외 7명|arXiv (Cornell University)|2024. 02. 03.

Brain Tumor Detection and ClassificationNeuroscience인용 수 3

한 줄 요약

ScribFormer는 다중 척도 주의(multi-scale attention)를 활용한 새로운 CNN-Transformer 하이브리드 아키텍처를 제안하여, 얕은 및 깊은 Transformer 블록으로부터 국소적이고 전역적인 특징 학습을 향상시킨다. 이는 ACDC, MSCMRseg 및 HeartUII 데이터셋에서 약 0.833의 Dice 점수를 기록하며, 풀 레이블링된 방법들을 능가하는 최신 기술 성능(SOTA)을 달성한다.

ABSTRACT

Most recent scribble-supervised segmentation methods commonly adopt a CNN framework with an encoder-decoder architecture. Despite its multiple benefits, this framework generally can only capture small-range feature dependency for the convolutional layer with the local receptive field, which makes it difficult to learn global shape information from the limited information provided by scribble annotations. To address this issue, this paper proposes a new CNN-Transformer hybrid solution for scribble-supervised medical image segmentation called ScribFormer. The proposed ScribFormer model has a triple-branch structure, i.e., the hybrid of a CNN branch, a Transformer branch, and an attention-guided class activation map (ACAM) branch. Specifically, the CNN branch collaborates with the Transformer branch to fuse the local features learned from CNN with the global representations obtained from Transformer, which can effectively overcome limitations of existing scribble-supervised segmentation methods. Furthermore, the ACAM branch assists in unifying the shallow convolution features and the deep convolution features to improve model's performance further. Extensive experiments on two public datasets and one private dataset show that our ScribFormer has superior performance over the state-of-the-art scribble-supervised segmentation methods, and achieves even better results than the fully-supervised segmentation methods. The code is released at https://github.com/HUANGLIZI/ScribFormer.

연구 동기 및 목표

• 의료 영상 분할에서 희박한 스크래치 애너테이션으로부터 CNN이 전역적 형태 정보를 포착하는 데에 한계가 있음을 해결한다.
• 노이즈가 많고 완전하지 않은 스크래치 레이블로 인해 약한 감독 방법의 일반화 능력이 떨어지는 문제를 해결한다.
• CNN의 국소적 특징 추출 능력과 Transformer의 장거리 의존성 모델링 능력을 융합하여, 약한 감독 하에서 분할 성능을 향상시킨다.
• 픽셀 수준의 애너테이션이 비용이 많이 들기 때문에, 오직 스크래치 수준의 감독만으로도 높은 성능을 내는 분할을 가능하게 하여 의존도를 줄인다.
• 주어진 여러 네트워크 깊이에서 주의 지도 기반 특징 개선을 통해 통합된 프레임워크를 개발하여 특징 표현을 향상시킨다.

제안 방법

• 삼중 분기 아키텍처를 제안한다: 국소적 특징 학습을 위한 CNN 분기, 전역적 맥락 모델링을 위한 Transformer 분기, 그리고 주의 지도 기반 클래스 활성화 맵(ACAM) 분기로 특징 융합을 수행한다.
• 얕은 Transformer 블록의 주의 가중치를 사용하여 저수준의 컨볼루션 특징을 개선함으로써 공간적 세부 정보 유지 능력을 향상시킨다.
• 깊은 Transformer 블록의 주의 가중치를 사용하여 고수준 의미적 특징과 ACAM을 개선함으로써 전역적 형태 일관성을 향상시킨다.
• CNN과 Transformer 분기의 특징을 상호 주의(mechanism)를 통해 융합하여 국소적 세부 정보와 전역적 맥락을 결합한다.
• ACAM 분기를 도입하여 주의 기반 집계를 통해 얕은 및 깊은 컨볼루션 특징을 통합함으로써 특징 표현 품질을 향상시킨다.
• 완전한 마스크가 필요 없이, 희박한 스크래치 애너테이션에서 표준 분할 손실(예: Dice 손실)을 사용해 모델을 엔드 투 엔드로 훈련한다.

실험 결과

연구 질문

• RQ1 순수한 CNN과 비교했을 때, CNN-Transformer 하이브리드 아키텍처가 스크래치 감독 기반 의료 영상 분할에서 전역적 형태 학습을 향상시킬 수 있는가?
• RQ2 Transformer의 다양한 깊이 수준에서의 주의 맵은 약한 감독 환경에서 국소적 및 전역적 특징을 개선하는 데 어떻게 기여하는가?
• RQ3 CNN과 Transformer 분기 간의 주의 지도 기반 특징 융합이, 표준 특징 연결 또는 원소별 연산보다 성능이 뛰어나게 되는가?
• RQ4 제안된 방법이 오직 스크래치 애너테이션으로만 훈련되었을 때, 풀 레이블링된 방법들과 비교해 성능이 유사하거나 뛰어나게 되는가?
• RQ5 삼중 분기 아키텍처, 주의 기반 개선 등 아키텍처 설계 선택 사항이 모델의 강인성과 추론 시 분산에 미치는 영향은 무엇인가?

주요 결과

• ScribFormer는 HeartUII 데이터셋에서 0.833(Dice 점수, 95% CI: 0.808–0.854)의 성능을 기록하여, UNet++ 및 CycleMix S를 포함한 모든 비교 방법들을 뛰어넘었다.
• UNet pce 및 UNet ustr와 비교했을 때 p-value < 10^-9를 기록하여, 분할 성능 향상이 통계적으로 유의미하다는 것을 입증했다.
• UNet crf 및 CycleMix S와의 비교에서 유의미한 p-value가 없었지만, ScribFormer는 더 좁은 95% 신뢰구간과 높은 중앙값 성능를 보이며 더 높은 강인성을 나타냈다.
• ScribFormer는 동일한 데이터셋에서 풀 레이블링된 UNet++를 능가했으며, 주의 기반 특징 학습과 결합된 스크래치 기반 약한 감독 학습이 최신 기술 성능(SOTA)을 달성할 수 있음을 보여주었다.
• 제거 분석(ablation study) 결과, 얕은 및 깊은 Transformer 주의 헤드 모두 고유하게 기여함을 확인했다—얕은 헤드는 공간적 세부 정보를 향상시키고, 깊은 헤드는 의미 일관성을 향상시켰다.
• 상자 그림(Box plot) 분석 결과, ScribFormer의 추론 분산은 CycleMix S보다 낮았으며, 더 좁은 신뢰구간과 높은 중앙값 성능로 이를 뒷받침했다.