[论文解读] Memory-efficient Learning for High-Dimensional MRI Reconstruction

Ke Wang, Michael Kellman+6 more|arXiv (Cornell University)|Mar 6, 2021

Advanced MRI Techniques and Applications参考文献 21被引用 5

一句话总结

本文提出了一种内存高效的训练（MEL）方法，仅使用单张12GB GPU即可在高维数据（3D和2D+时间）上训练基于深度学习的展开式MRI重建模型。通过在反向传播过程中重新计算中间激活值而非存储它们，MEL将GPU内存使用量降低了N倍（展开层数），仅带来可忽略的计算量增加，从而在体内3D和心脏电影MRI数据集上实现了更优的图像质量，表现为边缘更清晰、各项指标更优。

ABSTRACT

Deep learning (DL) based unrolled reconstructions have shown state-of-the-art performance for under-sampled magnetic resonance imaging (MRI). Similar to compressed sensing, DL can leverage high-dimensional data (e.g. 3D, 2D+time, 3D+time) to further improve performance. However, network size and depth are currently limited by the GPU memory required for backpropagation. Here we use a memory-efficient learning (MEL) framework which favorably trades off storage with a manageable increase in computation during training. Using MEL with multi-dimensional data, we demonstrate improved image reconstruction performance for in-vivo 3D MRI and 2D+time cardiac cine MRI. MEL uses far less GPU memory while marginally increasing the training time, which enables new applications of DL to high-dimensional MRI.

研究动机与目标

• 克服限制基于深度学习的展开式MRI重建在3D和2D+时间等高维数据上训练的GPU内存瓶颈。
• 使展开重建网络中可使用更多展开层数，从而通过捕捉更精细的空间和时间纹理提升图像质量。
• 证明内存高效的训练（MEL）可在标准硬件（单张12GB GPU）上训练复杂且高维的MRI重建模型，而无需牺牲性能。
• 验证通过MEL利用高维冗余可带来定量指标的提升以及临床相关的图像质量改进。

提出的方法

• MEL通过在反向传播期间仅存储单层激活值，而非所有N个展开层的激活值，从而减少内存使用；中间特征通过反向顺序重新计算。
• 该方法依赖于展开网络中每一层的可逆性，使得梯度可通过从输出反向重新计算每一层的前向传播来获得。
• 展开重建采用MoDL框架建模，其中每个展开步骤交替包含基于CNN的正则化层和通过共轭梯度求解的数据一致性（DC）层。
• DC层通过包含编码算子的伴随算子和Tikhonov正则化项的闭式解实现。
• 正则化层中使用的残差CNN去噪器具备可逆性，这对MEL的重新计算策略至关重要。
• 训练在单张12GB GPU上进行，MEL使原本超出内存限制的模型得以训练。

实验结果

研究问题

• RQ1 内存高效的训练（MEL）是否能够实现对3D和2D+时间MRI等高维数据上基于深度学习的展开式MRI重建模型的训练？
• RQ2 通过MEL实现更多展开层数的展开重建，是否能带来高维MRI中图像质量与定量指标的提升？
• RQ3 MEL训练的模型是否能在无全采样真实值的情况下泛化到前瞻性 undersampled 数据，表明其临床鲁棒性？
• RQ4 在高维MRI重建中，MEL与标准反向传播相比，在内存使用和训练时间方面有何差异？

主要发现

• 3D MoDL结合MEL在体内3D膝关节MRI上达到32.11 ± 2.05 dB的pSNR，优于2D MoDL和PICS。
• 2D+时间MoDL结合MEL与10次展开在心脏电影MRI上达到27.42 ± 3.21 dB的pSNR，较4次展开基线提升0.6 dB。
• 2D+时间MoDL结合MEL与10次展开的版本相比4次展开版本，生成了更自然的运动轨迹，并更清晰地分辨了乳头肌结构。
• 3D MoDL结合MEL在3D MRI上相比PICS和2D MoDL展现出更优的SSIM（0.830 ± 0.038）和FID（41.48）。
• 在前瞻性undersampled的3D FSE膝关节扫描中，3D MoDL结合MEL即使在扫描参数不匹配的情况下，仍能保持更清晰的边缘和更精细的纹理。
• 视觉与定量结果共同证实，MEL可实现更高维重建，且图像质量更优，包括更佳的对比度、纹理连续性与运动保真度。