[논문 리뷰] Cross-modal Attention for MRI and Ultrasound Volume Registration

Xinrui Song, Hengtao Guo외 7명|arXiv (Cornell University)|2021. 07. 09.

Advanced Neural Network Applications참고 문헌 19인용 수 6

한 줄 요약

이 논문은 MRI와 초음파 영상 간의 공간적 대응을 명시적으로 모델링하는 크로스모달 어텐션 메커니즘을 제안하여 의료 영상 정렬 성능을 향상시킨다. 경량 컨볼루션 네트워크에 이 어텐션 블록을 통합함으로써 기존 최고 성능(SOTA) 모델 대비 모델 크기의 1/10, 런타임의 절반을 사용하면서도 정확도를 최고 수준으로 끌어올렸다. 정렬 오차는 10.17±5.75 mm에서 3.71±1.99 mm로 감소하였다.

ABSTRACT

Prostate cancer biopsy benefits from accurate fusion of transrectal ultrasound (TRUS) and magnetic resonance (MR) images. In the past few years, convolutional neural networks (CNNs) have been proved powerful in extracting image features crucial for image registration. However, challenging applications and recent advances in computer vision suggest that CNNs are quite limited in its ability to understand spatial correspondence between features, a task in which the self-attention mechanism excels. This paper aims to develop a self-attention mechanism specifically for cross-modal image registration. Our proposed cross-modal attention block effectively maps each of the features in one volume to all features in the corresponding volume. Our experimental results demonstrate that a CNN network designed with the cross-modal attention block embedded outperforms an advanced CNN network 10 times of its size. We also incorporated visualization techniques to improve the interpretability of our network. The source code of our work is available at https://github.com/DIAL-RPI/Attention-Reg .

연구 동기 및 목표

• 전립선 생검 유도에서 MRI와 직장내초음파(TRUS) 영상 간 정확한 크로스모달 정렬 문제를 해결한다.
• 다른 영상 모odal리티의 특징 간 공간적 대응을 모델링하는 데 한계가 있는 전통적 CNN의 한계를 극복한다.
• MRI와 TRUS 영상 간의 전역적 특징 대응을 명시적으로 캡처하는 새로운 어텐션 메커니즘을 개발한다.
• 기존 딥러닝 방법에 비해 훨씬 작고 효율적인 네트워크 아키텍처를 통해 정렬 성능을 향상시킨다.
• Grad-CAM과 같은 시각화 기법을 활용하여 정렬 네트워크의 해석 가능성을 향상시킨다.

제안 방법

• 고정 영상인 MRI와 이동 영상인 TRUS 영상의 특징 간 작용하는 비국소 어텐션을 확장한 크로스모달 어텐션 블록을 제안한다.
• 한 영상의 각 특징과 다른 영상의 모든 특징 간 유사도를 계산함으로써 전역적 특징 매칭을 가능하게 한다.
• 3D 영상 정렬을 위해 유사 U-Net 아키텍처의 인코더-디코더 네트워크에 크로스모달 어텐션 블록을 통합한다.
• 입력 영상의 다운샘플링 및 계층적 특징 추출을 위해 컨볼루션 및 맥스풀링 레이어를 갖춘 특징 추출기 사용한다.
• 어텐션 처리된 특징을 바탕으로 다중모odal 정보를 융합하고 공간적 변환을 예측하는 깊은 정렬기로 전달한다.
• 왜곡된 이동 영상과 고정 영상 간 박스별 평균 제곱오차를 기반으로 한 손실 함수를 사용해 네트워크를 엔드 투 엔드로 훈련시킨다.

실험 결과

연구 질문

• RQ1 자기어텐션 메커니즘이 MRI와 TRUS 영상 간 크로스모달 영상 정렬을 향상시킬 수 있는가? 특히 공간적 대응을 명시적으로 모델링함으로써.
• RQ2 제안된 크로스모달 어텐션 블록은 기존의 더 큰 표준 CNN에 비해 훨씬 작은 네트워크로도 더 나은 정렬 성능을 달성할 수 있는가?
• RQ3 입력으로 세그멘테이션 마스크를 사용할 경우, 원본 MRI 영상에 비해 정렬 정확도와 강건성이 향상되는가?
• RQ4 어텐션 메커니즘이 특징 시각화를 통해 네트워크의 해석 가능성에 얼마나 기여하는가?
• RQ5 기존 최고 성능(SOTA) 방법에 비해 더 낮은 추론 시간과 모델 크기로도 뛰어난 성능을 달성할 수 있는가?

주요 결과

• 제안된 Attention-Reg 네트워크는 목표 정렬 오차(TRE)를 3.71±1.99 mm로 감소시켜 기준값 10.17±5.75 mm보다 뚜렷한 향상을 이룩했다.
• 모델 파라미터 수가 1,248,777개에 불과한 이 방법은 MSReg(16M 파라미터)와 DVNet(5.3M 파라미터)를 능가하여 모델 크기의 1/10로 더 높은 정확도를 달성했다.
• 3ms의 추론 시간을 기록하여 MSReg(6ms)의 절반 수준으로 매우 높은 효율성을 입증했다.
• 제거 실험에서 크로스모달 어텐션 블록을 제거할 경우 성능 저하가 심각하게 발생했으며(p < 0.001), 이는 블록의 핵심적 역할을 뒷받침했다.
• 세그멘테이션 마스크를 입력으로 사용할 경우 정확도가 추가로 향상되어 TRE가 3.60±2.01 mm로 감소하였고, 어 attention 블록이 존재할 경우 레이블 기반 입력이 원본 MRI 입력보다 우월한 성능을 보였다.
• Grad-CAM 시각화 결과 어텐션 메커니즘이 해부학적으로 관련 있는 영역—예를 들어 전립선 경계—에 집중하고 있음을 확인하여 의미 있는 특징 학습과 공간 정렬을 입증했다.