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

技术摘要：用于展开式 MRI 重建的 SO-Mamba

问题陈述
加速磁共振成像（MRI）重建旨在从欠采样的 k 空间测量中恢复高质量图像，同时保持大空间范围内的解剖结构连贯性。尽管深度学习通过数据驱动的先验推动了该领域的发展，但当前的架构在展开式、数据一致性（DC）耦合求解器中面临特定的局限性。卷积神经网络（CNN）由于感受野有限，难以处理长程解剖依赖；而视觉 Transformer 则因二次计算复杂度而阻碍了其向高分辨率临床输入的扩展。

选择性状态空间模型（SSM），如 Mamba，提供了线性复杂度的长程建模能力，使其成为展开式重建中极具吸引力的正则化器。然而，在 DC 耦合的展开式求解器的状态空间计算中，如何分配重建证据存在一个关键的设计缺口。在此类求解器中，每个阶段操作于不同的重建迭代结果。若将所有证据同等对待——将持久的“驻留”内容与阶段相关的“非驻留”更新置于相同的循环内容路径中——会导致需要不同处理的角色的纠缠。具体而言，尚不明确哪些证据应成为循环内容（即状态轨迹），而哪些证据应仅调节该内容的写入/读取方式，或在未进入循环回路的情况下修正输出。

方法：SO-Mamba
作者提出了SO-Mamba（状态所有权 Mamba），这是一种将上述分配问题形式化为状态所有权的正则化器。其核心原则是将每个展开阶段内重建证据的功能角色分离为三条不同的路径：

1. 循环驻留：成为选择性扫描状态空间计算直接内容的证据。
2. 状态接口访问：控制状态写入和读取方式的证据（调节 B 和 C 接口）。
3. 非状态输出修正：在不进入循环内容路径的情况下对最终输出做出贡献的证据。

关键组件：

• 状态所有权路由器（SOR）：该模块接收当前重建迭代结果的特征表示，并将其拆分为两个流：
  • 驻留载体（$L_k$）：通过固定的、非自适应的二项式载体投影生成。该流被分配循环驻留权，为 Mamba 内容路径提供令牌。它代表了稳定且空间连贯的解剖布局。
  • 非驻留证据（$G_k$）：由剩余证据（载体的局部偏差）和特定的非驻留特征组成。该流被排除在循环内容源之外（$G_k \nrightarrow u_k$）。相反，它被路由至：
    • 状态接口调制：$G_k$ 预测仿射参数以调制 B 和 C 状态接口，从而调整驻留内容的处理方式。
    • 非状态输出出口：$G_k$ 通过非状态细化（NSR）模块，并与 Mamba 读出结果拼接，以修正最终输出。
• 所有权感知 Mamba-3 更新：Mamba-3 块将驻留载体作为输入令牌进行处理，同时利用非驻留流来调节状态空间转换。这确保了循环状态轨迹主要由稳定的解剖内容主导，而动态更新则通过受控接口和输出修正影响系统。
• 两级外带泄漏诊断：为了验证角色分离的有效性，作者引入了一种诊断指标，用于测量“外带”（高频）能量在两个位置的表现：原始隐藏状态轨迹（$H_{Leak}$）和扫描后读出结果（$R_{Leak}$）。这将隐藏状态存储与读出表达分离开来，验证了非驻留证据不会破坏循环状态轨迹，而是在输出层面得到有效表达。

主要贡献

1. 问题形式化：本文识别出“阶段局部状态分配”是基于 Mamba 的展开式 MRI 重建中的一个独特设计问题，并将其形式化为状态所有权。
2. SO-Mamba 架构：一种基于 Mamba-3 的新型正则化器，在每个阶段内明确将证据分配给循环驻留、状态接口访问和非状态输出修正。
3. SOR 机制：引入了状态所有权路由器，该路由器构建驻留载体并将非驻留证据路由至接口调制和输出修正，从而强制执行所有权契约。
4. 综合评估：在五个涵盖多种解剖结构和采样模式的公共基准（fastMRI、CC359、ACDC、M4Raw、Prostate158）上进行了评估，并辅以新颖的状态行为诊断分析。

实验结果

• 性能：SO-Mamba 在五个数据集上始终优于基于 CNN、Transformer 和 Mamba 的基线模型（包括 UNet、TransUNet、ReconFormer、LMO 和 HiFi-Mamba）。
  • 在 fastMRI（加速因子 AF=4）上，其实现了 34.65 dB PSNR 和 0.855 SSIM。
  • 在 CC359（AF=4）上，其实现了 36.69 dB PSNR 和 0.954 SSIM。
  • 在径向（ACDC）和多线圈（M4Raw）设置中也处于领先地位。
• 效率：SO-Mamba 在高效模型中实现了最佳精度，同时在 fastMRI 的 8 倍加速下，利用了最低的 FLOPs（45.59G），低于 HiFi-Mamba（67.87G）和 LMO（484.98G）。
• 消融研究：
  • 移除状态接口访问或输出出口会导致性能下降。
  • 允许非驻留证据进入循环内容路径（违反所有权契约）会导致性能急剧下降（从 34.65 dB 降至 31.71 dB），证实了分离的必要性。
  • 泄漏诊断显示，移除输出出口会迫使外带内容进入循环读出（使 $R_{Leak}$ 增加约 5 倍），而完整的 SO-Mamba 设计则保持循环状态清洁，同时在输出层面表达必要的修正。

意义与主张
本文主张，SO-Mamba 解决了展开式 MRI 重建中的一个基本架构约束：即区分应形成持久状态轨迹的证据与仅应影响系统动态或输出的证据。通过强制执行状态所有权，该方法防止了稳定解剖结构与阶段相关更新的纠缠，从而产生了更清晰的边界和更均匀的误差图。

作者强调，状态所有权路由器（SOR）是一种架构分配规则，而非对解剖结构与伪影的语义分解。其意义在于证明，通过仔细管理信息流，线性复杂度的 SSM 可以被有效地适应于展开式重建，从而以高计算效率实现最先进的性能。这项工作表明，未来的改进可能在于自适应所有权路由和更广泛的临床采集设置，但当前的贡献在于确立了状态所有权原则本身。