[Paper Review] 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 ApplicationsMedicine3 citations

TL;DR

This paper proposes SFT-KD-Recon, a student-friendly knowledge distillation framework for MRI reconstruction that jointly trains a teacher network with student network branches to align representations before distillation. By optimizing teacher-reconstruction, student-reconstruction, and teacher-student imitation losses, the method reduces the performance gap between teacher and student from 0.53 dB to 0.03 dB, significantly improving reconstruction quality and distillation efficiency in 4x and 5x undersampled brain and cardiac MRI datasets.

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.

Motivation & Objective

• To address the diminishing returns in deep cascaded MRI reconstruction networks due to excess model capacity.
• To improve knowledge distillation in MRI reconstruction by making the teacher network aware of the student’s architecture and capacity.
• To enable better feature alignment between teacher and student through joint training prior to distillation.
• To achieve high-fidelity, lightweight MRI reconstruction with minimal performance degradation.
• To demonstrate consistent improvements across multiple KD methods and datasets (brain and cardiac).

Proposed method

• Jointly train the teacher and student networks using three loss terms: teacher-reconstruction loss, student-reconstruction loss, and teacher-student imitation loss.
• Integrate multiple student network branches into each corresponding teacher block to enable feature-level knowledge transfer during pre-training.
• Use L1 loss for both teacher and student reconstruction to ensure image-domain fidelity during training.
• Apply knowledge distillation after SFT training, leveraging the aligned representations for improved student performance.
• Employ a block-structured DC-CNN architecture with data fidelity layers after each block to maintain k-space consistency.
• Use D5C5 (141,765 parameters) as teacher and D3C5 (49,285 parameters) as student, achieving a 2.87:1 compression ratio.

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

Experimental results

Research questions

• RQ1 Can a teacher network trained in collaboration with the student network improve knowledge distillation performance in MRI reconstruction?
• RQ2 Does aligning teacher and student representations through joint training reduce the performance gap between teacher and distilled student models?
• RQ3 How does SFT-KD-Recon compare to standard KD methods in terms of image quality and reconstruction fidelity on 4x and 5x undersampled MRI data?
• RQ4 Can the proposed method consistently improve multiple KD techniques across diverse MRI datasets (brain and cardiac)?
• RQ5 To what extent does student-aware teacher training enhance the initialization and convergence of the student network?

Key findings

• The proposed SFT-KD-Recon framework consistently improves the performance of five different knowledge distillation methods across 4x and 5x undersampled brain and cardiac MRI datasets.
• The performance gap between the teacher and distilled student model is reduced from 0.53 dB to 0.03 dB, demonstrating near-teacher-level reconstruction quality.
• The student network trained with SFT-KD-Recon achieves competitive image quality, with PSNR and SSIM values close to those of the full teacher model.
• Joint training with student branches enables better feature alignment, especially in low-level and intermediate representations, improving distillation efficiency.
• The method maintains robustness across diverse anatomical structures and acceleration factors, showing generalization across datasets.
• The use of three loss terms—teacher-reconstruction, student-reconstruction, and imitation—leads to more stable and effective knowledge transfer than standard 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}$