💻 computer science

SO-Mamba: State-Ownership Mamba for Unrolled MRI Reconstruction

The paper proposes SO-Mamba, a novel state-ownership Mamba regularizer for unrolled MRI reconstruction that enhances anatomical coherence and detail recovery by distinctively routing persistent reconstruction evidence to a recurrent carrier while adapting state interfaces with stage-specific updates, outperforming existing CNN, Transformer, and Mamba baselines across diverse benchmarks.

Original authors: Pengcheng Fang, Hongli Chen, Fangfang Tang, Feng Liu, Xiaohao Cai, Shanshan Shan

Published 2026-05-22

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Pengcheng Fang, Hongli Chen, Fangfang Tang, Feng Liu, Xiaohao Cai, Shanshan Shan

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: Fixing Blurry MRI Scans

Imagine you are trying to take a photo of a busy city, but you only have time to take a few snapshots of the streets. When you try to put the full picture together later, it looks blurry, and there are weird ghostly lines (artifacts) where the missing pieces should be.

In medicine, MRI machines face this same problem. To save time and make patients more comfortable, they often scan only a fraction of the data needed. The computer then has to "guess" the missing parts to create a clear image.

The paper introduces a new AI tool called SO-Mamba that acts like a super-smart editor. It doesn't just guess; it knows exactly which parts of the image are the "real story" (the anatomy) and which parts are just "noise" or "updates" needed to fix the picture.

The Problem: Mixing Up the Roles

To fix these blurry images, AI models usually work in steps (like layers of a cake). In previous models using a technology called Mamba (which is great at looking at long distances in an image), everything was thrown into one big "recycling bin" of information.

Think of it like a construction site:

The Resident Carrier: This is the solid foundation and the main structure of the building (the stable anatomy). It needs to stay put and be consistent.
The Non-Resident Evidence: This is the daily work order, the new blueprints, or the specific corrections needed for this step (like fixing a specific aliasing artifact).

The Mistake: Previous models treated the "solid foundation" and the "daily work order" as the same thing. They put them in the same bucket. This caused confusion. The AI got mixed up about what was permanent structure and what was just a temporary update, leading to blurry or distorted images.

The Solution: SO-Mamba (State-Ownership)

The authors created SO-Mamba, which introduces a strict rule called "State Ownership."

Imagine a high-security office building with three distinct zones:

The Vault (Recurrent Residency): Only the "Resident Carrier" (the stable, coherent anatomy) is allowed inside here. This is the core memory that the AI carries forward. It must remain clean and stable.
The Control Panel (State-Interface Access): The "Non-Resident Evidence" (the updates and corrections) is not allowed in the Vault. Instead, it goes to the Control Panel. Here, it can tweak the settings (like turning a dial or adjusting a switch) to tell the Vault how to read or write information, but it doesn't become the information itself.
The Exit Door (Non-State Output Correction): The "Non-Resident Evidence" also has a direct line to the final output. It can make immediate corrections to the image before it leaves the building, without ever entering the Vault.

The "State-Ownership Router" (SOR):
This is the security guard at the door. When information comes in, the SOR checks it:

"Are you stable and part of the main structure?" -> Yes? You go to the Vault.
"Are you a temporary update or a specific correction?" -> Yes? You go to the Control Panel and the Exit Door, but you are banned from the Vault.

By keeping these roles separate, the AI doesn't get confused. It preserves the clear, stable anatomy while still being able to apply sharp, specific fixes to the blurry parts.

How They Tested It

The researchers didn't just guess this would work; they built a "leak detector."

Imagine the "Vault" is supposed to be quiet and stable. If you hear loud, chaotic noise coming from inside the Vault, it means the "updates" are leaking in where they shouldn't be.
They measured how much "noise" (high-frequency details) was leaking into the main memory versus how much was staying in the control panel.
The Result: In SO-Mamba, the noise stayed in the control panel and exit door. In older models, the noise leaked into the Vault, making the final image messy.

The Results

They tested this on five different types of MRI scans (knees, brains, hearts, prostates) with different scanning patterns.

The Winner: SO-Mamba beat all the other top AI models (including CNNs, Transformers, and other Mamba versions).
The Proof: The images were sharper, the edges were clearer, and the "ghost" artifacts were gone.
Efficiency: It did all this while using less computer power than many of its competitors.

Summary

SO-Mamba is like a smart traffic controller for MRI reconstruction. It realized that the "permanent structure" of an image and the "temporary fixes" needed to be kept in separate lanes. By building a system that strictly enforces this separation, it creates clearer, more accurate medical images faster than before.

Technical Summary: SO-Mamba for Unrolled MRI Reconstruction

Problem Statement
Accelerated Magnetic Resonance Imaging (MRI) reconstruction aims to recover high-quality images from undersampled k-space measurements while preserving anatomical coherence across large spatial regions. While deep learning has advanced this field through data-driven priors, current architectures face specific limitations in unrolled, data-consistency (DC)-coupled solvers. Convolutional Neural Networks (CNNs) struggle with long-range anatomical dependencies due to limited receptive fields, while Vision Transformers suffer from quadratic computational complexity that hinders scalability to high-resolution clinical inputs.

