[論文レビュー] SFT-KD-Recon: Learning a Student-friendly Teacher for Knowledge Distillation in Magnetic Resonance Image Reconstruction

Matcha Naga Gayathri, Sriprabha Ramanarayanan+4 more|arXiv (Cornell University)|Apr 11, 2023

Advanced MRI Techniques and ApplicationsMedicine被引用数 3

ひとこと要約

本稿では、教師ネットワークと生徒ネットワークのブランチを併用して、知識蒸留の前段階で表現を整列させる、MRI再構成のための生徒フレンドリーな知識蒸留フレームワーク、SFT-KD-Reconを提案する。教師再構成、生徒再構成、教師-生徒模倣の損失を最適化することで、教師と生徒間の性能差を0.53 dBから0.03 dBにまで低減し、4倍および5倍のアンダーサンプリングされた脳および心臓MRIデータセットにおいて、再構成品質と知識蒸留の効率が著しく向上した。

ABSTRACT

Deep cascaded architectures for magnetic resonance imaging (MRI) acceleration have shown remarkable success in providing high-quality reconstruction. However, as the number of cascades increases, the improvements in reconstruction tend to become marginal, indicating possible excess model capacity. Knowledge distillation (KD) is an emerging technique to compress these models, in which a trained deep teacher network is used to distill knowledge to a smaller student network such that the student learns to mimic the behavior of the teacher. Most KD methods focus on effectively training the student with a pre-trained teacher unaware of the student model. We propose SFT-KD-Recon, a student-friendly teacher training approach along with the student as a prior step to KD to make the teacher aware of the structure and capacity of the student and enable aligning the representations of the teacher with the student. In SFT, the teacher is jointly trained with the unfolded branch configurations of the student blocks using three loss terms - teacher-reconstruction loss, student-reconstruction loss, and teacher-student imitation loss, followed by KD of the student. We perform extensive experiments for MRI acceleration in 4x and 5x under-sampling on the brain and cardiac datasets on five KD methods using the proposed approach as a prior step. We consider the DC-CNN architecture and setup teacher as D5C5 (141765 parameters), and student as D3C5 (49285 parameters), denoting a compression of 2.87:1. Results show that (i) our approach consistently improves the KD methods with improved reconstruction performance and image quality, and (ii) the student distilled using our approach is competitive with the teacher, with the performance gap reduced from 0.53 dB to 0.03 dB.

研究の動機と目的

• 過剰なモデル容量に起因する深層カスケードMRI再構成ネットワークにおける収益の減少を解消すること。
• 教師ネットワークが生徒のアーキテクチャと容量を認識することで、MRI再構成における知識蒸留を改善すること。
• 蒸留の前段階で共同訓練により、教師と生徒の特徴をより良いレベルで整列させること。
• 最小限の性能劣化で高精度で軽量なMRI再構成を実現すること。
• 脳および心臓の複数のデータセットにおいて、一貫した改善を示すこと。

提案手法

• 教師再構成損失、生徒再構成損失、教師-生徒模倣損失の3つの損失項を用いて、教師および生徒ネットワークを共同で訓練する。
• 各教師ブロックに複数の生徒ネットワークブランチを統合し、事前訓練中に特徴レベルの知識移行を可能にする。
• 教師および生徒の再構成にL1損失を適用し、訓練中に画像ドメインの忠実性を保証する。
• SFT訓練の後に知識蒸留を実施し、整列された表現を活用して生徒の性能を向上させる。
• 各ブロックの後にデータ忠実性層を設けたブロック構造のDC-CNNアーキテクチャを採用し、k-spaceの一貫性を維持する。
• 教師としてD5C5（141,765パラメータ）、生徒としてD3C5（49,285パラメータ）を用い、2.87:1の圧縮比を達成した。

Figure 1: Comparison between the standard KD and SFT-KD-Recon. (a) The standard KD trains teacher alone and distills knowledge to student. (b) SFT-KD-Recon trains the teacher along with the student branches and then distills effective knowledge to student. (c) SFT Vs SFT-KD-Recon, the former learns

実験結果

リサーチクエスチョン

• RQ1 生徒ネットワークと協働して訓練された教師ネットワークは、MRI再構成における知識蒸留性能を向上させることができるか？
• RQ2 共同訓練による教師と生徒の表現整列は、教師と蒸留された生徒モデル間の性能差を低減するか？
• RQ3 SFT-KD-Reconは、4倍および5倍のアンダーサンプリングされたMRIデータにおいて、標準的なKD手法と比較して画像品質と再構成忠実度で優れているか？
• RQ4 提案手法は、脳および心臓の多様なMRIデータセットにおいて、複数のKD技術を一貫して改善できるか？
• RQ5 生徒に配慮した教師訓練は、生徒ネットワークの初期化および収束性をどの程度向上させるか？

主な発見

• 提案されたSFT-KD-Reconフレームワークは、4倍および5倍のアンダーサンプリングされた脳および心臓MRIデータセットにおいて、5種類の異なる知識蒸留手法の性能を一貫して向上させた。
• 教師と蒸留された生徒モデル間の性能差は、0.53 dBから0.03 dBにまで低減され、教師に匹敵する再構成品質を示した。
• SFT-KD-Reconで訓練された生徒ネットワークは、競争力のある画像品質を達成し、PSNRおよびSSIM値が完全な教師モデルに近い値を示した。
• 生徒ブランチを併用した共同訓練により、低レベルおよび中間レベルの表現における特徴整列が向上し、蒸留効率が向上した。
• 多様な解剖的構造および加速係数においてもロバスト性を維持し、データセット全体にわたる一般化性を示した。
• 教師再構成、生徒再構成、模倣の3つの損失項の使用により、標準的なKDよりも安定的かつ効果的な知識移行が実現された。

Figure 2: Student-Friendly training of the teacher. The teacher DC-CNN has five blocks, each having CNN with five convolution layers and DF layer, and the student DC-CNN has five blocks, each having three convolution layers and a DF layer. The teacher is trained with three loss terms - $L_{rec}^{T}$