[논문 리뷰] Affine Medical Image Registration with Coarse-to-Fine Vision Transformer

Tony C. W. Mok, Albert C. S. Chung|arXiv (Cornell University)|2022. 03. 29.

Medical Image Segmentation Techniques인용 수 5

한 줄 요약

이 논문은 3D 애핀 의료 영상 정렬을 위한 새로운 코arse-to-fine 비전 트랜스포머인 C2FViT를 제안한다. 이는 전역 자기주의(global self-attention)와 다중 해상도 특징 학습을 활용하여 CNN 기반 방법에 비해 정확도, 강건성, 일반화 능력을 향상시킨다. 이는 ANTs 및 Elastix와 같은 전통적 방법과 유사한 성능을 달성하면서도 학습 기반 접근법의 빠른 추론 속도를 유지한다.

ABSTRACT

Affine registration is indispensable in a comprehensive medical image registration pipeline. However, only a few studies focus on fast and robust affine registration algorithms. Most of these studies utilize convolutional neural networks (CNNs) to learn joint affine and non-parametric registration, while the standalone performance of the affine subnetwork is less explored. Moreover, existing CNN-based affine registration approaches focus either on the local misalignment or the global orientation and position of the input to predict the affine transformation matrix, which are sensitive to spatial initialization and exhibit limited generalizability apart from the training dataset. In this paper, we present a fast and robust learning-based algorithm, Coarse-to-Fine Vision Transformer (C2FViT), for 3D affine medical image registration. Our method naturally leverages the global connectivity and locality of the convolutional vision transformer and the multi-resolution strategy to learn the global affine registration. We evaluate our method on 3D brain atlas registration and template-matching normalization. Comprehensive results demonstrate that our method is superior to the existing CNNs-based affine registration methods in terms of registration accuracy, robustness and generalizability while preserving the runtime advantage of the learning-based methods. The source code is available at https://github.com/cwmok/C2FViT.

연구 동기 및 목표

• 대규모 초기 비일치(initial misalignments)를 다루는 데 어려움을 겪는 CNN 기반 애핀 정렬 방법의 한계를 해결하기 위해.
• 지역적 컨볼루션 인덕티브 바이어스보다 전역 기하학적 의존성을 더 효과적으로 포착할 수 있는 학습 기반 애핀 정렬 프레임워크를 개발하기 위해.
• 특히 도전적인 공간 초기화 조건 하에서도 3D 뇌 영상 정렬 작업의 정확도와 강건성을 향상시키기 위해.
• 딥 러닝 기반 방법의 런타임 효율성을 유지하면서도 전통적인 최적화 기반 애핀 정렬의 성능을 matching하거나 초월하기 위해.
• 최소한의 아키텍처 수정으로 다양한 매개변수 기반 정렬 작업에 본 연구의 패러다임이 적응 가능함을 보여주기 위해.

제안 방법

• C2FViT는 다중 해상도에서 영상 쌍을 처리하여 애핀 변환 추정치를 점진적으로 개선하는 코arse-to-fine 아키텍처를 사용한다.
• 장거리 공간적 의존성을 개선하기 위해 국소 인덕티브 바이어스와 전역 자기주의를 결합한 컨볼루션 비전 트랜스포머(CvT) 백본을 사용한다.
• 다중 해상도 전략을 통해 스케일 간 전역 방향성과 세밀한 공간 정렬 특징을 모두 포착할 수 있다.
• 특히 대규모 초기 비일치 조건에서 학습 안정성 향상과 수렴성 향상을 위해 추론 중에 점진적인 공간 변환을 적용한다.
• 학습 안정성 향상과 무작위 공간 초기화에 대한 의존도 감소를 위해 중심질량/중심점(Com) 초기화를 기하학적 사전 지식으로 사용한다.
• 모델 파rameter 수나 계산 비용 증가 없이 해부학적 레이블을 활용한 준지도 학습을 도입하여 일반화 능력을 향상시킨다.

실험 결과

연구 질문

• RQ1 대규모 초기 비일치 조건 하에서 비전 트랜스포머 기반 아키텍처가 CNN 기반 방법보다 성능을 뛰어넘을 수 있는가?
• RQ2 제안된 C2FViT 모델은 훈련 분포 외의 새로운 데이터셋에 대해 얼마나 잘 일반화되는가?
• RQ3 전역 자기주의와 다중 해상도 처리의 통합이 정렬 정확도와 강건성에 얼마나 기여하는가?
• RQ4 전통적인 최적화 기반 애핀 정렬의 성능에 도달하거나 초월하면서도 빠른 추론 속도를 유지할 수 있는가?
• RQ5 진행적 공간 변환과 중심질량 초기화와 같은 아키텍처 구성 요소가 모델 성능과 안정성에 미치는 영향은 어떠한가?

주요 결과

• LPBA 데이터셋에서 C2FViT는 Dice 스코어 0.66과 HD95 2.96를 기록하여 ANTs 및 Elastix와 같은 전통적 방법과 동등한 성능을 달성한다.
• 준지도 학습을 적용한 모델은 OASIS 데이터셋에서 Dice 스코어 0.69와 HD95 2.81을 기록하여 비지도 및 지도 학습 기반 베이스라인을 모두 초월한다.
• C2FViT는 평균 0.09초의 빠른 런타임을 유지하여 ANTs 및 Elastix보다 현저히 빠른 추론 속도를 확보한다.
• 제거 실험(ablation study) 결과, 진행적 공간 변환과 중심질량 초기화가 각각 Dice 스코어 +0.02, HD95 -0.37의 성능 향상을 이끌어냈다.
• 직접 매트릭스 추정보다 분리된 애핀 변환 모델이 약간 더 높은 정확도(Dice +0.02)를 기록했으며, 런타임 오버헤드도 최소한이었다.
• 모델는 훈련 분포 외의 데이터셋에 대해서도 잘 일반화되어 있으며, 훈련 데이터 외부의 데이터에 대해 테스트한 결과 강건성이 입증되었다.