[논문 리뷰] KD-MRI: A knowledge distillation framework for image reconstruction and image restoration in MRI workflow

Balamurali Murugesan, Sricharan Vijayarangan외 3명|arXiv (Cornell University)|2020. 04. 11.

Medical Imaging Techniques and Applications참고 문헌 27인용 수 7

한 줄 요약

이 논문은 심장, 뇌, 무릎 MRI 데이터셋에서 4x–8x 가속도 요인에서 성능 저하를 최소화하면서 65%의 파rameter 감소, CPU에서 2배 빠른, GPU에서 1.5배 빠른 추론를 달성하는, 깊이 학습 모델을 MRI 복원 및 복구에 압축하기 위한 지식 정렬 프레임워크인 KD-MRI를 제안한다. 이는 대규모 교사 네트워크에서 소형 학생 네트워크로 주의 기반 특징과 모방 손실을 전달한다.

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.

연구 동기 및 목표

• 모델 복잡도 증가로 인해 증가하는 MRI 워크플로우에서의 계산 및 저장 요구량을 해결하기 위해.
• 자원이 제한된 임상 환경에서 깊이 학습 모델의 효율적 배포를 가능하게 하는 모델 압축 전략을 개발하기 위해.
• 이번에 처음으로 MRI의 이미지 간 작업, 특히 복원 및 초해상도 증강에 지식 정렬을 적용하기 위해.
• 주의 기반 특징 정렬과 모방 손실이 학생 네트워크 성능 향상에 기여하는지 평가하기 위해.
• 소형 학생 네트워크가 훨씬 적은 파라미터와 추론 시간으로 교사 네트워크 성능을 따라하거나 초월할 수 있는지 보여주기 위해.

제안 방법

• MRI 복원 및 복구 작업을 위한 엔드 투 엔드로 훈련 가능한 지식 정렬 프레임워크를 제안한다.
• 교사에서 학생으로 공간적으로 가중된 특징 맵을 전달하는 새로운 주의 기반 특징 정렬(AT) 방법을 도입한다.
• 주의 전달과 모방 손실을 조합하여 학생 네트워크가 교사의 출력 분포를 모방하도록 이끈다.
• 두 단계 훈련 과정을 사용한다: 특징 정렬을 통한 사전 훈련 후, 복원 손실을 사용한 파생 훈련.
• 평가를 위한 기준으로 DC-CNN 아키텍처를 사용하고, 복원 및 초해상도 증강 작업에 KD 프레임워크를 적용한다.
• 여러 MRI 데이터셋(심장, 뇌, 무릎)과 가속도 요인(4x, 5x, 8x)에 걸쳐 프레임워크를 적용한다.

실험 결과

연구 질문

• RQ1 지식 정렬이 성능 저하 없이 MRI 복원을 위한 깊이 학습 모델을 효과적으로 압축할 수 있는가?
• RQ2 다른 특징 정렬 방법(FSP, FN, SP, AH 등)과 비교해 주의 기반 특징 정렬이 MRI 복원에서 어떻게 성능을 냈는가?
• RQ3 모방 손실과 주의 전달이 각각 및 함께 학생 네트워크 성능 향상에 기여하는 정도는 어떠한가?
• RQ4 제안된 KD 프레임워크는 초해상도와 같은 다른 MRI 작업으로 확장될 수 있는가?
• RQ5 모델 압축이 진단 이미지 품질을 유지하면서 추론 시간과 파라미터 수를 얼마나 줄일 수 있는가?

주요 결과

• 제안된 KD 프레임워크로 훈련된 학생 네트워크는 무릎 MRI 데이터셋에서 교사 네트워크 대비 65%의 파라미터 감소를 달성했다.
• 단일 이미지 복원에서 학생 네트워크는 CPU에서 2배, GPU에서 1.5배 더 빠른 속도로 실행되었다.
• 주의 기반 정렬 방법(AT)은 검증 손실과 복원 품질 측면에서 다른 특징 정렬 기법(FN, FSP, SP, AH)보다 뛰어났다.
• 제거 실험을 통해 주의 전달(L_AT)이 모방 손실(L_imit)만으로는 성능 향상에 더 크게 기여하는 것으로 확인되었다.
• L_AT와 L_imit의 조합은 가장 낮은 검증 오차를 기록하여 상호 보완적 이점이 있음을 입증했다.
• 프레임워크는 MRI 초해상도 작업으로 성공적으로 확장되었으며, 부록에서 유망한 결과를 도출했다.