[论文解读] KD-MRI: A knowledge distillation framework for image reconstruction and image restoration in MRI workflow

Balamurali Murugesan, Sricharan Vijayarangan+3 more|arXiv (Cornell University)|Apr 11, 2020

Medical Imaging Techniques and Applications参考文献 27被引用 7

一句话总结

本文提出KD-MRI，一种知识蒸馏框架，通过将注意力机制特征和模仿损失从大型教师网络迁移至紧凑的学生网络，实现对磁共振成像（MRI）重建与修复任务中深度学习模型的压缩。该方法在4倍至8倍加速因子下，于心脏、脑部和膝关节MRI数据集上实现了65%的参数减少、CPU推理速度提升2倍、GPU推理速度提升1.5倍，且性能下降可忽略不计。

ABSTRACT

Deep learning networks are being developed in every stage of the MRI workflow and have provided state-of-the-art results. However, this has come at the cost of increased computation requirement and storage. Hence, replacing the networks with compact models at various stages in the MRI workflow can significantly reduce the required storage space and provide considerable speedup. In computer vision, knowledge distillation is a commonly used method for model compression. In our work, we propose a knowledge distillation (KD) framework for the image to image problems in the MRI workflow in order to develop compact, low-parameter models without a significant drop in performance. We propose a combination of the attention-based feature distillation method and imitation loss and demonstrate its effectiveness on the popular MRI reconstruction architecture, DC-CNN. We conduct extensive experiments using Cardiac, Brain, and Knee MRI datasets for 4x, 5x and 8x accelerations. We observed that the student network trained with the assistance of the teacher using our proposed KD framework provided significant improvement over the student network trained without assistance across all the datasets and acceleration factors. Specifically, for the Knee dataset, the student network achieves $65\%$ parameter reduction, 2x faster CPU running time, and 1.5x faster GPU running time compared to the teacher. Furthermore, we compare our attention-based feature distillation method with other feature distillation methods. We also conduct an ablative study to understand the significance of attention-based distillation and imitation loss. We also extend our KD framework for MRI super-resolution and show encouraging results.

研究动机与目标

• 为应对由于模型复杂度增加而导致的深度学习在MRI工作流程中计算与存储需求不断增长的问题。
• 开发一种模型压缩策略，以实现在计算资源有限的临床环境中高效部署深度学习模型。
• 首次将知识蒸馏应用于MRI中的图像到图像任务，具体为重建与超分辨率。
• 评估基于注意力的特征蒸馏与模仿损失在提升学生网络性能方面的有效性。
• 证明紧凑的学生网络可在显著减少参数量与推理时间的前提下，达到甚至超越教师网络的性能。

提出的方法

• 提出一种端到端可训练的知识蒸馏框架，用于MRI重建与修复任务。
• 提出一种新颖的基于注意力的特征蒸馏（AT）方法，将教师网络的空间加权特征图迁移至学生网络。
• 结合注意力迁移与模仿损失，引导学生网络模仿教师网络的输出分布。
• 采用两阶段训练流程：先通过特征蒸馏进行预训练，再使用重建损失进行微调。
• 以DC-CNN架构作为基线进行评估，并将KD框架应用于重建与超分辨率任务。
• 在多个MRI数据集（心脏、脑部、膝关节）和加速因子（4倍、5倍、8倍）上应用该框架。

实验结果

研究问题

• RQ1 知识蒸馏能否在不造成显著性能下降的前提下，有效压缩用于MRI重建的深度学习模型？
• RQ2 在MRI重建任务中，基于注意力的特征蒸馏方法相较于其他特征蒸馏方法（如FSP、FN、SP、AH）表现如何？
• RQ3 模仿损失与注意力迁移在提升学生网络性能方面各自及联合的贡献分别是什么？
• RQ4 所提出的KD框架能否扩展至其他MRI任务（如超分辨率）？
• RQ5 模型压缩在多大程度上可减少推理时间与参数量，同时保持诊断级图像质量？

主要发现

• 在膝关节MRI数据集上，采用所提KD框架训练的学生网络相比教师网络实现了65%的参数减少。
• 在单张图像重建任务中，学生网络在CPU上推理速度提升2倍，在GPU上提升1.5倍。
• 基于注意力的蒸馏方法（AT）在验证损失与重建质量方面优于其他特征蒸馏技术（FN、FSP、SP、AH）。
• 消融实验表明，注意力迁移（L_AT）对性能提升的贡献显著高于仅使用模仿损失（L_imit）的情况。
• L_AT与L_imit的联合使用实现了最低的验证误差，证明了二者具有协同增益效应。
• 该框架成功扩展至MRI超分辨率任务，在附录中取得了令人鼓舞的结果。