[논문 리뷰] UCTransNet: Rethinking the Skip Connections in U-Net from a Channel-wise Perspective with Transformer

Haonan Wang, Peng Cao외 2명|arXiv (Cornell University)|2021. 09. 09.

Advanced Neural Network Applications인용 수 44

한 줄 요약

UCTransNet은 평범한 U-Net의 스킵 연결을 채널 트랜스포머(CTrans)로 바꿔 다중 스케일 인코더 피처를 융합하고 디코더 피처와 정렬하여 데이터 세트 전반의 의료 영상 분할 성능을 향상합니다.

ABSTRACT

Most recent semantic segmentation methods adopt a U-Net framework with an encoder-decoder architecture. It is still challenging for U-Net with a simple skip connection scheme to model the global multi-scale context: 1) Not each skip connection setting is effective due to the issue of incompatible feature sets of encoder and decoder stage, even some skip connection negatively influence the segmentation performance; 2) The original U-Net is worse than the one without any skip connection on some datasets. Based on our findings, we propose a new segmentation framework, named UCTransNet (with a proposed CTrans module in U-Net), from the channel perspective with attention mechanism. Specifically, the CTrans module is an alternate of the U-Net skip connections, which consists of a sub-module to conduct the multi-scale Channel Cross fusion with Transformer (named CCT) and a sub-module Channel-wise Cross-Attention (named CCA) to guide the fused multi-scale channel-wise information to effectively connect to the decoder features for eliminating the ambiguity. Hence, the proposed connection consisting of the CCT and CCA is able to replace the original skip connection to solve the semantic gaps for an accurate automatic medical image segmentation. The experimental results suggest that our UCTransNet produces more precise segmentation performance and achieves consistent improvements over the state-of-the-art for semantic segmentation across different datasets and conventional architectures involving transformer or U-shaped framework. Code: https://github.com/McGregorWwww/UCTransNet.

연구 동기 및 목표

• U-Net에서의 간단한 스킵 연결이 보편적으로 유익하지 않으며 일부 데이터 세트에서 성능을 해칠 수 있음을 시연합니다.
• 다중 스케일 인코더 피처를 디코더 피처와 더 잘 융합하기 위한 채널 단위 Transformer 기반 스킵 메커니즘(CCT + CCA)을 제안합니다.
• 채널 단위 융합이 의미적 및 해상도 간극을 줄이고 데이터 세트 전반의 분할 정확도를 향상시키는지 보여줍니다.
• GlaS, MoNuSeg, Synapse 데이터 세트에서 UCTransNet을 평가하고 강력한 U-Net 및 Transformer 기반 기준과 비교합니다.

제안 방법

• 표준 스킵 연결을 Channel Transformer(CTrans)로 대체합니다. 이는 다중 스케일 인코더 피처 융합을 위한 CCT(Channel-wise Cross Fusion Transformer)와 변환기 출력과 함께 디코더 피처 융합을 위한 CCA(Channel-wise Cross Attention)로 구성됩니다.
• CCT는 네 가지 스킵 레이어를 패치로 토큰화하고, 연결된 키/값을 사용한 다중 헤드 채널-와이즈 교차 어텐션을 수행하며, L 레이어에 걸쳐 잔여 연결이 있는 MLP를 적용하여 피처를 융합합니다.
• CCA는 O_i와 D_i를 풀링하여 채널 어텐션 맵을 계산하고, 업샘플링된 디코더 피처와의 연결 전에 O_i를 재캘리브레이션합니다.

실험 결과

연구 질문

• RQ1 채널-와이즈 트랜스포머 기반의 스킵 연결이 의료 영상 분할에서 전통적 스킵 연결을 능가할 수 있는가?
• RQ2 다중 스케일 채널-와이즈 융합(CCT)이 디코더 인지 융합(CCA)과 어떻게 상호작용하여 인코더와 디코더 간 의미적 및 해상도 간극을 좁히는가?
• RQ3 UCTransNet과 그 CTrans 모듈이 여러 의료 영상 데이터 세트 및 아키텍처에서 최첨단 기준 대비 일관된 개선을 제공하는가?

주요 결과

Method	GlaS Dice (%)	GlaS IoU (%)	MoNuSeg Dice (%)	MoNuSeg IoU (%)	Synapse Dice (%)	Synapse HD (mm)
U-Net	85.45	74.78	76.45	62.86	-	-
UNet++	87.56	79.13	77.01	63.04	-	-
AttUNet	88.80	80.69	76.67	63.47	-	-
MRUNet	88.73	80.89	78.22	64.83	-	-
TransUNet	88.40	80.40	78.53	65.05	-	-
Swin-Unet	89.58	82.06	77.69	63.77	-	-
Ours (UCTransNet w/o CCA)	78.99	30.29	78.23	26.75	-	-
Ours (UCTransNet)	90.18	82.96	79.08	65.50	-	-

• UCTransNet은 GlaS와 MoNuSeg에서 Dice 및 IoU 이득으로 기준선보다 현저한 개선을 달성하고 Synapse에서도 Dice 및 Hausdorff 거리가 개선된 것으로 보고됩니다.
• Ablation 연구에서 Baseline+CCT+CCA가 Dataset 전반에서 Baseline, Baseline+CCT, Baseline+CCA보다 일관되게 우수합니다.
• CCT의 스킵-스케일 입력 수를 늘리면 다중 스케일 피처 융합의 가치를 시사합니다.
• 교차 어텐션 시각화는 어떤 인코더 레벨이 분할에 가장 크게 기여하는지 보여주며, 이는 스킵 연결 분석 결과와 일치합니다.
• UCTransNet의 사전 학습은 MoNuSeg 및 Synapse에서 수렴 속도 및 최종 성능을 추가로 향상시킵니다.