DH-Mamba: Exploring Dual-domain Hierarchical State Space Models for MRI Reconstruction

TL;DR

This work tackles the ill-posed problem of accelerated MRI reconstruction by extending selective state space models (Mamba) to dual domains. It introduces DH-Mamba, featuring a k-space specific circular scanning strategy, a dual-domain hierarchical processing scheme (Hi-scan/HiRe) to balance global context with efficiency, and a Local Enhancement Module to boost local diversity. Across CC359, fastMRI, and SKM-TEA, DH-Mamba achieves significant PSNR and SSIM gains over CNN-, ViT-, and other Mamba-based methods while reducing computational cost. These results highlight the importance of dual-domain modeling and hierarchical scanning for fast, high-fidelity MRI reconstruction, offering a scalable framework for clinical deployment and future Mamba-based MRI research.

Abstract

The accelerated MRI reconstruction poses a challenging ill-posed inverse problem due to the significant undersampling in k-space. Deep neural networks, such as CNNs and ViTs, have shown substantial performance improvements for this task while encountering the dilemma between global receptive fields and efficient computation. To this end, this paper explores selective state space models (Mamba), a new paradigm for long-range dependency modeling with linear complexity, for efficient and effective MRI reconstruction. However, directly applying Mamba to MRI reconstruction faces three significant issues: (1) Mamba typically flattens 2D images into distinct 1D sequences along rows and columns, disrupting k-space's unique spectrum and leaving its potential in k-space learning unexplored. (2) Existing approaches adopt multi-directional lengthy scanning to unfold images at the pixel level, leading to long-range forgetting and high computational burden. (3) Mamba struggles with spatially-varying contents, resulting in limited diversity of local representations. To address these, we propose a dual-domain hierarchical Mamba for MRI reconstruction from the following perspectives: (1) We pioneer vision Mamba in k-space learning. A circular scanning is customized for spectrum unfolding, benefiting the global modeling of k-space. (2) We propose a hierarchical Mamba with an efficient scanning strategy in both image and k-space domains. It mitigates long-range forgetting and achieves a better trade-off between efficiency and performance. (3) We develop a local diversity enhancement module to improve the spatially-varying representation of Mamba. Extensive experiments are conducted on three public datasets for MRI reconstruction under various undersampling patterns. Comprehensive results demonstrate that our method significantly outperforms state-of-the-art methods with lower computational cost.

Figures (9)

• Figure 1: Our main idea. (a) Applying vanilla Mamba into MRI reconstruction faces three challenges: the disruption of k-space spectrum, the long-range forgetting, and the lack of local diversity. (b) To address these, we pioneer a dual-domain hierarchical Mamba for MRI reconstruction. It tackles above issues through circular scanning in k-space, hierarchical Mamba operation in both image and k-space domains, and local diversity enhancement.
• Figure 2: (a) The overall architecture of the proposed DH-Mamba, which can be divided into three stages, i.e., shallow extraction, deep extraction, and high quality reconstruction. Given the low-quality MRI image $I_{in}$ as input, we first obtain the shallow features $F_{s}$. Then, we send $F_{s}$ into the network backbone, i.e., stacked DHM Groups, to extract deep features $F_{d}$. The core design of DHM Group is stacked DHM Blocks, which consists of DHM and LEM. Finally, the fused $F_{s}$ and $F_{d}$ is send into reconstruction head to obtain the high-quality output $I_{out}$. (b) The proposed DHM consists of two branches, i.e., k-space branch and image space branch, processing in the k-space and image domains, respectively . (c) The motivation of our k-space scanning. General scanning disrupts the k-space's structure, then we customize the circular scanning scheme to rearrange frequencies.
• Figure 3: Architecture of the proposed local enhancement module (LEM). Given the input features, after expending the channels, they are split into two parts along the channel dimension and processed by two branches, respectively. Then, the resulting outputs are combined via Hadamard product and convolution to obtain the final output.
• Figure 4: Visualization comparison on the single-coil datasets, including (a) the CC359 dataset and (b) the fastMRI dataset. The first row of each subplot shows the magnified results of the corresponding red boxes, while the second row shows the error maps. The yellow ellipses highlight the details in the results.
• Figure 5: Visualization comparison on the multi-coil SKM-TEA dataset.