[논문 리뷰] NexToU: Efficient Topology-Aware U-Net for Medical Image Segmentation

Pengcheng Shi, Xutao Guo외 3명|arXiv (Cornell University)|2023. 05. 25.

Medical Imaging and Analysis인용 수 5

한 줄 요약

NexToU는 효율적이고 토폴로지 인식형 의료 영상 분할을 위한 새로운 하이브리드 U-Net 아키텍처를 제안한다. 이는 시각 GNN 모듈인 Pool GNN와 Swin GNN을 통합하여 전역적이고 국소적인 토폴로지 표현을 활용하며, 계산 비용을 $c(c-1)$에서 $2(c-1)$로 감소시키는 이진 토폴로지 상호작용(BTI) 모듈을 도입한다. NexToU는 세 가지 다양한 데이터셋에서 최신 기술(SOTA) 성능을 달성하여 더 높은 Dice 점수와 더 낮은 Hausdorff 거리로 개선된 결과를 보였다.

ABSTRACT

Convolutional neural networks (CNN) and Transformer variants have emerged as the leading medical image segmentation backbones. Nonetheless, due to their limitations in either preserving global image context or efficiently processing irregular shapes in visual objects, these backbones struggle to effectively integrate information from diverse anatomical regions and reduce inter-individual variability, particularly for the vasculature. Motivated by the successful breakthroughs of graph neural networks (GNN) in capturing topological properties and non-Euclidean relationships across various fields, we propose NexToU, a novel hybrid architecture for medical image segmentation. NexToU comprises improved Pool GNN and Swin GNN modules from Vision GNN (ViG) for learning both global and local topological representations while minimizing computational costs. To address the containment and exclusion relationships among various anatomical structures, we reformulate the topological interaction (TI) module based on the nature of binary trees, rapidly encoding the topological constraints into NexToU. Extensive experiments conducted on three datasets (including distinct imaging dimensions, disease types, and imaging modalities) demonstrate that our method consistently outperforms other state-of-the-art (SOTA) architectures. All the code is publicly available at https://github.com/PengchengShi1220/NexToU.

연구 동기 및 목표

• 기존의 CNN 및 Transformer 백본이 전역적 맥락과 토폴로지 일관성에서 어려움을 겪는 복잡하고 비정형적인 해부학적 구조—특히 관형 혈관—을 분할하는 데 도전한다.
• 현재 모델들이 해부학적 구조 간 포함 및 배제 관계를 모델링하지 못해 잘못된 분할과 공간적 혼동을 일으키는 한계를 극복한다.
• 그래프 신경망(GNNs)을 U-Net 아키텍처에 통합하여 의료 영상 내 비유클리드적이고 토폴로지적 관계를 포착하면서도 계산 효율성을 유지한다.
• 이진 트리 구조를 사용해 계층적 관계를 인코딩하는 경량이며 미분 가능한 토폴로지 상호작용 메커니즘(BTI)을 개발하여 이전의 TI 모듈 대비 FLOPs를 감소시킨다.
• 더 적은 파라미터로 다수의 2D 및 3D 의료 영상 데이터셋에서 최신 기술 성능을 달성하며, 경계 선명도를 향상시킨다.

제안 방법

• Pool GNN는 전역적 토폴로지 표현을 위해, Swin GNN는 국소적 창문 기반 특징 학습을 위해 통합한 하이브리드 U-Net 아키텍처인 NexToU를 제안한다.
• Pool GNN를 활용해 먼 해부학적 영역 간 장거리 의존성과 구조적 맥락을 포착하기 위해 그래프 표현에서 핵심 노드를 식별한다.
• 이동 창 분할 및 역연산을 적용한 Swin GNN를 사용해 혈관과 같은 비정형으로 생긴 국소적 구조를 효율적으로 모델링한다.
• 이중 트리 논리로 재구성된 이진 토폴로지 상호작용(BTI) 모듈을 도입하여 해부학적 구조 간 계층적 포함 및 배제 관계를 인코딩한다.
• c-클래스 분할에서 기존 TI 모듈의 계산 복잡도를 $c imes (c-1)$에서 $2 imes (c-1)$로 감소시켜 FLOPs를 크게 낮추면서도 토폴로지 인식 능력을 유지한다.
• BTI 모듈을 손실 함수에 통합하여 훈련 중에 토폴로지 일관성을 강제함으로써 클래스 간 경계에서의 분할 정확도를 향상시킨다.

실험 결과

연구 질문

• RQ1 하이브리드 CNN-GNN 아키텍처는 의료 영상 분할에서 전역적이고 국소적인 토폴로지 표현을 효과적으로 학습할 수 있는가?
• RQ2 해부학적 구조 간의 토폴로지 제약—예를 들어 포함 및 배제 관계—는 어떻게 효율적으로 딥 러닝 모델에 인코딩할 수 있는가?
• RQ3 제안된 BTI 모듈은 표준 토폴로지 상호작용 모듈 대비 계산 비용을 줄이고 분할 정확도를 향상시키는가?
• RQ4 NexToU는 다양한 영상 모odalities와 해부학적 구조—2D 및 3D 데이터셋, 다양한 질환 유형 및 해상도를 포함—에서 일반화 가능한가?
• RQ5 토폴로지 인식 학습을 통합할 경우, 뇌 혈관과 같은 복잡한 구조에서 경계 선명도 향상과 개인 간 변동성 감소에 어느 정도 기여하는가?

주요 결과

• BTCV 데이터셋에서 NexToU는 평균 Dice 스코어 계수(DSC) 87.84%와 Hausdorff 거리(HD) 6.33 mm를 기록하여 nnU-Net, Swin-Unet, PHTrans를 포함한 모든 SOTA 방법을 초월했다.
• ICA 데이터셋에서 NexToU는 74.19%의 DSC를 기록하여 비교된 모든 방법 중 최고 수준이었으며, AcomA 및 PcomA와 같은 작은 복잡한 혈관의 분할에서 뚜렷한 향상이 있었다.
• RAVIR 데이터셋에서 NexToU는 DSC 78.21%를 기록하여 nnU-Net(77.54%) 및 U-Net++(75.35%)를 능가했으며, 망막 혈관 분할에서 뛰어난 성능을 보였다.
• 절단 실험 결과, Pool GNN 및 Swin GNN 모듈을 추가함으로써 DSC가 각각 1.48%와 1.78% 향상되었고, BTI 모듈을 통한 추가 개선으로 DSC는 0.61% 향상되고 HD는 1.02 mm 감소했다.
• 전체 NexToU 모델는 오직 23.06M 파라미터로도 뛰어난 성능을 달성하여 높은 파라미터 효율성과 다양한 해부학적 구조 및 영상 모달리티에 대한 강력한 일반화 능력을 입증했다.
• 시각적 비교 결과, NexToU는 중간 뇌 혈관 및 비장에서의 잘못된 분류와 공간적 혼동을 줄였고, PHTrans 및 nnFormer보다 작은 혈관의 연결성을 더 잘 유지했다.