[논문 리뷰] PC-SwinMorph: Patch Representation for Unsupervised Medical Image Registration and Segmentation

Lihao Liu, Zhening Huang외 3명|arXiv (Cornell University)|2022. 03. 10.

Advanced Neural Network Applications인용 수 13

한 줄 요약

PC-SwinMorph는 두 가지 새로운 피처 기반 전략을 사용하여 의료 영상 정렬 및 분할을 위한 통합된 비지도 프레임워크를 제안한다: 국소적이고 전역적인 특징 표현을 향상시키기 위한 피처 기반 대비 학습 방법과 피처 분할 잔상 문제를 제거하기 위한 3D 윈도우/시프트된 윈도우 다중 헤드 자기주의 모듈. 이 방법은 CANDI 및 LPBA40 데이터셋에서 기준 모델 대비 최대 4% 향상된 Dice 점수를 달성하여 최신 기술 수준의 성능을 보였다.

ABSTRACT

Medical image registration and segmentation are critical tasks for several clinical procedures. Manual realisation of those tasks is time-consuming and the quality is highly dependent on the level of expertise of the physician. To mitigate that laborious task, automatic tools have been developed where the majority of solutions are supervised techniques. However, in medical domain, the strong assumption of having a well-representative ground truth is far from being realistic. To overcome this challenge, unsupervised techniques have been investigated. However, they are still limited in performance and they fail to produce plausible results. In this work, we propose a novel unified unsupervised framework for image registration and segmentation that we called PC-SwinMorph. The core of our framework is two patch-based strategies, where we demonstrate that patch representation is key for performance gain. We first introduce a patch-based contrastive strategy that enforces locality conditions and richer feature representation. Secondly, we utilise a 3D window/shifted-window multi-head self-attention module as a patch stitching strategy to eliminate artifacts from the patch splitting. We demonstrate, through a set of numerical and visual results, that our technique outperforms current state-of-the-art unsupervised techniques.

연구 동기 및 목표

• 지도 학습이 부족하여 미세한 해부학적 세부 사항을 포착하지 못하는 비지도 의료 영상 정렬 및 분할의 한계를 해결하기 위해.
• 국소적이고 전역적인 해부학적 모델링을 향상시키기 위해 피처 기반 공간 구조를 활용하여 비지도 학습에서의 특징 표현을 개선하기 위해.
• 3D 의료 영상에서 피처 기반 표현으로 인한 잔상 문제를 해결하기 위해 3D W-MSA 및 SW-MSA 모듈을 활용한 새로운 피처 스티칭 기법을 제안하기 위해.
• 지식 기반 변형 필드나 분할 마스크가 필요 없이 정렬 및 분할 양쪽에서 최신 기술 수준의 성능을 달성하기 위해.
• 피처 기반 표현이 복잡한 의료 영상 작업에서 전역 토큰 임베딩보다 더 의미 있는가를 입증하기 위해.

제안 방법

• 프레임워크는 VoxelMorph를 백본으로 사용하며, 특징 학습을 향상시키기 위해 인코더를 피처 기반 표현 전략으로 대체한다.
• 이미지 간 피처를 비교하여 국소 일관성과 풍부한 특징 표현을 강화하기 위해 피처 단위 대비 학습 전략을 도입한다.
• 피처 스티칭 및 장거리 의존성 복원을 위해 인코더/디코더 이후에 3D 윈도우/시프트된 윈도우 다중 헤드 자기주의 모듈(3D W-MSA 및 3D SW-MSA)을 적용한다.
• 대비 손실과 피처 스티칭을 통합된 프레임워크에 통합하여 동시에 정렬 및 분할 성능을 향상시킨다.
• 3D SW-MSA 모듈은 메인 네트워크 이후 단일 레이어로 적용되어 계산 비용은 최소화하면서도 최대한의 특징 융합을 달성한다.
• 지식 기반 변형 필드나 분할 마스크가 필요 없이 비지도 손실 함수를 사용하여 엔드 투 엔드로 모델을 훈련시킨다.

실험 결과

연구 질문

• RQ1 전역 토큰 기반 방법에 비해 피처 기반 표현이 비지도 의료 영상 정렬 및 분할 성능을 크게 향상시킬 수 있는가?
• RQ2 피처 단위 대비 학습 전략을 도입함으로써 특징 표현과 세밀한 해부학적 정렬이 향상되는가?
• RQ3 3D 윈도우/시프트된 윈도우 다중 헤드 자기주의 기법이 피처 분할 잔상 문제를 효과적으로 줄이고 피처 간 교환을 향상시킬 수 있는가?
• RQ4 피처 대비 학습과 피처 스티칭의 조합이 실제 의료 데이터셋에서 정렬 및 분할 성능을 얼마나 향상시키는가?
• RQ5 제안된 방법은 최신 기술 수준의 성능을 달성하면서도 계산 효율성이 뛰어나다고 할 수 있는가?

주요 결과

• 제안된 PC-SwinMorph 프레임워크는 CANDI 및 LPBA40 데이터셋에서 기준 모델 대비 최대 4% 향상된 Dice 점수를 달성하였다.
• 피처 단위 대비 학습 전략만으로도 대비 학습 기준 모델 대비 2% 향상된 Dice 점수를 기록하여 특징 풍부화에 효과적임을 입증하였다.
• 3D W-MSA 및 3D SW-MSA 모듈의 통합은 CANDI 데이터셋에서 정렬 성능을 3.06% 향상시키고 분할 성능을 2.65% 향상시켰다.
• CANDI 데이터셋에서 분할 성능의 Dice 점수는 0.817, 정렬 성능은 0.812를 기록하여 모든 최신 기술 수준의 비지도 방법을 능가하였다.
• 계산 비용 증가는 극히 미미하여 총 0.6 GFlops만 추가되었으며(416.04에서 416.10 GFlops로 증가), 이는 방법의 효율성을 확인한다.
• 시각적 및 수치적 결과는 피처 기반 접근 방식이 전역 표현 방법보다 더 견고하게 미세한 해부학적 세부 사항을 포착함을 확인하였다.