点击下方卡片,关注 「AI视界引擎」 公众号
( 添加时备注:方向+学校/公司+昵称/姓名 )
MRI(磁共振成像)是一种最重要的医学影像方法之一,它提供了软组织的高分辨率,虽然扫描速度有限。深度学习的出现推动了加速MRI扫描重建的创新方法的快速发展,这些方法利用卷积神经网络以及最近出现的视觉变形器。
最近提出的结构状态空间模型(如Mamba)由于其效率和与变形器模型相比的低计算要求而得到了一定的关注。
作者提出了一种创新性的MRI重建框架,该框架的核心采用了结构状态空间模型,旨在增强长程上下文敏感性和重建效率。
在公开的大脑MRI数据集上的全面实验表明,作者的模型为最先进的重建 Baseline 设定了新的里程碑。
1 Introduction
磁共振成像(MRI)是一种发病率较高的成像技术,因其卓越的软组织可视化能力而备受青睐。然而,其使用受到扫描过程长且昂贵的限制,因此迫切需要加速MRI技术以提高在临床环境中的实用性。在奈奎斯特率以下取样的k空间点(即傅里叶系数)较少,可以加快扫描速度,但需要解决一个欠定问题的重建问题 [28]。压缩感知(CS)解决方案已在重建问题中得到证明,具有较强的效力 。然而,尽管它们有效,CS方法受到高参数敏感性的限制,需要精心调整以实现最佳性能。此外,大多数CS算法的迭代性质导致其计算效率低下,与其他方法相比较慢。
自深度学习出现以来,迭代方法被深度学习方法逐步取代 。这些方法主要利用卷积神经网络(CNNs),已证明在具有数据驱动的各种图像处理任务上效果卓著。除了纯数据驱动方法,提出了多种物理引导的重建模型 [46, 1, 38, 34]。在这些模型中,包括下采样的模式和线圈灵敏度在多线圈设置下,也使用正向编码算子以增强重建的鲁棒性和优越性能 [46]。然而,CNNs受到有限感受野的限制,其能力捕捉长程依赖关系与上下文信息的能力有限 [53]。
近年来,随着 Transformer 在下游计算机视觉任务中的成功, Transformer 基础的方法已经被引入到磁共振重建中以解决自注意力 Transformer 与输入序列长度的平方复杂度的限制。通过集成 Transformer 架构,增强对长程依赖关系和扩展有效感受野的敏感性得到了实现。然而,自注意力 Transformer 的平方复杂性与输入序列长度成正比,会产生计算负担 。
结构化状态空间模型(SSM)是解决自注意力 Transformer 与平方复杂性问题的一个解决方案,可以有效地模拟长序列 [12]。最近,一种新的结构化状态空间模型(即Mamba)被提出,具有输入相关的参数化,且随着序列长度的线性缩放,其复杂性可以减小到平方 [11]。Mamba在大型语言模型中展示出前景,并被用于几个视觉任务,包括分类 [26],图像分割 [44, 31]和图像恢复 [14]。
受到这些发展的启发,在本文中,作者提出了一种以结构化状态空间模型为核心,具有非常少可训练参数的物理引导MRI重建框架。作者的模型通过在数据一致性和Mamba块之间交替,在提高long-range敏感性同时不遗漏底层物理模型的条件下实现亚平方复杂性。综合实验证明本模型在复杂多线圈和大小单线圈公共脑MRI数据集的优越性能,优于最先进的MRI重建模型。
本论文的贡献如下:
- 作者提出了一种新型的轻量级物理引导MRI重建模型,利用结构化状态空间模型作为核心块,该模型通过显著增加有效感受野提高了模型的 long-range 敏感性。
- 作者在复杂多线圈和大小单线圈公共脑MRI数据集的比较中取得了优越性能,超过了最先进的重建 Baseline 。
2 Related Works
基于数据驱动卷积神经网络(CNN)的模型.卷积神经网络(CNNs)在深度MRI重建中主要被用于学习下采样和全书采样之间的底层映射。王等人[42]提出了一种基于CNN的MRI重建框架。叶等人[49]提出了一种基于CNN的通用反问题高效解决模型。李等人[25]提出了一种深度CNN辅助的CNN-模型用于MRI重建。현等人[21]利用了一种U型网络的重建模型。达等人[6]提出了一个基于条件生成对抗网络(GAN)的模型,这种模型是在GAN在图像生成方面的成功基础上提出的。冈戈尔等人[13]提出了一种基于扩散模型的自适应到测试时的方法。
基于物理引导的卷积神经网络(CNN)模型.阿盖尔等人[1]提出了一种有卷积神经网络(CNN)背书的无展开模型。沙莱珀等人[38]利用了深度级联CNN以高效集成前向算子。额等人[9]提出了一种双域物理引导的CNN模型。亚曼等人[46]提出了一种通过无展开的CNN背书新颖的解决方案,该方案可以消除监督需求。杨等人[47]提出了一种在交替方向乘法器(ADMM)算法训练下的统一物理引导模型。比施瓦斯等人[3]提出了一种添加物理引导网络的额外平滑性正则化方法。秦等人[35]提出了一种用于动态MRI重构的基于循环卷积模型的物理引导方法。斯特拉姆等人[39]提出了一种利用卷积展开和图像域的变分模型。亚西美斯等人[51]提出了一种以图像域为优化的混合架构正在k空间进行优化。
基于 Transformer 的模型.在基于 Transformer 方法在计算机视觉领域成功之后,科克马等人[24]提出了一种零样本学习的 Transformer 模型用于无监督重建。郭等人[15]提出了一种有物理引导的 Transformer 背骨的物理引导模型。周等人[56]提出了一种使用频率和图像域双域的 Transformer 模型。黄等人[19]将计算高效的Swin Transformer模型适用于磁共振成像(MRI)重建。胡等人[18]提出了一种具有防止在保持高频细节的物理引导的 Transformer 模型。法比安等人[10]提出了一种使用 Transformer 与卷积神经网络结合的无展开混合架构。刘等人[30]提出了一种以 Transformer 为骨架的聚焦分区GAN。赵等人[55]提出了一种具有Swin Transformer背骨的MRI重建的对抗模型。
Mamba在医学影像中的应用。更近一步,一些Mamba基于的深度学习模型已在各种医学影像任务中得到利用。于等人[52]提出了一种用于医学图像分类的视觉Mamba模型。
谢等人[45]提出了一种U型的Mamba模型用于3D图像分割。杨等人[48]将Mamba应用于医学视频目标分割。
王等人[44]和阮等人[37]提出了一种用于医学图像分割的U型的Mamba-基于模型。阿特里等人[2]提出了一种用于医学图像合成的Mamba-基于的合成模型。
同时,黄等人[20]提供了一种任意屏蔽的Mamba-基于模型,用于医学图像的重建和无物理模型考虑的不确定性估计。邹等人[58]提出了一种多域的Mamba-基础重建框架。然而,这项工作与作者的方法不同,因为它使用了完全采样的辅助模态与目标重建相结合。据作者所知,这是第一个物理引导的SSM(物理引导的随机采样模型)基于的MRI重建模型。
3 Background
Accelerated MRI Reconstruction
MRI图像重建从部分k空间可以定义如下:
其中
是部分傅里叶算子,是目标MRI图像,是线圈灵敏度图,表示部分的k空间点采样。从重建会导致无穷多解,因此需要正则化处理。带有正则化的重建目标可以表示为
Structured State Space Models
结构化状态空间模型(SSMs)用于将一维序列 转换为序列 ,使用隐含潜在空间 如图所示,
其中 , 和 . 为了将连续的系统集成到现有的深度学习框架中,需要进行离散化。离散化通常通过零阶保持(ZOH)实现,其中 参数引入空间,用于重新定义离散形式下的连续参数 和 如下
采用以上方法,作者可以重新编写方程(3)和(4)如下
Mamba中提出的选择性扫描机制使得可学习参数输入依赖(动态),其中 , 和 定义使用了输入序列 的线性投影,其中 对应批次大小, 对应序列长度, 对应输入通道数, 对应内部状态维度。
4 Methodology
作者将[26]中提出的视觉状态空间模型(VSSM)模块应用于作者的物理引导模型,如图1所示。遵循VSSM,作者简单地将输入线圈合成的复杂MR图像分割成大小为p的 patches,并使用VSSM在4种不同的方向(从左上角到右下角、从右下角到左上角、从右上角到左下角和从左下角到右上角)上展开,如图1所示。在视觉数据上使用SSM需要展开不同方向,因为原始的SSM是因果的,并且在输入序列中以单一方向遍历,这适合于时间序列数据或语言处理,但与图像的非因果性质相矛盾。作者使用重构和完全采样线圈合成的图像之间的L1范数差异作为训练目标,同时采用adamW优化器在网络权重上施加解耦权重衰减[27]。作者的简化训练目标可以表示为以下形式:
其中和分别是将两个采样方案的线圈合成后的匹配图像,指的是重构图像,是下采样 Mask ,是线圈灵敏度,而是模型参数。
为了包含复杂MR数据,作者使用两个独立通道进行实部和虚部的转换。在数据一致性模块中,作者通过全连接层以2个通道的方式,将patchified图像还原为原始大小,从而转换为复数形式,并进行中心傅立叶变换。然后,作者使用下采样 Mask 和线圈灵敏度映射替换生成的k空间点,这些映射使用ESPIRIT[41]进行估计,使用默认参数,如图1中的Eqs.(10)和(11)所示。在第i个数据一致性模块中的步骤可以表示为:
在这里,和分别表示前后2D傅立叶变换,是基于卷积的patch嵌入层,是非patchily层,是第i个数据一致性层的patchified输入,是冗余中间输出,是第i个数据一致性层的patchified输出,而SiLU是Swish[36]激活函数。
这种架构设计使得信息可以通过patches在所有可能的方向上传递,同时保持对物理模型的依赖,通过在VSSM模块之间施加严格的硬数据一致性。
Network and Training Details
AdamW 优化器用于使用默认参数进行100,000次迭代训练,批量大小为4。学习率在 Warm up 阶段使用半周期余弦退火,从 衰减到 。在超参数调优中,尽管在ablation研究中考虑了1、4和8的大小,但选择了2作为划痕大小。连续的VSSM和数据一致性块数量选择为6;4和8也作为在ablation研究中展示。VSSM的内部状态维数选择为16,参考[26],特征的隐藏维数选择为128。在整个网络中使用SiLU激活[8]。
5 Experiments
每一种竞争方法都是通过使用下采样和完全采样图像对进行端到端训练。加速通过使用变量密度的2D高斯 Mask 实现,其中方差调整为实现4和8种加速率。所有实验都在一个PyTorch框架的单个NVIDIA RTX A5000 GPU上进行。每个模型的参数数量如表3所示。
Datasets
作者在实验中使用了以下的大脑 MRI 数据集。
- fastMRI:考虑多联脑 MRI 数据集 [23]。参与者被分为训练、验证和测试组:(100, 10, 40)。T-、T-加权和平面回波加权和 Flair 采集都被考虑在内。所有参与者采用不同的协议。使用 [54] 将联数减少到 5 以降低计算成本。使用 ESPIRIT 估计默认参数 [41] 中的敏感度图。
- IXI:考虑来自IXI http://brain-development.org/ixi-dataset/ 的单联大脑 MRI 数据。T-、T 加权和点扩散加权成像都被采用。所有实验中,25个参与者在所有实验中的训练、5个参与者的验证和10个参与者的测试分别使用了 T-、T 和 PD-加权成像。
Competing Methods
图2:来自快速MRI的T2图像的重建,加速比为4。重构区域和错误地图附在顶部下方。
* UNET: 使用 [23] 中的U型网络模型。网络超参数和实现代码收集自 [50]。使用默认参数的Adam优化器,学习率0.002,迭代次数为150,000。
- E2E-Varnet: 考虑端到端变分网络 [40]。网络超参数和实现代码收集自 [50]。使用默认参数的Adam优化器,学习率0.0005,迭代次数为500,000。
- SwinUnet: 适应于MRI重建的基于 Transformer 的U型分割模型 [5]。使用AdamW优化器进行训练,学习率0.0002,在第五个周期后使用余弦退化为0.0002。训练过程持续100,000次迭代。
- SwinMR: 考虑基于Swin- Transformer 的MRI重建模型 [19]。使用默认参数的Adam优化器,学习率0.0002,迭代次数为100,000。
Ablation Studies
作者进行了多种消融实验来证明超参数和设计选择的单一效应。消融实验的性能指标如表4所示。正如指标所显示的,使用过拟合深的模型或非常小的分块大小会导致性能损失,此外还需要增加计算预算。下面列出进行了的消融实验:
表2:在R = 4和8时,IXI数据集上的重构度量。
图4:来自快速MRI的Flair图像的加速重构。将重构区域和误差图附加在重构的顶部下方。
1. 只有DC:只考虑MambaRecon,其中删除了模型中的VSSM块。
2. SwinRecon:考虑MambaRecon,其中使用了与VSSM块具有相同隐层维数(128)和分块大小()的Swin Transformer 块替换VSSM块。窗口大小已指定为8。
- MambaRecon具有更浅和更深的前端是考虑的,深度等于4和8。其中深度表示连续的VSSM和数据一致性块对的数量。
- MambaRecon考虑了不同的分块大小,其中从1到8。
6 Results
作者对MambaRecon与一系列模型和数据集进行了全面的比较,以彻底评估其性能。作者的基准包括最先进的物理指导CNN模型,E2E-Varnet和RecurrentVarnet,它们利用物理模型进行增强图像重建。
作者还包括使用先进的 Transformer 架构进行重建的变形体 Baseline ,由SwinUNET和SwinMR代表,这些方法捕捉长程依赖关系。最后,作者考虑了一个纯数据驱动的CNN模型,如UNET,它仅依靠数据进行图像重建。作者考虑了由fastMRI和IXI表示的多线圈复杂和单线圈幅度MRI数据集。
这项多样选择使作者能够在各种重建范式下评估模型的能力,从物理指导方法到先进的 Transformer 方法再到传统的CNN方法。
图5:IXI中速度为4的PD图像的重建结果。对齐了缩放区域和错误图在重建之上和之下。
峰值信噪比(PSNR)和结构相似性指数测量(SSIM)[43]被视为比较指标。为每个加速率和对比度对单独呈现结果。每个测试案例的最佳结果用粗体突出显示。表1和表2分别给出了fastMRI和IXI的性能指标。图2显示了加速率为4的T2加权切片的代表性重建结果,图4是从fastMRI数据集的8倍加速率获得的Flair加权切片。图5和图6分别表示从IXI数据集的加速率为4的PD和T1加权重建。
定量化地,与次佳方法相比,MambaRecon在fastMRI数据集上平均PSNR提高了0.72dB,在IXI数据集上平均PSNR提高了1.01dB。从视觉效果上看,MambaRecon捕捉了几乎所有的高频细节,因此其在重建图像中的错误图最暗,没有明显的过度平滑效应或噪声。
作者还说明了每个竞争方法的有效感受野在图3中进行可视化。作者观察到,最强大的感受野对应于MambaRecon,它产生了整个图像的全局感受野。作者在表3中呈现了每个竞争方法中的参数数量。
7 Discussion and Future Work
作者提出了一种轻量级且高效的重建模型,该模型结合了自编码器(SSM)与物理引导的MRI模型。为了克服原始自编码器单向跨越的限制,作者采用了多向跨越的方法。然而,这种方法可能不如自注意力 Transformer 在完全捕捉图像块之间的交互上达到最优。实现更高级的扫描技术可能有助于提高模型的性能。
作者的方法目前依赖于ESPRIT[41]进行线圈敏感性估计。通过将此过程集成到统一的框架中,作者可以实现端到端的训练,从而可能提高整个模型的效率和准确性。这种统一的途径将简化工作流程,允许同时优化线圈敏感性估计以及随后的图像重建,从而提高性能。
作者的模型,与其他所有竞争方法一样,使用来自公共MRI数据集的欠采样和全采样图像进行监督训练。为了消除对监督培训的依赖,可以采用自监督训练损失,如[46]所示。此外,在测试时采用适应性方法可能有助于纳入物理模型(如加速率或欠采样模式)的变化。作者将这方面的探索留待未来工作。
8 Conclusion
作者提出了一种物理引导的MRI重建框架,其核心利用了结构化的状态空间模型。
该模型在不增加计算负担的情况下,从长距离上下文的敏感性中受益,并且与现有最先进的重建基线相比,提供了更好的重建质量。
参考
[1].MambaRecon: MRI Reconstruction with Structured State Space Models.
点击上方卡片,关注 「AI视界引擎」 公众号