[논문 리뷰] TransClaw U-Net: Claw U-Net with Transformers for Medical Image Segmentation

Yao‐Wen Chang, Menghan Hu외 2명|arXiv (Cornell University)|2021. 07. 12.

Advanced Neural Network Applications인용 수 10

한 줄 요약

이 논문은 인코더에서 국소적 특징 추출을 위한 합성곱 신경망(CNNs)과 전역적 맥락 모델링을 위한 비전 트랜스포머(ViT) 블록을 통합한 하이브리드 딥 러닝 아키텍처인 TransClaw U-Net을 제안한다. CNN 기반의 특징 학습과 자기주의 기반 메커니즘을 결합함으로써, 512×512 입력 해상도에서 Synapse 다기관 CT 분할 데이터셋에서 최신 기술(SOTA) 성능을 달성하며, 평균 DSC는 80.39%를 기록한다.

ABSTRACT

In recent years, computer-aided diagnosis has become an increasingly popular topic. Methods based on convolutional neural networks have achieved good performance in medical image segmentation and classification. Due to the limitations of the convolution operation, the long-term spatial features are often not accurately obtained. Hence, we propose a TransClaw U-Net network structure, which combines the convolution operation with the transformer operation in the encoding part. The convolution part is applied for extracting the shallow spatial features to facilitate the recovery of the image resolution after upsampling. The transformer part is used to encode the patches, and the self-attention mechanism is used to obtain global information between sequences. The decoding part retains the bottom upsampling structure for better detail segmentation performance. The experimental results on Synapse Multi-organ Segmentation Datasets show that the performance of TransClaw U-Net is better than other network structures. The ablation experiments also prove the generalization performance of TransClaw U-Net.

연구 동기 및 목표

• 의료 영상 분할에서 CNN의 장거리 공간적 의존성 파악 능력의 한계를 해결하기 위해.
• 자기주의 기반 메커니즘을 통해 인코더에 전역적 맥락 모델링을 통합하여 분할 정확도를 향상시키기 위해.
• 디코더에서 원래의 Claw U-Net 아키텍처의 세밀한 특징 복원 능력을 유지하기 위해.
• 입력 해상도, 패치 크기, 스킵 커넥션 수가 모델 성능에 미치는 영향을 평가하기 위해.
• Abdominal CT 스캔에서 다양한 기관에 걸쳐 제안된 아키텍처의 일반화 능력과 강건성을 입증하기 위해.

제안 방법

• 인코더는 이중 스트림 아키텍처를 사용한다: 한 분지는 표준 합성곱 레이어를 통해 얕은 공간적 특징을 추출하고, 다른 분지는 특징 맵을 패치로 처리하여 자기주의 기반의 전역 표현 학습을 수행한다.
• 비전 트랜스포머 모듈은 인코딩된 특징 맵 내 이미지 패치 간의 장거리 의존성을 학습하기 위해 다중 헤드 자기주의를 적용한다.
• 디코더는 분해 해상도 복원 과정에서 세밀한 공간적 세부 정보를 유지하기 위해 Claw U-Net의 하향식 업샘플링 구조를 유지한다.
• 해당 인코더 및 디코더 레이어 간에 스킵 커넥션을 설정하여 특징 융합을 촉진하고 정위치 정확도를 향상시킨다.
• 모델는 Synapse 데이터셋에서 최적의 성능를 달성하기 위해 초모수를 조정한 교차 엔트로피 손실과 Dice 손실을 사용하여 엔드 투 엔드로 훈련된다.
• 제거 분석을 통해 입력 크기, 패치 크기, 스킵 커넥션 수가 분할 성능에 미치는 영향을 체계적으로 평가한다.

Fig. 1 : The architecture of TransClaw U-Net

실험 결과

연구 질문

• RQ1 U-Net 인코더에 트랜스포머를 통합함으로써 표준 CNN과 비교해 장거리 특징 학습 능력이 어떻게 향상되는가?
• RQ2 분할 정확도와 계산 비용을 균형 잡기 위해 최적의 입력 영상 해상도는 무엇인가?
• RQ3 패치 크기가 비전 트랜스포머 구성 요소의 시퀀스 길이와 모델 성능에 어떤 영향을 미치는가?
• RQ4 스킵 커넥션 수를 다양하게 조절했을 때, 다양한 기관의 분할 정확도에 어떤 영향을 미치는가?
• RQ5 제안된 TransClaw U-Net은 복부 CT 스캔의 다양한 해부학적 구조에 대해 잘 일반화되는가?

주요 결과

• 512×512 입력 해상도를 사용할 경우 TransClaw U-Net은 Synapse 데이터셋에서 평균 DSC 80.39%를 달성하며, 기준 모델인 224×224 해상도 대비 78.09% DSC를 기록한 것보다 뛰어난 성능을 보였다.
• 입력 해상도를 224×224에서 512×512로 증가시키면 DSC가 유의미하게 향상되며, 이는 고해상도 입력이 특징 표현과 분할 정확도를 향상시킨다는 것을 시사한다.
• 16×16 패치 크기에서 모델은 최고의 DSC 78.09%를 기록했고, 더 큰 패치 크기(24×24)는 계산 비용은 증가시키지만 성능은 77.27%로 감소시켰다.
• 스킵 커넥션 수는 강력한 정성적 영향을 미친다: 스킵 커넥션을 완전히 제거하면 대동맥의 DSC가 10.8% 감소하고 췌장의 DSC는 10.16% 감소하여 그 중요성을 확인한다.
• 제거 분석 결과, 더 큰 입력 크기, 더 작은 패치 크기, 더 많은 스킵 커넥션은 성능 향상에 기여하지만, 이는 메모리 소비와 훈련 시간 증가를 수반한다.
• 모델는 Synapse 데이터셋의 여덟 기관 전반에서 강력한 일반화 능력을 보였으며, 최적 설정에서 간(95.06%), 비장(91.16%), 위(82.82%)의 DSC가 일관되게 향상되었다.

Fig. 2 : Comparison of segmentation results of different models on the Synapse multi-organ CT dataset.