[논문 리뷰] CoTr: Efficiently Bridging CNN and Transformer for 3D Medical Image Segmentation

Yutong Xie, Jianpeng Zhang외 2명|arXiv (Cornell University)|2021. 03. 04.

Advanced Neural Network Applications참고 문헌 28인용 수 56

한 줄 요약

CoTr은 CNN 인코더와 Deformable Transformer를 하이브리드로 결합하여 3D 의료 영상 분할에서 장거리 맥락을 효율적으로 모델링하고 BCV 11-organ 분할에서 최첨단 성능을 달성하는 한편 고해상도 다중 스케일 특징을 처리합니다.

ABSTRACT

Convolutional neural networks (CNNs) have been the de facto standard for nowadays 3D medical image segmentation. The convolutional operations used in these networks, however, inevitably have limitations in modeling the long-range dependency due to their inductive bias of locality and weight sharing. Although Transformer was born to address this issue, it suffers from extreme computational and spatial complexities in processing high-resolution 3D feature maps. In this paper, we propose a novel framework that efficiently bridges a {\bf Co}nvolutional neural network and a {\bf Tr}ansformer {\bf (CoTr)} for accurate 3D medical image segmentation. Under this framework, the CNN is constructed to extract feature representations and an efficient deformable Transformer (DeTrans) is built to model the long-range dependency on the extracted feature maps. Different from the vanilla Transformer which treats all image positions equally, our DeTrans pays attention only to a small set of key positions by introducing the deformable self-attention mechanism. Thus, the computational and spatial complexities of DeTrans have been greatly reduced, making it possible to process the multi-scale and high-resolution feature maps, which are usually of paramount importance for image segmentation. We conduct an extensive evaluation on the Multi-Atlas Labeling Beyond the Cranial Vault (BCV) dataset that covers 11 major human organs. The results indicate that our CoTr leads to a substantial performance improvement over other CNN-based, transformer-based, and hybrid methods on the 3D multi-organ segmentation task. Code is available at \def\UrlFont{ m\small tfamily} \url{https://github.com/YtongXie/CoTr}

연구 동기 및 목표

• CNN의 지역적 귀납 편향을 Transformer로 연결하여 3D 의료 영상 분할의 가능성을 모티브로 삼는다.
• 다중 스케일 특징 맵에서 장거리 의존성을 효율적으로 모델링하기 위한 Deformable Transformer(DeTrans)를 개발한다.
• 고해상도 디테일을 보존하면서 글로벌 컨텍스트를 포착하는 CNN-encoder–DeTrans-encoder–decoder 아키텍처를 설계한다.
• BCV 데이터셋에서 CNN 기반, Transformer 기반, 및 기타 하이브리드 메서드에 비해 분할 성능 향상을 입증한다.

제안 방법

• 다중 스케일 3D 특징 맵을 추출하기 위한 CNN-encoder를 사용한다.
• 다중 스케일 변형 가능한 자기 주의(multi-scale deformable self-attention)를 갖춘 DeTrans-encoder를 도입하여 장거리 의존성을 효율적으로 캡처한다.
• CNN 특징을 3D 위치 인코딩과 함께 평탄화하고 DeTrans 층을 통해 처리한다.
• 복잡도 감소를 위해 제한된 샘플 포인트를 갖는 다중 헤드 변형 가능한 자기 주의를 적용한다.
• DeTrans 출력을 CNN 기반 디코더와 skip 연결 및 심층 감독으로 융합한다.
• 공동 Dice 및 cross-entropy 손실로 최적화하고, 데이터 증강 및 인스턴스 정규화를 사용한다.

실험 결과

연구 질문

• RQ1 제한된 계산 자원으로 구성된 하이브리드 CNN–Transformer 인코더가 단독 CNN 또는 단독 Transformer 접근법보다 3D 의료 영상 분할에서 우수성을 보일 수 있는가?
• RQ2 다중 스케일 변형 가능한 자기 주의가 고해상도 3D 특징 맵에서 효과적인 장거리 모델링을 가능하게 하는가?
• RQ3 DeTrans 하이퍼파라미터 및 다중 스케일 특징 통합이 분할 성능에 어떤 영향을 미치는가?
• RQ4 CoTr이 BCV 다기관 분할 태스크에서 기존의 CNN-, Transformer-, 및 하이브리드 기반 방법과 어떻게 비교되는가?

주요 결과

Methods	Param (M)	Organs	Ave	Sp	Ki	Gb	Es	Li	St	Ao	IVC	PSV	Pa	AG
SETR (ViT-B/16-rand)	100.5	95.2	92.3	55.6	71.3	96.2	80.2	89.7	83.9	68.9	68.7	60.5	78.4
SETR (ViT-B/16-pre)	100.5	94.8	91.7	55.2	70.9	96.2	76.9	89.3	82.4	69.6	70.7	58.7	77.8
CoTr w/o CNN-encoder	21.9	95.2	92.8	59.2	72.2	96.3	81.2	89.9	85.1	71.9	73.3	61.0	79.8
CoTr w/o DeTrans	32.6	96.0	92.6	63.8	77.9	97.0	83.6	90.8	87.8	76.7	81.2	72.6	83.6
APSS	45.5	96.5	93.8	65.6	78.1	97.1	84.0	91.1	87.9	77.0	82.6	73.9	84.3
PP	33.9	96.1	93.1	64.3	77.4	97.0	85.3	90.8	87.4	77.2	81.9	72.8	83.9
Non-local	32.8	96.3	93.7	64.6	77.9	97.1	84.1	90.8	87.7	77.2	82.1	73.3	84.1
TransUnet	43.5	95.9	93.7	63.1	77.8	97.0	86.2	91.0	87.8	77.8	81.6	73.9	84.2
CoTr ∗	27.9	96.4	94.0	66.2	76.4	97.0	84.2	90.3	87.6	76.3	80.8	72.9	83.8
CoTr †	36.9	96.2	93.8	66.5	78.6	97.1	86.9	90.8	87.8	77.7	82.8	73.2	84.7
CoTr	41.9	96.3	93.9	66.6	78.0	97.1	88.2	91.2	88.0	78.1	83.1	74.1	85.0

• CoTr은 BCV 3D 다기관 분할에서 CNN만 사용한 경우, Transformer만 사용한 경우, 및 다른 하이브리드 기준선들을 능가한다.
• Deformable self-attention을 갖춘 CoTr은 다중 스케일 고해상도 특징 맵의 처리에서 계산 및 공간 복잡성을 줄이면서도 성능을 달성한다.
• 작은 CNN 인코더를 갖는 CoTr ∗, CoTr † 버전은 하이브리드 인코더의 이점을 보여주며 순수 Transformer 기반 인코더 대비 강력한 성능을 보인다.
• DeTrans를 전통적 맥락 모듈(ASPP, PP, Non-local)로 대체하면 Dice 점수가 낮아져 변형 가능한 Transformer의 이점이 강조된다.
• CoTr은 11개의 장기에 걸쳐 평균 Dice를 안정적으로 향상시키며, 담낭 및 췌장에 특히 강한 성능을 보이고 3D 설정에서 TransUNet과 경쟁력 있거나 우수한 성능을 달성한다.
• 학습 시간 및 추론 시간 효율성: 학습 ~2일 GTX 2080Ti에서; 추론은 48×192×192 부피당 30 ms 미만.