Selective State-Space Models (SSMs), such as Mamba, offer linear-complexity long-range modeling, making them attractive regularizers for unrolled reconstruction. However, a critical design gap exists in how reconstruction evidence is assigned within the state-space computation of a DC-coupled unrolled solver. In such solvers, each stage operates on a different reconstruction iterate. Treating all evidence uniformly—placing both persistent "resident" content and stage-dependent "non-resident" updates into the same recurrent content route—can entangle roles that require distinct handling. Specifically, it is unclear which evidence should become the recurrent content (the state trajectory) versus which evidence should merely modulate how that content is written/read or correct the output without entering the recurrent loop.

Methodology: SO-Mamba
The authors propose SO-Mamba (State-Ownership Mamba), a regularizer that formalizes this allocation problem as state ownership. The core principle separates the functional roles of reconstruction evidence within each unrolled stage into three distinct routes:

Recurrent Residency: Evidence that becomes the direct content of the selective-scan state-space computation.
State-Interface Access: Evidence that controls how the state is written and read (modulating the B and C interfaces).
Non-State Output Correction: Evidence that contributes to the final output without entering the recurrent content path.

Key Components:

State-Ownership Router (SOR): This module takes the feature representation of the current reconstruction iterate and splits it into two streams:
- Resident Carrier ( $L_k$ ): Generated via a fixed, non-adaptive binomial carrier projection. This stream is assigned recurrent residency, supplying the tokens for the Mamba content route. It represents stable, spatially coherent anatomical layouts.
- Non-Resident Evidence ( $G_k$ ): Composed of the remaining evidence (local deviations from the carrier) and specific non-resident features. This stream is excluded from the recurrent content source ( $G_k \nrightarrow u_k$ $G_{k} ↛ u_{k}$ ). Instead, it is routed to:
  - State-Interface Modulation: $G_k$ predicts affine parameters to modulate the B and C state interfaces, adapting how the resident content is processed.
  - Non-State Output Outlet: $G_k$ passes through a Non-State Refinement (NSR) block and is concatenated with the Mamba readout to correct the final output.
Ownership-Aware Mamba-3 Update: The Mamba-3 block processes the resident carrier as the input tokens while using the non-resident stream to condition the state-space transitions. This ensures the recurrent state trajectory remains dominated by stable anatomical content, while dynamic updates influence the system through controlled interfaces and output corrections.
Two-Level Outer-Band Leakage Diagnostic: To validate the separation of roles, the authors introduce a diagnostic metric that measures "outer-band" (high-frequency) energy in two places: the raw hidden-state trajectory ( $H_{Leak}$ ) and the post-scan readout ( $R_{Leak}$ ). This separates hidden-state storage from readout expression, verifying that non-resident evidence does not corrupt the recurrent state trajectory but is effectively expressed at the output level.

Key Contributions

Problem Formulation: The paper identifies "stage-local state allocation" as a distinct design problem in Mamba-based unrolled MRI reconstruction and formulates it as state ownership.
SO-Mamba Architecture: A novel regularizer built on Mamba-3 that explicitly assigns evidence to recurrent residency, state-interface access, and non-state output correction within each stage.
SOR Mechanism: The introduction of the State-Ownership Router, which constructs a resident carrier and routes non-resident evidence to interface modulation and output correction, enforcing the ownership contract.
Comprehensive Evaluation: Evaluation on five public benchmarks (fastMRI, CC359, ACDC, M4Raw, Prostate158) covering diverse anatomies and sampling patterns, alongside a novel diagnostic analysis of state behavior.

Experimental Results

Performance: SO-Mamba consistently outperforms CNN-, Transformer-, and Mamba-based baselines (including UNet, TransUNet, ReconFormer, LMO, and HiFi-Mamba) across five datasets.
- On fastMRI (AF=4), it achieves 34.65 dB PSNR and 0.855 SSIM.
- On CC359 (AF=4), it achieves 36.69 dB PSNR and 0.954 SSIM.
- It also leads in radial (ACDC) and multi-coil (M4Raw) settings.
Efficiency: SO-Mamba achieves the best accuracy among efficient models while utilizing the lowest FLOPs (45.59G) compared to HiFi-Mamba (67.87G) and LMO (484.98G) on fastMRI at 8× acceleration.
Ablation Studies:
- Removing the state-interface access or output outlet degrades performance.
- Allowing non-resident evidence to enter the recurrent content route (violating the ownership contract) causes a sharp performance drop (from 34.65 dB to 31.71 dB), confirming the necessity of the separation.
- The leakage diagnostic shows that removing the output outlet forces outer-band content into the recurrent readout (increasing $R_{Leak}$ by ~5×), whereas the full SO-Mamba design keeps the recurrent state clean while expressing necessary corrections at the output.

Significance and Claims
The paper claims that SO-Mamba addresses a fundamental architectural constraint in unrolled MRI reconstruction: the need to distinguish between evidence that should form the persistent state trajectory and evidence that should only influence the system's dynamics or output. By enforcing state ownership, the method prevents the entanglement of stable anatomical structures with stage-dependent updates, leading to sharper boundaries and more uniform error maps.

The authors emphasize that the State-Ownership Router (SOR) is an architectural assignment rule rather than a semantic decomposition of anatomy versus artifacts. The significance lies in demonstrating that linear-complexity SSMs can be effectively adapted for unrolled reconstruction by carefully managing the flow of information, achieving state-of-the-art performance with high computational efficiency. The work suggests that future improvements may lie in adaptive ownership routing and broader clinical acquisition settings, but the current contribution is the establishment of the state-ownership principle itself.