[논문 리뷰] Efficient Folded Attention for 3D Medical Image Reconstruction and Segmentation

Hang Zhang, Jinwei Zhang외 5명|arXiv (Cornell University)|2020. 09. 13.

Medical Image Segmentation Techniques참고 문헌 34인용 수 8

한 줄 요약

이 논문은 3차원 의료 영상 복원 및 분할을 위한 새로운 접힘 어텐션(FA) 기반 메커니즘을 제안하며, 표준 어텐션 대비 계산 비용과 GPU 메모리 사용량을 크게 줄이면서도 정확도를 유지하거나 향상시킨다. 전체 유사도 행렬을 근사하기 위해 네 가지 순열을 사용한 텐서 접기 및 펼치기 연산을 적용함으로써 FA는 공간적 및 채널 의존성을 효율적으로 공동 모델링할 수 있게 하며, 양자화된 순환자기율도 및 다발성 경화증 병변 분할 작업에서 최신 기술 수준의 성능을 달성한다. 표준 어텐션 대비 최대 97.9%의 메모리 절감과 88.9%의 FLOPs 절감을 기록한다.

ABSTRACT

Recently, 3D medical image reconstruction (MIR) and segmentation (MIS) based on deep neural networks have been developed with promising results, and attention mechanism has been further designed to capture global contextual information for performance enhancement. However, the large size of 3D volume images poses a great computational challenge to traditional attention methods. In this paper, we propose a folded attention (FA) approach to improve the computational efficiency of traditional attention methods on 3D medical images. The main idea is that we apply tensor folding and unfolding operations with four permutations to build four small sub-affinity matrices to approximate the original affinity matrix. Through four consecutive sub-attention modules of FA, each element in the feature tensor can aggregate spatial-channel information from all other elements. Compared to traditional attention methods, with moderate improvement of accuracy, FA can substantially reduce the computational complexity and GPU memory consumption. We demonstrate the superiority of our method on two challenging tasks for 3D MIR and MIS, which are quantitative susceptibility mapping and multiple sclerosis lesion segmentation.

연구 동기 및 목표

• 3차원 의료 영상 처리에서 전통적인 자기어텐션 메커니즘의 높은 계산 비용과 메모리 사용량 문제를 해결하기 위해.
• 계산 복잡도를 증가시키지 않으면서도 공간적 및 채널 방향 의존성을 모두 모델링할 수 있는 효율적인 어텐션 메커니즘을 개발하기 위해.
• 3D U-Net 기반 네트워크에서 어텐션 모듈의 메모리 및 FLOPs 오버헤드를 줄이기 위해.
• QSM 복원 및 MS 병변 분할과 같은 대규모 3차원 의료 영상 작업에서 어텐션 메커니즘의 실용적 구현을 가능하게 하기 위해.

제안 방법

• 3D 특징 텐서를 더 작은 부분 텐서로 변환하기 위해 텐서 접기 및 펼침 연산을 도입함으로써, 하나의 큰 유사도 행렬 대신 네 개의 하위유사도 행렬을 계산할 수 있도록 한다.
• 특징 텐서의 네 가지 순열을 사용하여 전체 전역 어텐션 행렬을 근사하는 하위유사도 행렬을 생성함으로써 복잡도를 감소시킨다.
• 각 하위유사도 행렬은 쿼리-키 내적 연산을 통해 계산된 후 소프트맥스를 거치며, 결과는 학습 가능한 투영을 통해 집계된다.
• 각 특징 원소를 모든 채널을 통해 벡터로 간주함으로써 FA 모듈은 공간적 및 채널 의존성을 통합적으로 모델링할 수 있다.
• 유사도 행렬에 랭크-1 제약 조건을 적용하여 일반화 성능을 향상시키고 과적합을 줄인다.
• FA 모듈은 플러그 앤 플레이 방식으로 구현되어 기존의 3D CNN 아키텍처(예: U-Net)에 최소한의 아키텍처 변경으로 통합 가능하다.

실험 결과

연구 질문

• RQ1 3차원 의료 영상에서 공간적 및 채널 의존성을 모두 모델링할 수 있는 계산적으로 효율적인 어텐션 메커니즘을 설계할 수 있는가?
• RQ2 텐서 접기 및 펼침 기법이 3차원 의료 영상 작업에서 자기어텐션의 FLOPs 및 GPU 메모리 사용량을 얼마나 줄일 수 있는가?
• RQ3 제안된 접힘 어텐션 기반 메커니즘이 도전적인 3차원 의료 영상 복원 및 분할 벤치마크에서 성능을 유지하거나 향상시키는가?
• RQ4 특히 임상 상황에서 진단 기준 데이터가 없을 경우 접힘 어텐션 모듈이 어떻게 일반화되는가?

주요 결과

• QSM에서 FA-Net은 RMSE 31.18로 가장 낮은 수치를 기록하여 QSMnet, DA-Net, RSA-Net을 모두 능가했다.
• QSM에서 FA-Net은 PSNR 49.06을 기록하여 복원 정밀도가 뛰어나고 일반 오차가 감소함을 나타냈다.
• QSM에서 FA-Net은 HFEN 32.49로 가장 낮게 기록되어 고주파 성분과 구조 유사도를 더 잘 유지함을 보였다.
• QSM에서 FA-Net은 SSIM 0.9833으로 가장 높은 수치를 기록하여 기준값과 비교해 뛰어난 구조적 유사성과 대trast 유사성을 입증했다.
• 표준 자기어텐션(SA) 대비 FA 모듈은 GPU 메모리 사용량을 97.9% 감소시키고 FLOPs를 88.9% 절감하면서도 유사하거나 더 높은 정확도를 확보했다.
• 정성적 결과 분석에서 FA-Net은 기준값이 없는 COSMOS 데이터가 없는 MS 환자 스캔에서도 더 선명한 병변 경계와 더 많은 고강도 병변을 생성함을 확인했다.