[論文レビュー] 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

ひとこと要約

本稿では、深層学習モデルのMRI再構成および修復のための知識蒸留フレームワークKD-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）手法を導入する。
• 注目転送と模倣損失を組み合わせ、学生ネットワークが教師の出力分布を模倣するようにガイドする。
• 2段階のトレーニングプロセスを採用：まず特徴量蒸留による事前学習を行い、その後再構成損失を用いたファインチューニング。
• 評価の基盤としてDC-CNNアーキテクチャを用い、KDフレームワークを再構成とスーパーレゾリューションの両タスクに適用する。
• 複数のMRIデータセット（心臓、脳、膝）および加速要因（4倍、5倍、8倍）にわたってフレームワークを適用する。

実験結果

リサーチクエスチョン

• RQ1 知識蒸留は、MRI再構成のための深層学習モデルを、性能の著しい低下を伴わずに圧縮できるか？
• RQ2 注目に基づく特徴量蒸留は、他の特徴量蒸留手法（FSP、FN、SP、AH）と比較して、MRI再構成においてどのように優れているか？
• RQ3 模倣損失と注目転送が、個別および併用した場合に学生ネットワークの性能向上に果たす寄与は何か？
• RQ4 提案されたKDフレームワークは、スーパーレゾリューションを含む他のMRIタスクへも拡張可能か？
• RQ5 モデル圧縮は、診断用画像品質を維持したまま、推論時間をどの程度短縮し、パラメータ数を削減できるか？

主な発見

• 膝MRIデータセットにおいて、提案されたKDフレームワークで学習した学生ネットワークは、教師ネットワーク比で65%のパラメータ削減を達成した。
• 単一画像再構成において、学生ネットワークはCPUで2倍速、GPUで1.5倍速の推論速度を示した。
• 注目に基づく蒸留手法（AT）は、他の特徴量蒸留手法（FN、FSP、SP、AH）と比較して、検証損失と再構成品質の両面で優れた性能を示した。
• アブレーションスタディにより、模倣損失（L_imit）単体よりも注目転送（L_AT）が性能向上により顕著な寄与をしていることが確認された。
• L_ATとL_imitの併用により、最小の検証誤差が達成され、相乗効果が示された。
• フレームワークはMRIスーパーレゾリューションへも成功裏に拡張され、付録に好ましい結果が示された。