[논문 리뷰] Robustness of Deep Learning for Accelerated MRI: Benefits of Diverse Training Data

Kang Lin, Reinhard Heckel|arXiv (Cornell University)|2023. 12. 16.

Advanced X-ray and CT Imaging인용 수 4

한 줄 요약

이 논문은 다양한 데이터 분포 조합—예를 들어 다른 MRI 스캐너, 해부학적 부위, 대조도 및 자기장 강도—에서 딥러닝 모델을 훈련시킬 경우, 분포 이탈에 대한 강건성을 크게 향상시키면서도 내분포 성능을 손상시키지 않고도 가속 MRI에 대해 성능을 높일 수 있음을 보여준다. 핵심 결과는 하나의 모델을 대규모이고 다양한 데이터셋에서 훈련시킬 경우, 개별적으로 좁은 범위의 데이터셋에서 훈련된 모델보다 우수하며, 특히 외부 분포 데이터에서의 성능 격차를 크게 줄일 수 있다는 것이다.

ABSTRACT

Deep learning based methods for image reconstruction are state-of-the-art for a variety of imaging tasks. However, neural networks often perform worse if the training data differs significantly from the data they are applied to. For example, a model trained for accelerated magnetic resonance imaging (MRI) on one scanner performs worse on another scanner. In this work, we investigate the impact of the training data on a model's performance and robustness for accelerated MRI. We find that models trained on the combination of various data distributions, such as those obtained from different MRI scanners and anatomies, exhibit robustness equal or superior to models trained on the best single distribution for a specific target distribution. Thus training on such diverse data tends to improve robustness. Furthermore, training on such a diverse dataset does not compromise in-distribution performance, i.e., a model trained on diverse data yields in-distribution performance at least as good as models trained on the more narrow individual distributions. Our results suggest that training a model for imaging on a variety of distributions tends to yield a more effective and robust model than maintaining separate models for individual distributions.

연구 동기 및 목표

• 다양한 MRI 데이터 분포에서 딥러닝 모델을 훈련시키는 것이 분포 이탈에 대한 강건성을 향상시키는지 조사하기.
• 다양한 훈련 데이터가 개별 분포에서의 내분포 성능을 저하시키는지 평가하기.
• 다양한 데이터에서 훈련된 모델이 훈련 중에 볼 수 없었던 병변이나 비정상 환자에게도 일반화 가능한지 평가하기.
• 훈련 기간이 강건성에 미치는 영향, 특히 분포 과적합의 위험을 분석하기.
• 단일 모델이 대규모이고 다양한 데이터셋에서 훈련된 경우, fastMRI와 같은 고립된 데이터셋에서 훈련된 모델보다 성능이 뛰어나다는 것을 검증하기.

제안 방법

• fastMRI 뇌 및 무릎, CC-359, Stanford 2D, M4Raw GRE, NYU 데이터를 포함한 총 13개의 다양한 MRI 데이터셋을 결합한 데이터셋에서 U-Net, Vision Transformer(ViT), VarNet을 훈련하기.
• 모든 데이터셋을 fastMRI 규격에 맞추기 위해 전처리하며, 순수한 노이즈로만 구성된 슬라이스는 제외하여 총 413,000개의 훈련 슬라이스 확보.
• 분포 과적합을 방지하기 위해 fastMRI 무릎 검증 세트 성능 기반 조기 정지 전략 적용.
• fastMRI 무릎 및 뇌에서의 내분포 성능을 평가하고, 다양한 스캐너 제조사, 해부학적 시야, 대조도 및 스캔 시퀀스를 가진 4개의 추가 데이터셋에서의 외부 분포 성능을 평가.
• fastMRI 뇌 또는 무릎 데이터만으로 훈련된 베이스라인 모델들과, 각 외부 분포 데이터셋에서 미세조정된 모델들과의 성능 비교.
• 구현 및 훈련을 위해 fastMRI 레포지토리의 기존 코드 활용.

실험 결과

연구 질문

• RQ1 다양한 MRI 데이터 분포에서 훈련된 단일 딥러닝 모델이 개별 분포에서 훈련된 모델보다 분포 이탈에 대한 강건성이 향상되는가?
• RQ2 다양한 데이터에서 훈련하는 것이 개별 데이터 소스에서의 내분포 성능을 저하시키는가?
• RQ3 다양한 건강한 환자에서 훈련된 모델이 병변이 있는 이미지를 정확하게 재구성할 수 있는가? 이는 훈련 중에 해당 데이터를 볼 수 없었더라도 가능한가?
• RQ4 장기간의 훈련이 분포 과적합을 유도하는가? 즉, 내분포 성능은 계속 향상되지만 외부 분포 성능은 악화되는가?
• RQ5 대규모이고 다양한 데이터셋에서 훈련된 단일 모델이, 특정 데이터셋에서 전용으로 훈련된 모델과 비교해 외부 분포 데이터에서의 성능 격차를 줄일 수 있는가?

주요 결과

• 다양한 MRI 스캐너, 해부학적 부위, 대조도 및 스캔 시퀀스를 포함한 다각도의 데이터셋에서 훈련된 모델은 fastMRI와 같은 개별 데이터셋에서 훈련된 모델보다 외부 분포 성능이 뛰어나다.
• 다양한 데이터셋에서 훈련된 모델 $π_{P}$ 는 fastMRI 뇌 또는 무릎 데이터에서 훈련된 모델의 내분포 성능에 맞추거나 초월하였다.
• CC-359의 사시면 데이터셋에서, $π_{P}$ 에서 훈련된 VarNet은 분포 이탈 성능 격차를 거의 완전히 해소했으며, 동일 데이터셋에서 미세조정된 모델과 비슷한 재구성 품질을 달성했다.
• 다양한 데이터셋에서 훈련된 모델의 외부 분포 재구성 결과는 fastMRI에서 훈련된 모델보다 더 높은 이미지 품질과 더 적은 아티팩트를 보였다.
• 분포 과적합 현상이 관찰되었다: 장기간의 훈련이 진행될수록 외부 분포 성능이 급격히 악화되었으며, 내분포 성능는 약간 향상되는 추세였다.
• 조기 정지는 내분포 성능를 거의 최적 수준으로 유지하면서도 외부 분포 일반화 능력을 유지하는 데 효과적이었다.