[논문 리뷰] Physics-Driven Autoregressive State Space Models for Medical Image Reconstruction

Bilal Kabaş, Fuat Arslan외 4명|arXiv (Cornell University)|2024. 12. 12.

Neural Networks and Applications인용 수 9

한 줄 요약

이 논문은 물리 기반의 autoregressive 상태공간 모델인 MambaRoll을 도입하여 undersampled 데이터로부터 다중 스케일 컨텍스트 특징을 점진적으로 모으고 autoregressive 다음 스케일 예측으로 재구성함으로써 MRI와 sparse-view CT에서 컨볼루션, 트랜스포머, 그리고 기존 SSM PD 방법을 능가한다.

ABSTRACT

Medical image reconstruction from undersampled acquisitions is an ill-posed inverse problem requiring accurate recovery of anatomical structures from incomplete measurements. Physics-driven (PD) network models have gained prominence for this task by integrating data-consistency mechanisms with learned priors, enabling improved performance over purely data-driven approaches. However, reconstruction quality still hinges on the network's ability to disentangle artifacts from true anatomical signals-both of which exhibit complex, multi-scale contextual structure. Convolutional neural networks (CNNs) capture local correlations but often struggle with non-local dependencies. While transformers aim to alleviate this limitation, practical implementations involve design compromises to reduce computational cost by balancing local and non-local sensitivity, occasionally resulting in performance comparable to CNNs. To address these challenges, we propose MambaRoll, a novel physics-driven autoregressive state space model (SSM) for high-fidelity and efficient image reconstruction. MambaRoll employs an unrolled architecture where each cascade autoregressively predicts finer-scale feature maps conditioned on coarser-scale representations, enabling consistent multi-scale context propagation. Each stage is built on a hierarchy of scale-specific PD-SSM modules that capture spatial dependencies while enforcing data consistency through residual correction. To further improve scale-aware learning, we introduce a Deep Multi-Scale Decoding (DMSD) loss, which provides supervision at intermediate spatial scales in alignment with the autoregressive design. Demonstrations on accelerated MRI and sparse-view CT reconstructions show that MambaRoll consistently outperforms state-of-the-art CNN-, transformer-, and SSM-based methods.

연구 동기 및 목표

• Motivate improved reconstruction from undersampled measurements by leveraging physics-driven priors.
• Develop a multi-scale autoregressive framework that fuses context across spatial scales while enforcing data fidelity.
• Introduce novel physics-driven state-space modules (PSSMs) that operate across scales within an unrolled architecture.
• Demonstrate superior reconstruction quality on accelerated MRI and sparse-view CT datasets compared to existing PD methods.

제안 방법

• Propose MambaRoll, an unrolled PD architecture with K cascades that progressively reconstructs high-resolution feature maps across S spatial scales.
• Within each cascade, employ PSSM modules that include an encoder, a shuffled SSM, a decoder, and a residual data-consistency block to enforce fidelity to the imaging operator.
• Use autoregressive prediction across scales by feeding concatenated features from earlier scales to process subsequent scales.
• Train with a multi-term objective that includes the cascade output error and scale-specific decoded feature errors to encourage faithful scale-wise reconstructions.
• Evaluate on accelerated MRI (IXI, fastMRI) and sparse-view CT (LoDoPaB-CT) against PD-CNN, PD-TransUNet, PD-UNetMHA, PD-UNetMamba, and PD-UMamba.

실험 결과

연구 질문

• RQ1 Can a physics-driven autoregressive SSM framework improve reconstruction fidelity for undersampled medical imaging tasks?
• RQ2 Do multi-scale PSSMs with autoregressive scale predictions better capture long-range contextual information while maintaining data consistency?
• RQ3 How does MambaRoll compare to convolutional, transformer, and conventional SSM PD approaches in MRI and CT reconstructions?
• RQ4 What is the contribution of PSSM, autoregression, and data-consistency blocks to overall performance?

주요 결과

• MambaRoll consistently outperforms competing methods across MRI and CT tasks at R=4 and R=8 (e.g., MRI IXI/fastMRI; CT LoDoPaB-CT).
• Across evaluated tasks, MambaRoll yields substantial PSNR/SSIM gains over PD-CNN, PD-TransUNet, PD-UNetMHA, PD-UNetMamba, and PD-UMamba.
• Ablation studies show that removing PSSM, autoregression, or data-consistency blocks degrades reconstruction performance, highlighting the importance of each component.
• In MRI, MambaRoll achieves higher PSNR/SSIM and better tissue delineation with reduced artifacts and noise compared to baselines.
• In sparse-view CT, MambaRoll achieves larger PSNR/SSIM gains relative to baselines, indicating robust performance across modalities.