[논문 리뷰] TransAttUnet: Multi-level Attention-guided U-Net with Transformer for Medical Image Segmentation

Bingzhi Chen, Yishu Liu외 3명|arXiv (Cornell University)|2021. 07. 12.

Advanced Neural Network Applications참고 문헌 45인용 수 84

한 줄 요약

TransAttUnet은 Transformer 기반 셀프 어텐션과 글로벌 공간 어텐션을 다중 스케일 스킵 커넥션과 함께 U-Net에 결합하여 다중 모달리티의 의료 영상 분할 성능을 향상시키고, 최첨단 baselines를 능가한다.

ABSTRACT

Accurate segmentation of organs or lesions from medical images is crucial for reliable diagnosis of diseases and organ morphometry. In recent years, convolutional encoder-decoder solutions have achieved substantial progress in the field of automatic medical image segmentation. Due to the inherent bias in the convolution operations, prior models mainly focus on local visual cues formed by the neighboring pixels, but fail to fully model the long-range contextual dependencies. In this paper, we propose a novel Transformer-based Attention Guided Network called TransAttUnet, in which the multi-level guided attention and multi-scale skip connection are designed to jointly enhance the performance of the semantical segmentation architecture. Inspired by Transformer, the self-aware attention (SAA) module with Transformer Self Attention (TSA) and Global Spatial Attention (GSA) is incorporated into TransAttUnet to effectively learn the non-local interactions among encoder features. Moreover, we also use additional multi-scale skip connections between decoder blocks to aggregate the upsampled features with different semantic scales. In this way, the representation ability of multi-scale context information is strengthened to generate discriminative features. Benefitting from these complementary components, the proposed TransAttUnet can effectively alleviate the loss of fine details caused by the stacking of convolution layers and the consecutive sampling operations, finally improving the segmentation quality of medical images. Extensive experiments on multiple medical image segmentation datasets from different imaging modalities demonstrate that the proposed method consistently outperforms the state-of-the-art baselines. Our code and pre-trained models are available at: https://github.com/YishuLiu/TransAttUnet.

연구 동기 및 목표

• 합성곱 인코더의 지역 편향을 해소하여 의료 영상 분할의 개선을 이끌고자 한다.
• Transformer 기반의 self-aware attention(SAA) 모듈을 제안하며, Transformer Self Attention과 Global Spatial Attention을 결합한다.
• 다중 스케일 디코더 특징을 더 잘 융합하기 위해 다중 스케일 스킵 커넥션을 도입한다.
• 다양한 의학 영상 모달리티 전반에서 방법의 효과를 입증한다.
• 강력한 기저선과 차등 실험(아브레이션 연구)과의 공정한 비교를 제공한다.

제안 방법

• 인코더-디코더 다리 부분에 SAA 모듈을 통합하여 TSA와 GSA를 인코더 특징과 융합한다.
• 다중 헤드 어텐션과 학습된 위치 인코딩을 갖춘 Transformer Self Attention(TSA)을 사용하여 장거리 의존성을 모델링한다.
• 위치 인식 채널 상호작용을 통해 글로벌 맥락을 포착하기 위해 Global Spatial Attention(GSA)을 적용한다.
• 학습 가능한 가중 합성(F_SAA)을 통해 TSA와 GSA 출력을 인코더 특징과 융합한다.
• 해당 의미의 다양성을 가진 기능을 점진적으로 집계하기 위해 잔차 또는 Dense 연결을 사용하여 디코더 블록 간에 다중 스케일 스킵 커넥션을 구현한다.
• 픽셀 단위 정확도와 분할 중첩 간의 균형을 맞추기 위해 Dice와 BCE 손실의 결합(L = alpha*L_BCE + beta*L_Dice)로 학습한다.

실험 결과

연구 질문

• RQ1 Transformer 기반의 셀프 어텐션과 글로벌 공간 어텐션을 통합하는 것이 U-Net 기반의 의료 영상 분할을 개선할 수 있는가?
• RQ2 다중 스케일 스킵 커넥션(잔차/밀집)이 전통적인 캐스케이드 연결보다 미세한 디테일을 더 잘 보존하는가?
• RQ3 TransAttUnet은 다양한 모달리티 및 데이터셋에서 최첨단 기저선과 비교해 어떤 성능을 보이는가?
• RQ4 SAA 모듈과 다중 스케일 융합이 분할 정확도와 경계 정밀도에 어떤 영향을 미치는가?

주요 결과

방법	연도	DICE	IoU	ACC	REC	PRE
U-Net	2015	67.40	54.90	-	70.80	-
Attention U-Net	2018	66.50	56.60	-	71.70	-
R2U-Net	2018	67.90	58.10	-	79.20	-
Att R2UNet	2018	69.10	59.20	-	72.60	-
ResUNet*	2019	79.15	70.15	92.28	82.43	84.77
Channel-UNet*	2019	84.82	75.92	94.10	94.01	81.04
BCDU-Net	2019	85.10	-	-	-	-
FANet	2021	87.31	80.23	-	86.50	92.35
PraNet*	2021	87.46	80.23	95.37	91.28	87.59
DoubleU-Net	2020	89.62	82.12	-	87.80	94.59
Swin-Unet*	2021	89.72	82.90	-	90.32	92.04
SegFormer*	2021	90.24	83.60	-	91.12	92.10
MCTrans	2021	90.35	-	-	-	-
TransAttUnet_C	-	89.25	81.46	95.06	89.90	91.59
TransAttUnet_D	-	90.14	83.04	96.14	90.42	92.17
TransAttUnet_R	-	90.74	83.80	96.38	90.93	92.42

• TransAttUnet 변형이 다수의 데이터셋에서 기본 U-Net 및 여러 기저선들을 능가한다.
• TransAttUnet_R(밀집 스킵 커넥션)은 ISIC-2018에서 최고 Dice 점수인 90.74%를 달성한다.
• Transformer Self Attention(TSA)와 Global Spatial Attention(GSA)을 함께 사용하면 두 요소 중 하나를 단독으로 사용할 때보다 맥락 모델링이 더 향상된다.
• 다중 스케일 스킵 커넥션(잔차 또는 Dense)은 한 번의 캐스케이드 연결보다 특징 융합을 더 잘 수행하여 미세한 디테일 손실을 완화한다.
• ISIC-2018에서 MCTrans과 비교할 때 TransAttUnet_R은 Dice 지표에서 개선을 보여준다(90.74% vs. 90.35%).