[논문 리뷰] Medical Image Segmentation Using Squeeze-and-Expansion Transformers

Shaohua Li, Xiuchao Sui외 4명|arXiv (Cornell University)|2021. 05. 20.

Advanced Neural Network Applications참고 문헌 43인용 수 23

한 줄 요약

이 논문은 효과적이고 넓은 수신장( receptive field)을 확보하면서도 높은 공간 해상도를 유지하는 Squeeze-and-Expansion Transformer를 사용하는 트랜스포머 기반의 의료 영상 분할 프레임워크인 Segtran을 제안한다. 유연한 사전 학습된 사인파 위치 인코딩과 이중 경로 어텐션 메커니즘을 통합함으로써, Segtran은 2D 및 3D 의료 영상 작업 전반에서 U-Net 및 그 변종보다 뛰어난 성능을 보이며, 최소한의 하이퍼파rameter 튜닝으로도 최신 기술 수준의 정확도와 강력한 도메인 간 일반화 능력을 입증한다.

ABSTRACT

Medical image segmentation is important for computer-aided diagnosis. Good segmentation demands the model to see the big picture and fine details simultaneously, i.e., to learn image features that incorporate large context while keep high spatial resolutions. To approach this goal, the most widely used methods -- U-Net and variants, extract and fuse multi-scale features. However, the fused features still have small "effective receptive fields" with a focus on local image cues, limiting their performance. In this work, we propose Segtran, an alternative segmentation framework based on transformers, which have unlimited "effective receptive fields" even at high feature resolutions. The core of Segtran is a novel Squeeze-and-Expansion transformer: a squeezed attention block regularizes the self attention of transformers, and an expansion block learns diversified representations. Additionally, we propose a new positional encoding scheme for transformers, imposing a continuity inductive bias for images. Experiments were performed on 2D and 3D medical image segmentation tasks: optic disc/cup segmentation in fundus images (REFUGE'20 challenge), polyp segmentation in colonoscopy images, and brain tumor segmentation in MRI scans (BraTS'19 challenge). Compared with representative existing methods, Segtran consistently achieved the highest segmentation accuracy, and exhibited good cross-domain generalization capabilities. The source code of Segtran is released at https://github.com/askerlee/segtran.

연구 동기 및 목표

• U-Net 및 그 변종이 작은 효과적 수신장을 가짐으로써 장거리 맥락을 포착하지 못하는 한계를 해결하기 위해.
• 트랜스포머의 무한한 효과적 수신장을 활용하면서도 높은 공간 해상도를 유지함으로써 의료 영상 분할 성능을 향상시키기 위해.
• 자연어 사전 학습에서 비롯된 제약를 극복하고 의료 영상 분할에 특화된 트랜스포머 아키텍처를 설계하기 위해.
• 새로운 위치 인코딩 기반의 연속성 인덕티브 바이어스를 통합함으로써 도메인 간 일반화 능력을 향상시키기 위해.
• 다양한 의료 영상 작업 전반에서 최소한의 아키텍처 수정과 하이퍼파rameter 튜닝으로 최신 기술 수준의 성능을 달성하기 위해.

제안 방법

• 자기 어텐션 행렬을 정규화하여 계산 부담을 줄이고 학습 안정성을 향상시키는 압축된 어텐션 블록을 갖춘 압축-확장 트랜스포머를 제안한다.
• 특징을 다수의 헤드로 투영하고 어텐션 후 융합함으로써 다양화된 표현을 학습하는 확장 블록을 도입한다.
• 공간 연속성을 강조하는 학습 가능한 사인파 위치 인코딩을 사용하여 표준 학습 가능한 또는 고정된 사인파 인코딩보다 이미지 데이터에 대한 인덕티브 바이어스를 향상시킨다.
• 공간 해상도를 유지하고 복원하기 위해 트랜스포머 블록 이전 및 이후에 피처 피라미드 네트워크(FPNs)를 사용하는 U-Net 유사 인코더-디코더 아키텍처를 채택한다.
• 초기 특징을 추출하기 위해 CNN 백본(예: ResNet-101 또는 EfficientNet-B4)을 사용하며, 이를 트랜스포머 처리를 위해 시퀀스로 평탄화한다.
• 픽셀 좌표 기반의 위치 인코딩을 적용하여 자기 어텐션 계산 중 공간적 구조를 유지한다.

실험 결과

연구 질문

• RQ1 향상된 어텐션 메커니즘을 갖춘 트랜스포머 기반 아키텍처가 제한된 맥락 포착 능력을 가진 U-Net 및 그 변종보다 의료 영상 분할 작업에서 뛰어난 성능을 보일 수 있는가?
• RQ2 압축-확장 트랜스포머 설계가 표준 트랜스포머에 비해 의료 영상에서 특징 표현과 분할 정확도를 향상시키는가?
• RQ3 제안된 학습 가능한 사인파 위치 인코딩이 의료 영상 분할에서 인덕티브 바이어스와 모델 일반화 능력을 얼마나 향상시키는가?
• RQ4 REFUGE20과 RIM-One과 같은 다른 영상 특성을 가진 데이터셋에서 테스트했을 때, Segtran은 도메인 간 일반화 성능이 어떻게 나타나는가?
• RQ5 ImageNet에서의 사전 학습이 의료 영상 분할 모델에 어떤 영향을 미치며, 다양한 아키텍처에서 성능에 어떤 영향을 미치는가?

주요 결과

• Segtran은 REFUGE’20 챌린지(망막판/망막구 분할)에서 최고의 Dice 스코어를 기록하여 테스트 세트에서 평균 Dice가 0.938로 상위 5개 팀 내에 포함되었다.
• BraTS’19 챌린지(3D 뇌 종양 분할)에서 Segtran은 U-Net과 DeepLabV3+를 모두 능가하여 검증 세트에서 평균 Dice 스코어 0.875를 기록했다.
• EfficientNet-B4를 백본으로 사용했을 때, Segtran은 모든 작업의 평균 병합 결과에서 Dice 스코어 0.909를 기록하여 U-Net(0.875)과 TransU-Net(0.901)보다 뚜렷이 뛰어났다.
• Segtran은 도메인 간 일반화 능력이 가장 뛰어났으며, REFUGE20에서 RIM-One으로의 전이 시 Dice 스코어가 14.2% 감소에 그쳤다. 이는 U-Net의 18.2%와 DeepLabV3+의 19.3%보다 낮은 감소율이었다.
• ImageNet에서의 사전 학습은 Segtran의 평균 Dice 스코어를 약 2.5% 향상시켜, 데이터가 적은 환경에서 사전 학습의 강력한 이점을 입증했다.
• ResNet-101을 사용할 경우 FLOPs와 파라미터 수가 더 높지만, EfficientNet-B4와 조합했을 때 Segtran은 FLOPs를 71.3G로 줄이고 파라미터를 93.1M로 줄여 최적의 효율성과 성능을 달성했다.