[논문 리뷰] MambaMIR: An Arbitrary-Masked Mamba for Joint Medical Image Reconstruction and Uncertainty Estimation

Jiahao Huang, Liutao Yang외 6명|arXiv (Cornell University)|2024. 02. 28.

Medical Imaging and Analysis인용 수 11

한 줄 요약

MambaMIR은 빠른 MRI와 sparse-view CT 작업에서 Monte Carlo 기반 불확실성 맵을 가능하게 하는 임의 마스크(arbitrary-mask) 메커니즘을 사용하여 의학 영상 재구성 및 불확실성 추정을 함께 수행하는 Mamba 기반 모델이며, GAN 변형(MambaMIR-GAN)을 포함한다.

ABSTRACT

The recent Mamba model has shown remarkable adaptability for visual representation learning, including in medical imaging tasks. This study introduces MambaMIR, a Mamba-based model for medical image reconstruction, as well as its Generative Adversarial Network-based variant, MambaMIR-GAN. Our proposed MambaMIR inherits several advantages, such as linear complexity, global receptive fields, and dynamic weights, from the original Mamba model. The innovated arbitrary-mask mechanism effectively adapt Mamba to our image reconstruction task, providing randomness for subsequent Monte Carlo-based uncertainty estimation. Experiments conducted on various medical image reconstruction tasks, including fast MRI and SVCT, which cover anatomical regions such as the knee, chest, and abdomen, have demonstrated that MambaMIR and MambaMIR-GAN achieve comparable or superior reconstruction results relative to state-of-the-art methods. Additionally, the estimated uncertainty maps offer further insights into the reliability of the reconstruction quality. The code is publicly available at https://github.com/ayanglab/MambaMIR.

연구 동기 및 목표

• 의학 영상 재구성 및 불확실성 추정을 위한 Mamba를 동기 부여하고 확장한다.
• 전역 수용 영역(global receptive fields)을 갖춘 강인한 재구성을 위한 Arbitrary-Masked State Space (AMSS) 블록과 AMS6를 도입한다.
• 학습 및 추론 중 스캔의 임의 마스킹을 통해 MC 기반 불확실성 추정을 가능하게 한다.
• 여러 해부 영역에 걸친 빠른 MRI 및 sparse-view CT에서 경쟁력 있는 재구성 품질을 보여준다.
• 재구성의 지각 품질을 향상시키기 위한 GAN 변형(MambaMIR-GAN)을 제공한다.

제안 방법

• 고해상도 의학 영상을 위해 선형 복잡도와 전역 수용 영역을 갖춘 Mamba를 채택한다.
• Scan Expanding, Arbitrary-Masked, S6, Scan Merging 모듈로 구성된 Arbitrary-Masked S6 (AMS6) 블록을 도입한다.
• 두 개의 AMSS 블록으로 이루어진 AMSS Block Group과 잔차 연결을 사용한다.
• 이미지 도메인 Charbonnier 손실, 변환 도메인 Charbonnier 손실, VGG 잠재 공간의 지각 손실로 구성된 복합 손실로 학습하며; GAN 모드에서는 U-Net 판별자를 이용한 적대적 손실을 추가한다.
• Generator로서 U-Net 판별자를 갖춘 MambaMIR-GAN과 적대적 학습 목적을 제공한다.
• FastMRI 무릎 및 Low-Dose CT SVCT 데이터셋에서 최첨단 방법과 비교 평가한다.

실험 결과

연구 질문

• RQ1 MambaMIR이 fast MRI 및 SVCT 작업에서 최첨단 방법과 동등하거나 더 우수한 재구성 품질을 달성할 수 있는가?
• RQ2 임의 마스크 메커니즘이 재구성된 영상에 대해 의미 있는 Monte Carlo 기반 불확실성 추정을 가능하게 하는가?
• RQ3 GAN 기반 적대적 학습(MambaMIR-GAN)을 도입하면 재구성의 지각 품질이 충실도 손실 없이 향상되는가?
• RQ4 무릎, 흉부, 복부 등 서로 다른 해부 영역과 가속/뷰 샘플링 설정에서 MambaMIR 및 MambaMIR-GAN의 성능은 어떠한가?

주요 결과

• MambaMIR 및 MambaMIR-GAN은 fast MRI 및 SVCT 작업에서 최첨단 방법과 비교해 동등하거나 우수한 재구성 결과를 달성한다.
• MambaMIR-GAN은 특히 높은 가속 계수에서 MambaMIR보다 지각 점수(LPIPS)가 더 나은 경향이 있다.
• 임의 마스크 메커니즘에 의해 생성된 불확실성 맵은 재구성 신뢰도에 대한 시각적 인사이트를 제공하며, MRI의 고가속에서 더 높은 불확실성을 보인다.
• SVCT의 경우 보고된 부분집합에서 SSIM 및 PSNR 측면에서 MambaMIR과 MambaMIR-GAN이 모든 다른 방법을 능가하며, MambaMIR-GAN은 LPIPS에서도 경쟁력 있는 성능을 보인다.
• GAN 통합은 지각 품질을 향상시키고 불확실성 맵에 영역별 변화를 유발할 수 있어 적대적 학습 하에서도 강건한 재구성을 시사한다.