[논문 리뷰] SFT-KD-Recon: Learning a Student-friendly Teacher for Knowledge Distillation in Magnetic Resonance Image Reconstruction

Matcha Naga Gayathri, Sriprabha Ramanarayanan외 4명|arXiv (Cornell University)|2023. 04. 11.

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 재구성을 달성하기 위해.
• 뇌 및 심장 MRI 데이터셋에서 다양한 KD 방법에 대해 일관된 향상을 보여주기 위해.

제안 방법

• 교사-재구성 손실, 학생-재구성 손실, 교사-학생 모방 손실의 세 가지 손실 항목을 사용하여 교사 및 학생 네트워크를 공동으로 훈련시킨다.
• 각 교사 블록에 다수의 학생 네트워크 브랜치를 통합하여 사전 훈련 중 특징 수준의 지식 전달을 가능하게 한다.
• 훈련 중 이미지 도메인의 정밀도를 확보하기 위해 교사 및 학생 재구성에 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 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}$