[논문 리뷰] VMambaMorph: a Multi-Modality Deformable Image Registration Framework based on Visual State Space Model with Cross-Scan Module

Ziyang Wang, Jian-Qing Zheng외 2명|arXiv (Cornell University)|2024. 04. 07.

Advanced Image and Video Retrieval Techniques인용 수 4

한 줄 요약

VMambaMorph는 3D-최적화된 크로스스캔 모듈과 순환적 정밀조정 전략을 갖춘 하이브리드 시각 Mamba-CNN 아키텍처를 활용한 새로운 3D 다중 모odal성 변형 이미지 정렬 프레임워크를 제안한다. 공개된 MR-CT 데이터셋에서 최신 기술 수준의 정렬 정확도를 달성하며, VoxelMorph, TransMorph 및 MambaMorph를 능가하여 Dice 점수 82.94±2.01을 기록하고 계산 효율성도 향상시켰다.

ABSTRACT

Image registration, a critical process in medical imaging, involves aligning different sets of medical imaging data into a single unified coordinate system. Deep learning networks, such as the Convolutional Neural Network (CNN)-based VoxelMorph, Vision Transformer (ViT)-based TransMorph, and State Space Model (SSM)-based MambaMorph, have demonstrated effective performance in this domain. The recent Visual State Space Model (VMamba), which incorporates a cross-scan module with SSM, has exhibited promising improvements in modeling global-range dependencies with efficient computational cost in computer vision tasks. This paper hereby introduces an exploration of VMamba with image registration, named VMambaMorph. This novel hybrid VMamba-CNN network is designed specifically for 3D image registration. Utilizing a U-shaped network architecture, VMambaMorph computes the deformation field based on target and source volumes. The VMamba-based block with 2D cross-scan module is redesigned for 3D volumetric feature processing. To overcome the complex motion and structure on multi-modality images, we further propose a fine-tune recursive registration framework. We validate VMambaMorph using a public benchmark brain MR-CT registration dataset, comparing its performance against current state-of-the-art methods. The results indicate that VMambaMorph achieves competitive registration quality. The code for VMambaMorph with all baseline methods is available on GitHub.

연구 동기 및 목표

• 기존 딥러닝 기반 이미지 정렬 방법이 장거리 상관관계를 포착하는 데서와 복잡한 다중 모달리티 변형을 다루는 데에서 약점을 해결하기 위해.
• 3D 체적 영상 정렬에서 장거리 특징 모델링을 향상시키기 위해 크로스스캔 모듈을 갖춘 시각 상태공간모델(VMamba)의 적용을 탐색하기 위해.
• 3D 다중 모달리티 정렬에 특화된 향상된 특징 추출과 변형장 학습 능력을 갖춘 하이브리드 VMamba-CNN U-Net 아키텍처를 설계하기 위해.
• 정확도를 향상시키면서도 디퍼모르픽 성질을 유지하는 순환적 정밀조정 및 추론 프레임워크를 개발하기 위해.
• 공개된 MR-CT 정렬 벤치마크에서 제안된 방법을 검증하고, VoxelMorph, TransMorph 및 MambaMorph와 같은 최신 기술 수준의 방법들과 비교하기 위해.

제안 방법

• 3D VMamba 블록과 재설계된 2D 크로스스캔 모듈을 통합한 U-형 인코더-디코더 아키텍처를 활용하며, 이는 3D 체적 특징 처리에 적합하도록 조정되었다.
• 세밀한 특징 추출기가 도입되어 공간적 특징 표현을 향상시켜, 다양한 모달리티 간에 복잡한 해부학적 구조를 더 잘 포착할 수 있도록 모델 능력을 향상시켰다.
• 반복적 학습 및 추론 전략이 구현되었으며, 예측된 이동장이 피드백으로 사용되어 네트워크가 반복적으로 변형장을 정밀조정한다.
• 손실 함수는 정규화된 상관관계 비유사도 측정값과 부드러움 정규화 항을 조합하여 정확하고 디퍼모르픽 변환을 보장한다.
• 엔드 투 엔드로 학습이 이루어지며, 공개된 뇌 MR-CT 정렬 데이터셋을 사용하고, AdamW 최적화 및 학습률 스케줄링을 통해 학습이 진행된다.
• 변형장은 연속적인 이동 벡터장으로 예측되어, 자이아르디언 행렬식 정규화를 통해 디퍼모르픽 성질을 보장한다.

Figure 1: The Recursive Registration Framework of VMambaMorph.

실험 결과

연구 질문

• RQ1 크로스스캔 모듈을 갖춘 시각 Mamba 아키텍처가 3D 다중 모달리티 의료 영상 정렬에서 장거리 상관관계를 효과적으로 모델링할 수 있는가?
• RQ2 3D-최적화된 VMamba 블록을 3D CNN과 통합함으로써 기존 CNN 또는 트랜스포머 기반 접근법에 비해 정렬 정확도가 얼마나 향상되는가?
• RQ3 제안된 순환 정렬 프레임워크가 변형장 품질을 얼마나 향상시키고 정렬 오차를 줄이는가?
• RQ4 세밀한 특징 추출기가 MR 및 CT 영상 간의 구조적 및 강도 변화를 다루는 데에 상당한 기여를 하는가?
• RQ5 표준 벤치마크에서 VoxelMorph, TransMorph 및 MambaMorph와 같은 최신 기술 수준의 방법들과 비교했을 때 VMambaMorph의 성능과 효율성은 어떠한가?

주요 결과

• VMambaMorph는 공개된 MR-CT 정렬 벤치마크에서 Dice 점수 82.94±2.01을 기록하여, VoxelMorph(62.42±3.29), TransMorph(78.46±2.13), MambaMorph(82.49±1.99)를 모두 능가했다.
• 95% 허프만 거리가 1.35±0.18 mm로 감소하여 기준 모델 대비 뛰어난 공간 정렬 정확도를 보였다.
• 비양수 자이아르디언 행렬식 비율이 1.04±0.05%로 낮게 유지되어 변형장에서 디퍼모르픽 성질이 잘 유지됨을 확인했다.
• K=2인 순환 추론 전략이 정확도와 추론 시간 사이에서 최적의 균형을 이뤘으며, 런타임은 0.19초, GPU 메모리 사용량은 3.93 GB였다.
• 절단 실험 결과, 세밀한 특징 추출기가 기준 모델 대비 최대 4.5%까지 Dice 점수 향상을 이끌었다.
• VMambaMorph는 VoxelMorph, TransMorph 및 MambaMorph보다 손실 값이 더 빨리 안정화되고 더 낮은 수준에 도달하여 더 나은 수렴성과 일반화 능력을 보였다.

Figure 2: The Architecture of VMambaMorph, and the Details of Visual State Space Block.