[论文解读] 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

一句话总结

该论文提出SFT-KD-Recon，一种面向磁共振成像（MRI）重建的面向学生的知识蒸馏框架，通过联合训练教师网络与学生网络分支，在知识蒸馏前对齐表征。通过优化教师重建、学生重建和教师-学生模仿损失，该方法将教师与学生之间的性能差距从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重建，同时保持最小的性能损失。
• 在多个KD方法与数据集（脑部与心脏）上展示一致的性能提升。

提出的方法

• 联合训练教师与学生网络，采用三种损失项：教师重建损失、学生重建损失与教师-学生模仿损失。
• 在每个对应的教师模块中集成多个学生网络分支，以在预训练阶段实现实现特征级知识迁移。
• 对教师与学生重建任务均采用L1损失，以确保训练过程中图像域保真度。
• 在SFT训练后应用知识蒸馏，利用对齐的表征提升学生模型性能。
• 采用分块结构的DC-CNN架构，在每个模块后加入数据保真度层，以保持k空间一致性。
• 使用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 在4倍与5倍欠采样MRI数据上，SFT-KD-Recon与标准KD方法在图像质量与重建保真度方面相比如何？
• RQ4 所提出方法是否能在多种MRI数据集（脑部与心脏）上持续提升多种KD技术？
• RQ5 学生感知的教师训练在多大程度上提升了学生网络的初始化质量与收敛速度？

主要发现

• 所提出的SFT-KD-Recon框架在4倍与5倍欠采样脑部与心脏MRI数据集上，对五种不同知识蒸馏方法的性能均实现了持续提升。
• 教师与蒸馏后学生模型之间的性能差距从0.53 dB降低至0.03 dB，表明重建质量接近教师模型水平。
• 采用SFT-KD-Recon训练的学生网络实现了具有竞争力的图像质量，其PSNR与SSIM值接近全参数教师模型。
• 与学生分支联合训练可实现更优的特征对齐，尤其在低级与中级表征中，显著提升了知识蒸馏效率。
• 该方法在不同解剖结构与加速因子下均保持鲁棒性，展现出良好的跨数据集泛化能力。
• 采用教师重建、学生重建与模仿三项损失，相比标准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}$