[논문 리뷰] GCtx-UNet: Efficient Network for Medical Image Segmentation

Khaled Alrfou, Tian Zhao|arXiv (Cornell University)|2024. 06. 09.

Brain Tumor Detection and Classification인용 수 6

한 줄 요약

GCtx-UNet은 전역/로컬 주의 집중을 GC-ViT와 CNN 기반 하향/상향 샘플링을 결합한 경량형 UNet 유사 아키텍처로, 여러 의학 영상 데이터셋에서 모델 복잡도는 낮추면서 경쟁력 있는 분할 정확도를 달성합니다.

ABSTRACT

Medical image segmentation is crucial for disease diagnosis and monitoring. Though effective, the current segmentation networks such as UNet struggle with capturing long-range features. More accurate models such as TransUNet, Swin-UNet, and CS-UNet have higher computation complexity. To address this problem, we propose GCtx-UNet, a lightweight segmentation architecture that can capture global and local image features with accuracy better or comparable to the state-of-the-art approaches. GCtx-UNet uses vision transformer that leverages global context self-attention modules joined with local self-attention to model long and short range spatial dependencies. GCtx-UNet is evaluated on the Synapse multi-organ abdominal CT dataset, the ACDC cardiac MRI dataset, and several polyp segmentation datasets. In terms of Dice Similarity Coefficient (DSC) and Hausdorff Distance (HD) metrics, GCtx-UNet outperformed CNN-based and Transformer-based approaches, with notable gains in the segmentation of complex and small anatomical structures. Moreover, GCtx-UNet is much more efficient than the state-of-the-art approaches with smaller model size, lower computation workload, and faster training and inference speed, making it a practical choice for clinical applications.

연구 동기 및 목표

• 높은 계산 비용 없이 정확한 의료 영상 분할이 필요하다는 점을 다룬다.
• UNet과 같은 아키텍처에서 전역 컨텍스트 모델링과 로컬 주의 집중을 통합한다.
• MedNet 프리트레이닝을 활용해 도메인 내 성능을 ImageNet 프리트레이닝보다 향상시킨다.
• 여러 데이터셋(Synapse, ACDC, Polyp)에서 일반화 및 효율성을 보여준다.

제안 방법

• 로컬 자기 주의와 글로벌 컨텍스트 쿼리를 결합한 GC-ViT 블록으로 장단거리 의존성 모델링.
• 유도 편향 및 채널 간 모델링을 주입하기 위한 다운샘플링(Fused-MBConv) 모듈 도입.
• U자형 아키텍처에서 GC-ViT 기반 인코더–병목–디코더와 스킵 연결을 사용.
• MedNet(의료 영상)에서 GC-ViT를 사전 학습하고 ImageNet 사전 학습과 비교.
• 인코더에서 패치화 층을 사용하여 겹치는 패치를 생성하고 임베딩 투영을 수행.

실험 결과

연구 질문

• RQ1 GC-ViT 기반 블록이 최첨단 CNN/Transformer 기반 모델보다 적은 매개변수로도 경쟁력 있는 분할 성능을 달성할 수 있는가?
• RQ2 도메인 특화 의료 데이터(MedNet)에서의 사전 학습이 자연 이미지(ImageNet) 프리트레이닝보다 분할 정확도를 향상시키는가?
• RQ3 다양한 의료 영상 작업(CT, MRI, 폴립 이미지)에서 DSC와 HD 측면에서 GCtx-UNet의 성능은 어떠한가?
• RQ4 업샘플링/다운샘플링 설계 및 하이퍼파라미터가 분할 성능에 미치는 영향은?
• RQ5 제안된 아키텍처가 동료 대비 모델 크기, FLOPs, 훈련 시간 및 추론 속도 측면에서 효율적인가?

주요 결과

알고리즘	DSC	HD	대동맥	담낭	신장(왼쪽)	신장(오른쪽)	간	췌장	비장	위
U-Net	76.85	39.70	89.07	69.72	77.77	68.60	93.43	53.98	86.67	75.58
Att-UNet	77.77	36.02	89.55	68.88	77.98	71.11	93.57	58.04	87.30	75.75
Swin-UNet	79.13	21.55	85.47	66.53	83.2	79.61	94.29	56.58	90.66	76.60
TransDeepLab	80.16	21.25	86.04	69.16	84.08	79.88	93.53	61.19	89.00	78.40
MISSFormer	81.96	18.20	86.99	68.65	85.21	82.00	94.41	65.67	91.92	80.81
TransUNet	77.48	31.69	87.23	63.13	81.87	77.02	94.08	55.86	85.08	75.62
GPA-TUNet	80.37	20.55	88.74	65.63	83.51	80.37	94.84	63.89	87.58	78.40
HiFormer	80.39	14.70	86.21	65.69	85.23	79.77	94.61	59.52	90.99	81.08
CS-UNet	83.27	15.26	88.07	71.32	88.00	84.38	94.80	65.64	89.95	83.81
GCtx-UNet1	81.95	16.80	86.96	66.26	87.75	83.86	94.53	61.06	91.42	84.15
GCtx-UNet2	82.39	15.94	86.30	69.32	86.11	81.89	94.64	64.88	91.81	84.15

• GCtx-UNet은 최첨단 결과에 근접한 성능을 달성하면서도 가장 작은 모델 크기(12.34M 매개변수) 및 대상 조사 방법 중 가장 낮은 FLOPs를 기록했다.
• MedNet으로 사전 학습된 GCtx-UNet은 일반적으로 여러 데이터셋에서 ImageNet으로 사전 학습된 변형보다 우수한 성능을 보였다.
• Synapse에서 GCtx-UNet2(MedNet 사전 학습)은 DSC 82.39% 및 HD 15.94 mm를 달성하며, 동료보다 낮은 계산량으로도 최고 수준의 결과를 보였다.
• ACDC에서 GCtx-UNet2는 DSCs 91.23(RV), 89.88(Myocardium), 87.25(LV)를 달성하며 여러 Transformer 기반 및 하이브리드 모델을 능가했다.
• 폴립 데이터셋에서 GCtx-UNet2는 보지 못한 데이터셋에 대한 일반화가 강하고 종종 최상위 또는 거의 최상위 DSC 점수를 기록했다.
• 완화 연구에서 최적의 손실 조합(dice 0.3, cross-entropy 0.7) 및 학습률(0.0001)을 확인했고, SE 블록이 포함된 전치 합성(upsampling)으로 최상의 성능이 나타났다.