Physics-Driven Autoregressive State Space Models for Medical Image Reconstruction

2024-12-12

arXiv.org

AI Summary

Medical image undersampling reconstruction is a canonical ill-posed inverse problem; conventional CNNs struggle to capture non-local contextual dependencies, while Transformers face prohibitive computational costs and architectural compromises. This paper proposes MambaRoll, a physics-driven autoregressive state-space model. It integrates data consistency and multi-scale modeling within an unrolled architecture, introducing the first Physics-driven State-Space Module (PSSM) and a cross-scale autoregressive prediction mechanism—jointly enabling local fidelity preservation and non-local context awareness without computational redundancy. Crucially, MRI/CT forward operators are explicitly embedded into the model, overcoming the scale-sensitivity limitations inherent in both CNNs and Transformers. Evaluated on accelerated MRI and sparse-view CT reconstruction, MambaRoll achieves significantly higher PSNR and SSIM than state-of-the-art CNNs, Transformers, and conventional physics-driven SSM-based methods, with superior artifact suppression robustness.

Technology Category

• Computer Vision: Multi-modal Vision
• Machine Learning: Multimodal Learning
• Intelligent Robots: State Estimation

Application Category

• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• Responsible Web: Machine-in-the-loop, human agency and autonomy
• User Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendation

Abstract

Medical image reconstruction from undersampled acquisitions is an ill-posed problem that involves inversion of the imaging operator linking measurement and image domains. In recent years, physics-driven (PD) models have gained prominence in learning-based reconstruction given their enhanced balance between efficiency and performance. For reconstruction, PD models cascade data-consistency modules that enforce fidelity to acquired data based on the imaging operator, with network modules that process feature maps to alleviate image artifacts due to undersampling. Success in artifact suppression inevitably depends on the ability of the network modules to tease apart artifacts from underlying tissue structures, both of which can manifest contextual relations over broad spatial scales. Convolutional modules that excel at capturing local correlations are relatively insensitive to non-local context. While transformers promise elevated sensitivity to non-local context, practical implementations often suffer from a suboptimal trade-off between local and non-local sensitivity due to intrinsic model complexity. Here, we introduce a novel physics-driven autoregressive state space model (MambaRoll) for enhanced fidelity in medical image reconstruction. In each cascade of an unrolled architecture, MambaRoll employs an autoregressive framework based on physics-driven state space modules (PSSM), where PSSMs efficiently aggregate contextual features at a given spatial scale while maintaining fidelity to acquired data, and autoregressive prediction of next-scale feature maps from earlier spatial scales enhance capture of multi-scale contextual features. Demonstrations on accelerated MRI and sparse-view CT reconstructions indicate that MambaRoll outperforms state-of-the-art PD methods based on convolutional, transformer and conventional SSM modules.

Problem

• Improving medical image reconstruction from undersampled data
• Enhancing artifact suppression and data fidelity jointly
• Capturing multi-scale context efficiently without high complexity

Innovation

• Physics-driven state-space module for contextual aggregation
• Autoregressive multi-scale feature prediction
• Deep multi-scale decoding loss enhancement