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

TL;DR

This paper proposes KD-MRI, a knowledge distillation framework that compresses deep learning models for MRI reconstruction and restoration by transferring attention-based features and imitation losses from a large teacher network to a compact student network. The method achieves 65% parameter reduction, 2x faster CPU, and 1.5x faster GPU inference with minimal performance drop across cardiac, brain, and knee MRI datasets at 4x–8x acceleration factors.

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.

Motivation & Objective

• To address the growing computational and storage demands of deep learning in MRI workflows due to increasing model complexity.
• To develop a model compression strategy that enables efficient deployment of deep learning models in clinical settings with limited resources.
• To apply knowledge distillation to image-to-image tasks in MRI, specifically reconstruction and super-resolution, for the first time.
• To evaluate the effectiveness of attention-based feature distillation and imitation loss in improving student network performance.
• To demonstrate that compact student networks can match or exceed teacher network performance with significantly reduced parameters and inference time.

Proposed method

• Proposes an end-to-end trainable knowledge distillation framework for MRI reconstruction and restoration tasks.
• Introduces a novel attention-based feature distillation (AT) method that transfers spatially weighted feature maps from teacher to student.
• Combines attention transfer with an imitation loss to guide the student network in mimicking the teacher's output distribution.
• Employs a two-stage training process: pre-training with feature distillation, followed by fine-tuning using reconstruction loss.
• Uses the DC-CNN architecture as the baseline for evaluation and applies the KD framework to both reconstruction and super-resolution tasks.
• Applies the framework across multiple MRI datasets (cardiac, brain, knee) and acceleration factors (4x, 5x, 8x).

Experimental results

Research questions

• RQ1 Can knowledge distillation effectively compress deep learning models for MRI reconstruction without significant performance degradation?
• RQ2 How does attention-based feature distillation compare to other feature distillation methods (e.g., FSP, FN, SP, AH) in MRI reconstruction?
• RQ3 What is the contribution of imitation loss and attention transfer individually and jointly in improving student network performance?
• RQ4 Can the proposed KD framework be extended to other MRI tasks such as super-resolution?
• RQ5 To what extent can model compression reduce inference time and parameter count while maintaining diagnostic image quality?

Key findings

• The student network trained with the proposed KD framework achieved 65% parameter reduction compared to the teacher network on the knee MRI dataset.
• The student network ran 2 times faster on CPU and 1.5 times faster on GPU compared to the teacher network for single-image reconstruction.
• The attention-based distillation method (AT) outperformed other feature distillation techniques (FN, FSP, SP, AH) in terms of validation loss and reconstruction quality.
• The ablation study confirmed that attention transfer (L_AT) contributes more significantly to performance improvement than imitation loss (L_imit) alone.
• The combination of L_AT and L_imit yielded the lowest validation error, demonstrating synergistic benefits.
• The framework was successfully extended to MRI super-resolution, yielding encouraging results in the appendix.