[论文解读] MRI Image Reconstruction via Learning Optimization Using Neural ODEs

Eric Z. Chen, Terrence Chen+1 more|arXiv (Cornell University)|Jun 24, 2020

Model Reduction and Neural Networks参考文献 39被引用 6

一句话总结

该论文提出ReconODE，一种新颖的MRI重建方法，通过神经微分方程（Neural ODEs）对图像重建的连续优化轨迹进行建模。通过学习优化过程的动力学特性，并将ODE求解器的知识融入网络架构，该方法在参数量显著少于现有深度学习模型（如UNet和级联CNN）的情况下，实现了最先进的重建质量。

ABSTRACT

We propose to formulate MRI image reconstruction as an optimization problem and model the optimization trajectory as a dynamic process using ordinary differential equations (ODEs). We model the dynamics in ODE with a neural network and solve the desired ODE with the off-the-shelf (fixed) solver to obtain reconstructed images. We extend this model and incorporate the knowledge of off-the-shelf ODE solvers into the network design (learned solvers). We investigate several models based on three ODE solvers and compare models with fixed solvers and learned solvers. Our models achieve better reconstruction results and are more parameter efficient than other popular methods such as UNet and cascaded CNN. We introduce a new way of tackling the MRI reconstruction problem by modeling the continuous optimization dynamics using neural ODEs.

研究动机与目标

• 通过将图像重建表述为连续优化问题，解决加速MRI重建的挑战。
• 通过使用常微分方程（ODEs）将优化过程建模为连续动力系统，克服离散迭代方法和固定求解器的局限性。
• 通过在神经网络中学习ODE求解器的系数和步长（即‘可学习求解器’），而非使用固定求解器，提升参数效率和重建质量。
• 证明将求解器知识整合到网络设计中，相比标准深度学习架构（如UNet和级联CNN），可实现更优性能与更低模型复杂度。
• 研究不同ODE求解器的影响，以及注意力机制在提升重建保真度和泛化能力方面的作用。

提出的方法

• 使用ODE将MRI重建表述为初值问题：dx(t)/dt = f(x(t), t, y, θ)，其中f为建模连续优化动力学的神经网络。
• 使用现成的ODE求解器（如Euler、Runge-Kutta）从初始的欠采样图像积分至完全采样的重建图像。
• 通过允许网络学习ODE求解器的系数和步长，引入‘可学习求解器’（ReconODE-LT），使求解器具备自适应性和可训练性。
• 在神经ODE架构中引入注意力机制，以增强特征表示并提高重建精度。
• 使用通过ODE求解器的反向传播（ReconODE-FT）或伴随法（ReconODE-FA）进行训练，结合梯度检查点技术以减少内存使用。
• 在多个ODE求解器（Euler、RK2、RK4）和数据集上，对比固定求解器模型（ReconODE-FT、ReconODE-FA）与可学习求解器模型（ReconODE-LT）的性能。

实验结果

研究问题

• RQ1 通过Neural ODE将MRI重建建模为连续优化过程，是否能相比离散迭代方法显著提升重建质量？
• RQ2 学习ODE求解器的系数和步长（即‘可学习求解器’）是否能显著优于使用固定预设求解器的性能？
• RQ3 将ODE求解器知识整合到网络架构中，对参数效率和重建保真度有何影响？
• RQ4 注意力机制对Neural ODE-based MRI重建模型的性能有何影响？
• RQ5 梯度检查点技术是否能有效降低Neural ODE-based MRI重建训练中的内存使用，而不损害推理性能？

主要发现

• ReconODE-LT-RK4在fastMRI验证数据集上取得最佳PSNR（4X时为32.39，8X时为30.27）和SSIM（4X时为0.7333，8X时为0.6617），优于所有其他模型，包括UNet。
• ReconODE-LT-Euler仅使用UNet参数量的0.9%（0.03M vs. 3.35M），在4X下达到与UNet相当的SSIM，展现出极高的参数效率。
• ReconODE-LT-RK2使用D5C5参数量的64%，在4X和8X下均达到与D5C5相近的SSIM和PSNR，表明其具有更优的参数效率。
• 在ReconODE-FT-RK4中引入注意力机制后，4X下的SSIM从0.726（无注意力）提升至0.733，表明其对重建质量有积极影响。
• ReconODE-FT模型（通过求解器反向传播）优于ReconODE-FA模型（伴随法），表明前者在训练中具有更精确的梯度计算。
• 梯度检查点技术使ReconODE-LT-RK4的GPU内存使用量降低73%（从13.6G降至3.7G），训练时间仅增加0.13秒，且不影响推理性能。